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

     1  package gcp
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"io/ioutil"
     7  	"time"
     8  
     9  	"cloud.google.com/go/storage"
    10  	"github.com/pkg/errors"
    11  
    12  	"github.com/sequix/cortex/pkg/chunk"
    13  	"github.com/sequix/cortex/pkg/chunk/util"
    14  )
    15  
    16  type gcsObjectClient struct {
    17  	cfg       GCSConfig
    18  	schemaCfg chunk.SchemaConfig
    19  	client    *storage.Client
    20  	bucket    *storage.BucketHandle
    21  }
    22  
    23  // GCSConfig is config for the GCS Chunk Client.
    24  type GCSConfig struct {
    25  	BucketName      string        `yaml:"bucket_name"`
    26  	ChunkBufferSize int           `yaml:"chunk_buffer_size"`
    27  	RequestTimeout  time.Duration `yaml:"request_timeout"`
    28  }
    29  
    30  // RegisterFlags registers flags.
    31  func (cfg *GCSConfig) RegisterFlags(f *flag.FlagSet) {
    32  	f.StringVar(&cfg.BucketName, "gcs.bucketname", "", "Name of GCS bucket to put chunks in.")
    33  	f.IntVar(&cfg.ChunkBufferSize, "gcs.chunk-buffer-size", 0, "The size of the buffer that GCS client for each PUT request. 0 to disable buffering.")
    34  	f.DurationVar(&cfg.RequestTimeout, "gcs.request-timeout", 0, "The duration after which the requests to GCS should be timed out.")
    35  }
    36  
    37  // NewGCSObjectClient makes a new chunk.ObjectClient that writes chunks to GCS.
    38  func NewGCSObjectClient(ctx context.Context, cfg GCSConfig, schemaCfg chunk.SchemaConfig) (chunk.ObjectClient, error) {
    39  	option, err := gcsInstrumentation(ctx, storage.ScopeReadWrite)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	client, err := storage.NewClient(ctx, option)
    45  	if err != nil {
    46  		return nil, err
    47  	}
    48  	return newGCSObjectClient(cfg, schemaCfg, client), nil
    49  }
    50  
    51  func newGCSObjectClient(cfg GCSConfig, schemaCfg chunk.SchemaConfig, client *storage.Client) chunk.ObjectClient {
    52  	bucket := client.Bucket(cfg.BucketName)
    53  	return &gcsObjectClient{
    54  		cfg:       cfg,
    55  		schemaCfg: schemaCfg,
    56  		client:    client,
    57  		bucket:    bucket,
    58  	}
    59  }
    60  
    61  func (s *gcsObjectClient) Stop() {
    62  	s.client.Close()
    63  }
    64  
    65  func (s *gcsObjectClient) PutChunks(ctx context.Context, chunks []chunk.Chunk) error {
    66  	for _, chunk := range chunks {
    67  		buf, err := chunk.Encoded()
    68  		if err != nil {
    69  			return err
    70  		}
    71  		writer := s.bucket.Object(chunk.ExternalKey()).NewWriter(ctx)
    72  		// Default GCSChunkSize is 8M and for each call, 8M is allocated xD
    73  		// By setting it to 0, we just upload the object in a single a request
    74  		// which should work for our chunk sizes.
    75  		writer.ChunkSize = s.cfg.ChunkBufferSize
    76  
    77  		if _, err := writer.Write(buf); err != nil {
    78  			return err
    79  		}
    80  		if err := writer.Close(); err != nil {
    81  			return err
    82  		}
    83  	}
    84  	return nil
    85  }
    86  
    87  func (s *gcsObjectClient) GetChunks(ctx context.Context, input []chunk.Chunk) ([]chunk.Chunk, error) {
    88  	return util.GetParallelChunks(ctx, input, s.getChunk)
    89  }
    90  
    91  func (s *gcsObjectClient) getChunk(ctx context.Context, decodeContext *chunk.DecodeContext, input chunk.Chunk) (chunk.Chunk, error) {
    92  	if s.cfg.RequestTimeout > 0 {
    93  		// The context will be cancelled with the timeout or when the parent context is cancelled, whichever occurs first.
    94  		var cancel context.CancelFunc
    95  		ctx, cancel = context.WithTimeout(ctx, s.cfg.RequestTimeout)
    96  		defer cancel()
    97  	}
    98  
    99  	reader, err := s.bucket.Object(input.ExternalKey()).NewReader(ctx)
   100  	if err != nil {
   101  		return chunk.Chunk{}, errors.WithStack(err)
   102  	}
   103  	defer reader.Close()
   104  
   105  	buf, err := ioutil.ReadAll(reader)
   106  	if err != nil {
   107  		return chunk.Chunk{}, errors.WithStack(err)
   108  	}
   109  
   110  	if err := input.Decode(decodeContext, buf); err != nil {
   111  		return chunk.Chunk{}, err
   112  	}
   113  
   114  	return input, nil
   115  }