github.com/sequix/cortex@v1.1.6/pkg/chunk/baidu/bos_object_client.go (about)

     1  package baidu
     2  
     3  import (
     4  	"context"
     5  	"io/ioutil"
     6  
     7  	"github.com/baidubce/bce-sdk-go/services/bos"
     8  	bosAPI "github.com/baidubce/bce-sdk-go/services/bos/api"
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/weaveworks/common/instrument"
    11  
    12  	"github.com/sequix/cortex/pkg/chunk"
    13  	"github.com/sequix/cortex/pkg/chunk/util"
    14  )
    15  
    16  var (
    17  	bosRequestDuration = instrument.NewHistogramCollector(prometheus.NewHistogramVec(prometheus.HistogramOpts{
    18  		Namespace: "cortex",
    19  		Name:      "bos_request_duration_seconds",
    20  		Help:      "Time spent doing bos requests.",
    21  		Buckets:   []float64{.025, .05, .1, .25, .5, 1, 2},
    22  	}, []string{"operation", "status_code"}))
    23  )
    24  
    25  func init() {
    26  	bosRequestDuration.Register()
    27  }
    28  
    29  type objectStorageClient struct {
    30  	bos             *bos.Client
    31  	bucket          string
    32  	objectKeyPrefix string
    33  }
    34  
    35  // Configuration for Baidu Object Storage (BOS).
    36  type ObjectStorageConfig struct {
    37  	Endpoint        string `yaml:"endpoint"`
    38  	AccessKey       string `yaml:"ak"`
    39  	SecretKey       string `yaml:"sk"`
    40  	Bucket          string `yaml:"bucket"`
    41  	ObjectKeyPrefix string `yaml:"object_key_prefix"`
    42  }
    43  
    44  // New return bos-backed object client.
    45  func New(c ObjectStorageConfig) (chunk.ObjectClient, error) {
    46  	bosClient, err := bos.NewClient(c.AccessKey, c.SecretKey, c.Endpoint)
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  	return &objectStorageClient{
    51  		bos:             bosClient,
    52  		bucket:          c.Bucket,
    53  		objectKeyPrefix: c.ObjectKeyPrefix,
    54  	}, nil
    55  }
    56  
    57  // Stop implements chunk.ObjectClient interface.
    58  func (client *objectStorageClient) Stop() {
    59  	// nothing to do
    60  }
    61  
    62  // PutChunks implements chunk.ObjectClient interface, put chunks parallelly.
    63  func (client *objectStorageClient) PutChunks(ctx context.Context, chunks []chunk.Chunk) error {
    64  	bosChunkKeys := make([]string, 0, len(chunks))
    65  	bosChunkBufs := make([][]byte, 0, len(chunks))
    66  
    67  	for i := range chunks {
    68  		buf, err := chunks[i].Encoded()
    69  		if err != nil {
    70  			return err
    71  		}
    72  		key := client.objectKeyPrefix + chunks[i].ExternalKey()
    73  
    74  		bosChunkKeys = append(bosChunkKeys, key)
    75  		bosChunkBufs = append(bosChunkBufs, buf)
    76  	}
    77  
    78  	incomingErrors := make(chan error)
    79  	for i := range bosChunkBufs {
    80  		go func(i int) {
    81  			incomingErrors <- client.putChunk(ctx, bosChunkKeys[i], bosChunkBufs[i])
    82  		}(i)
    83  	}
    84  
    85  	var lastErr error
    86  	for range bosChunkKeys {
    87  		err := <-incomingErrors
    88  		if err != nil {
    89  			lastErr = err
    90  		}
    91  	}
    92  	return lastErr
    93  }
    94  
    95  func (client *objectStorageClient) putChunk(ctx context.Context, key string, buf []byte) error {
    96  	return instrument.CollectedRequest(ctx, "bos.PutObject", bosRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
    97  		_, err := client.bos.PutObjectFromBytes(client.bucket, key, buf, nil)
    98  		return err
    99  	})
   100  }
   101  
   102  // GutChunks implements chunk.ObjectClient interface, get chunks parallelly.
   103  func (client *objectStorageClient) GetChunks(ctx context.Context, chunks []chunk.Chunk) ([]chunk.Chunk, error) {
   104  	return util.GetParallelChunks(ctx, chunks, client.getChunk)
   105  }
   106  
   107  func (client *objectStorageClient) getChunk(ctx context.Context, decodeContext *chunk.DecodeContext, c chunk.Chunk) (chunk.Chunk, error) {
   108  	var result *bosAPI.GetObjectResult
   109  	key := client.objectKeyPrefix + c.ExternalKey()
   110  
   111  	err := instrument.CollectedRequest(ctx, "bos.GetObject", bosRequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   112  		var err error
   113  		result, err = client.bos.BasicGetObject(client.bucket, key)
   114  		return err
   115  	})
   116  	if err != nil {
   117  		return chunk.Chunk{}, err
   118  	}
   119  	defer result.Body.Close()
   120  
   121  	buf, err := ioutil.ReadAll(result.Body)
   122  	if err != nil {
   123  		return chunk.Chunk{}, err
   124  	}
   125  	if err := c.Decode(decodeContext, buf); err != nil {
   126  		return chunk.Chunk{}, err
   127  	}
   128  	return c, nil
   129  }