github.com/sequix/cortex@v1.1.6/pkg/chunk/aws/s3_storage_client.go (about)

     1  package aws
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  	"io/ioutil"
     8  	"strings"
     9  
    10  	"github.com/aws/aws-sdk-go/aws"
    11  	"github.com/aws/aws-sdk-go/aws/session"
    12  	"github.com/aws/aws-sdk-go/service/s3"
    13  	"github.com/aws/aws-sdk-go/service/s3/s3iface"
    14  	"github.com/davecgh/go-spew/spew"
    15  	"github.com/prometheus/client_golang/prometheus"
    16  
    17  	"github.com/sequix/cortex/pkg/chunk"
    18  	"github.com/sequix/cortex/pkg/chunk/util"
    19  	awscommon "github.com/weaveworks/common/aws"
    20  	"github.com/weaveworks/common/instrument"
    21  )
    22  
    23  var (
    24  	s3RequestDuration = instrument.NewHistogramCollector(prometheus.NewHistogramVec(prometheus.HistogramOpts{
    25  		Namespace: "cortex",
    26  		Name:      "s3_request_duration_seconds",
    27  		Help:      "Time spent doing S3 requests.",
    28  		Buckets:   []float64{.025, .05, .1, .25, .5, 1, 2},
    29  	}, []string{"operation", "status_code"}))
    30  )
    31  
    32  func init() {
    33  	s3RequestDuration.Register()
    34  }
    35  
    36  type s3ObjectClient struct {
    37  	bucketName string
    38  	S3         s3iface.S3API
    39  }
    40  
    41  // NewS3ObjectClient makes a new S3-backed ObjectClient.
    42  func NewS3ObjectClient(cfg StorageConfig, schemaCfg chunk.SchemaConfig) (chunk.ObjectClient, error) {
    43  	if cfg.S3.URL == nil {
    44  		return nil, fmt.Errorf("no URL specified for S3")
    45  	}
    46  	s3Config, err := awscommon.ConfigFromURL(cfg.S3.URL)
    47  	if err != nil {
    48  		return nil, err
    49  	}
    50  
    51  	s3Config = s3Config.WithS3ForcePathStyle(cfg.S3ForcePathStyle) // support for Path Style S3 url if has the flag
    52  
    53  	s3Config = s3Config.WithMaxRetries(0) // We do our own retries, so we can monitor them
    54  	s3Client := s3.New(session.New(s3Config))
    55  	bucketName := strings.TrimPrefix(cfg.S3.URL.Path, "/")
    56  	client := s3ObjectClient{
    57  		S3:         s3Client,
    58  		bucketName: bucketName,
    59  	}
    60  	return client, nil
    61  }
    62  
    63  func (a s3ObjectClient) Stop() {
    64  }
    65  
    66  func (a s3ObjectClient) GetChunks(ctx context.Context, chunks []chunk.Chunk) ([]chunk.Chunk, error) {
    67  	return util.GetParallelChunks(ctx, chunks, a.getChunk)
    68  }
    69  
    70  func (a s3ObjectClient) getChunk(ctx context.Context, decodeContext *chunk.DecodeContext, c chunk.Chunk) (chunk.Chunk, error) {
    71  	var resp *s3.GetObjectOutput
    72  	err := instrument.CollectedRequest(ctx, "S3.GetObject", s3RequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
    73  		var err error
    74  		resp, err = a.S3.GetObjectWithContext(ctx, &s3.GetObjectInput{
    75  			Bucket: aws.String(a.bucketName),
    76  			Key:    aws.String(c.ExternalKey()),
    77  		})
    78  		return err
    79  	})
    80  	if err != nil {
    81  		return chunk.Chunk{}, err
    82  	}
    83  	defer resp.Body.Close()
    84  	buf, err := ioutil.ReadAll(resp.Body)
    85  	if err != nil {
    86  		return chunk.Chunk{}, err
    87  	}
    88  	if err := c.Decode(decodeContext, buf); err != nil {
    89  		return chunk.Chunk{}, err
    90  	}
    91  	return c, nil
    92  }
    93  
    94  func (a s3ObjectClient) PutChunks(ctx context.Context, chunks []chunk.Chunk) error {
    95  	var (
    96  		s3ChunkKeys []string
    97  		s3ChunkBufs [][]byte
    98  	)
    99  
   100  	for i := range chunks {
   101  		buf, err := chunks[i].Encoded()
   102  		if err != nil {
   103  			return err
   104  		}
   105  		key := chunks[i].ExternalKey()
   106  
   107  		s3ChunkKeys = append(s3ChunkKeys, key)
   108  		s3ChunkBufs = append(s3ChunkBufs, buf)
   109  	}
   110  
   111  	incomingErrors := make(chan error)
   112  	for i := range s3ChunkBufs {
   113  		go func(i int) {
   114  			incomingErrors <- a.putS3Chunk(ctx, s3ChunkKeys[i], s3ChunkBufs[i])
   115  		}(i)
   116  	}
   117  
   118  	var lastErr error
   119  	for range s3ChunkKeys {
   120  		err := <-incomingErrors
   121  		if err != nil {
   122  			lastErr = err
   123  		}
   124  	}
   125  	return lastErr
   126  }
   127  
   128  func (a s3ObjectClient) putS3Chunk(ctx context.Context, key string, buf []byte) error {
   129  	return instrument.CollectedRequest(ctx, "S3.PutObject", s3RequestDuration, instrument.ErrorCode, func(ctx context.Context) error {
   130  		rsp, err := a.S3.PutObjectWithContext(ctx, &s3.PutObjectInput{
   131  			Body:   bytes.NewReader(buf),
   132  			Bucket: aws.String(a.bucketName),
   133  			Key:    aws.String(key),
   134  		})
   135  		spew.Dump(rsp, err)
   136  		return err
   137  	})
   138  }