github.com/sequix/cortex@v1.1.6/pkg/chunk/gcp/bigtable_object_client.go (about)

     1  package gcp
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"cloud.google.com/go/bigtable"
     8  	ot "github.com/opentracing/opentracing-go"
     9  	otlog "github.com/opentracing/opentracing-go/log"
    10  	"github.com/pkg/errors"
    11  
    12  	"github.com/sequix/cortex/pkg/chunk"
    13  	"github.com/sequix/cortex/pkg/util"
    14  )
    15  
    16  type bigtableObjectClient struct {
    17  	cfg       Config
    18  	schemaCfg chunk.SchemaConfig
    19  	client    *bigtable.Client
    20  }
    21  
    22  // NewBigtableObjectClient makes a new chunk.ObjectClient that stores chunks in
    23  // Bigtable.
    24  func NewBigtableObjectClient(ctx context.Context, cfg Config, schemaCfg chunk.SchemaConfig) (chunk.ObjectClient, error) {
    25  	opts := toOptions(cfg.GRPCClientConfig.DialOption(bigtableInstrumentation()))
    26  	client, err := bigtable.NewClient(ctx, cfg.Project, cfg.Instance, opts...)
    27  	if err != nil {
    28  		return nil, err
    29  	}
    30  	return newBigtableObjectClient(cfg, schemaCfg, client), nil
    31  }
    32  
    33  func newBigtableObjectClient(cfg Config, schemaCfg chunk.SchemaConfig, client *bigtable.Client) chunk.ObjectClient {
    34  	return &bigtableObjectClient{
    35  		cfg:       cfg,
    36  		schemaCfg: schemaCfg,
    37  		client:    client,
    38  	}
    39  }
    40  
    41  func (s *bigtableObjectClient) Stop() {
    42  	s.client.Close()
    43  }
    44  
    45  func (s *bigtableObjectClient) PutChunks(ctx context.Context, chunks []chunk.Chunk) error {
    46  	keys := map[string][]string{}
    47  	muts := map[string][]*bigtable.Mutation{}
    48  
    49  	for i := range chunks {
    50  		buf, err := chunks[i].Encoded()
    51  		if err != nil {
    52  			return err
    53  		}
    54  		key := chunks[i].ExternalKey()
    55  		tableName, err := s.schemaCfg.ChunkTableFor(chunks[i].From)
    56  		if err != nil {
    57  			return err
    58  		}
    59  		keys[tableName] = append(keys[tableName], key)
    60  
    61  		mut := bigtable.NewMutation()
    62  		mut.Set(columnFamily, column, 0, buf)
    63  		muts[tableName] = append(muts[tableName], mut)
    64  	}
    65  
    66  	for tableName := range keys {
    67  		table := s.client.Open(tableName)
    68  		errs, err := table.ApplyBulk(ctx, keys[tableName], muts[tableName])
    69  		if err != nil {
    70  			return err
    71  		}
    72  		for _, err := range errs {
    73  			if err != nil {
    74  				return err
    75  			}
    76  		}
    77  	}
    78  	return nil
    79  }
    80  
    81  func (s *bigtableObjectClient) GetChunks(ctx context.Context, input []chunk.Chunk) ([]chunk.Chunk, error) {
    82  	sp, ctx := ot.StartSpanFromContext(ctx, "GetChunks")
    83  	defer sp.Finish()
    84  	sp.LogFields(otlog.Int("chunks requested", len(input)))
    85  
    86  	chunks := map[string]map[string]chunk.Chunk{}
    87  	keys := map[string]bigtable.RowList{}
    88  	for _, c := range input {
    89  		tableName, err := s.schemaCfg.ChunkTableFor(c.From)
    90  		if err != nil {
    91  			return nil, err
    92  		}
    93  		key := c.ExternalKey()
    94  		keys[tableName] = append(keys[tableName], key)
    95  		if _, ok := chunks[tableName]; !ok {
    96  			chunks[tableName] = map[string]chunk.Chunk{}
    97  		}
    98  		chunks[tableName][key] = c
    99  	}
   100  
   101  	outs := make(chan chunk.Chunk, len(input))
   102  	errs := make(chan error, len(input))
   103  
   104  	for tableName := range keys {
   105  		var (
   106  			table  = s.client.Open(tableName)
   107  			keys   = keys[tableName]
   108  			chunks = chunks[tableName]
   109  		)
   110  
   111  		for i := 0; i < len(keys); i += maxRowReads {
   112  			page := keys[i:util.Min(i+maxRowReads, len(keys))]
   113  			go func(page bigtable.RowList) {
   114  				decodeContext := chunk.NewDecodeContext()
   115  
   116  				var processingErr error
   117  				var receivedChunks = 0
   118  
   119  				// rows are returned in key order, not order in row list
   120  				err := table.ReadRows(ctx, page, func(row bigtable.Row) bool {
   121  					chunk, ok := chunks[row.Key()]
   122  					if !ok {
   123  						processingErr = errors.WithStack(fmt.Errorf("Got row for unknown chunk: %s", row.Key()))
   124  						return false
   125  					}
   126  
   127  					err := chunk.Decode(decodeContext, row[columnFamily][0].Value)
   128  					if err != nil {
   129  						processingErr = err
   130  						return false
   131  					}
   132  
   133  					receivedChunks++
   134  					outs <- chunk
   135  					return true
   136  				})
   137  
   138  				if processingErr != nil {
   139  					errs <- processingErr
   140  				} else if err != nil {
   141  					errs <- errors.WithStack(err)
   142  				} else if receivedChunks < len(page) {
   143  					errs <- errors.WithStack(fmt.Errorf("Asked for %d chunks for Bigtable, received %d", len(page), receivedChunks))
   144  				}
   145  			}(page)
   146  		}
   147  	}
   148  
   149  	output := make([]chunk.Chunk, 0, len(input))
   150  	for i := 0; i < len(input); i++ {
   151  		select {
   152  		case c := <-outs:
   153  			output = append(output, c)
   154  		case err := <-errs:
   155  			return nil, err
   156  		case <-ctx.Done():
   157  			return nil, ctx.Err()
   158  		}
   159  	}
   160  
   161  	return output, nil
   162  }