github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/bucketindex/storage.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/storage/tsdb/bucketindex/storage.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package bucketindex
     7  
     8  import (
     9  	"bytes"
    10  	"compress/gzip"
    11  	"context"
    12  	"encoding/json"
    13  	"time"
    14  
    15  	"github.com/go-kit/log"
    16  	"github.com/grafana/dskit/runutil"
    17  	"github.com/pkg/errors"
    18  
    19  	"github.com/grafana/pyroscope/pkg/objstore"
    20  )
    21  
    22  var (
    23  	ErrIndexNotFound  = errors.New("bucket index not found")
    24  	ErrIndexCorrupted = errors.New("bucket index corrupted")
    25  )
    26  
    27  // ReadIndex reads, parses and returns a bucket index from the bucket.
    28  // ReadIndex has a one-minute timeout for completing the read against the bucket.
    29  // One minute is hard-coded to a reasonably high value to protect against operations that can take unbounded time.
    30  func ReadIndex(ctx context.Context, bkt objstore.Bucket, userID string, cfgProvider objstore.TenantConfigProvider, logger log.Logger) (*Index, error) {
    31  	ctx, cancel := context.WithTimeout(ctx, time.Minute)
    32  	defer cancel()
    33  
    34  	userBkt := objstore.NewTenantBucketClient(userID, bkt, cfgProvider)
    35  
    36  	// Get the bucket index.
    37  	reader, err := userBkt.WithExpectedErrs(userBkt.IsObjNotFoundErr).Get(ctx, IndexCompressedFilename)
    38  	if err != nil {
    39  		if userBkt.IsObjNotFoundErr(err) {
    40  			return nil, ErrIndexNotFound
    41  		}
    42  		return nil, errors.Wrap(err, "read bucket index")
    43  	}
    44  	defer runutil.CloseWithLogOnErr(logger, reader, "close bucket index reader")
    45  
    46  	// Read all the content.
    47  	gzipReader, err := gzip.NewReader(reader)
    48  	if err != nil {
    49  		return nil, ErrIndexCorrupted
    50  	}
    51  	defer runutil.CloseWithLogOnErr(logger, gzipReader, "close bucket index gzip reader")
    52  
    53  	// Deserialize it.
    54  	index := &Index{}
    55  	d := json.NewDecoder(gzipReader)
    56  	if err := d.Decode(index); err != nil {
    57  		return nil, ErrIndexCorrupted
    58  	}
    59  
    60  	return index, nil
    61  }
    62  
    63  // WriteIndex uploads the provided index to the storage.
    64  func WriteIndex(ctx context.Context, bkt objstore.Bucket, userID string, cfgProvider objstore.TenantConfigProvider, idx *Index) error {
    65  	bkt = objstore.NewTenantBucketClient(userID, bkt, cfgProvider)
    66  
    67  	// Marshal the index.
    68  	content, err := json.Marshal(idx)
    69  	if err != nil {
    70  		return errors.Wrap(err, "marshal bucket index")
    71  	}
    72  
    73  	// Compress it.
    74  	var gzipContent bytes.Buffer
    75  	gzip := gzip.NewWriter(&gzipContent)
    76  	gzip.Name = IndexFilename
    77  
    78  	if _, err := gzip.Write(content); err != nil {
    79  		return errors.Wrap(err, "gzip bucket index")
    80  	}
    81  	if err := gzip.Close(); err != nil {
    82  		return errors.Wrap(err, "close gzip bucket index")
    83  	}
    84  
    85  	// Upload the index to the storage.
    86  	if err := bkt.Upload(ctx, IndexCompressedFilename, &gzipContent); err != nil {
    87  		return errors.Wrap(err, "upload bucket index")
    88  	}
    89  
    90  	return nil
    91  }
    92  
    93  // DeleteIndex deletes the bucket index from the storage. No error is returned if the index
    94  // does not exist.
    95  func DeleteIndex(ctx context.Context, bkt objstore.Bucket, userID string, cfgProvider objstore.TenantConfigProvider) error {
    96  	bkt = objstore.NewTenantBucketClient(userID, bkt, cfgProvider)
    97  
    98  	err := bkt.Delete(ctx, IndexCompressedFilename)
    99  	if err != nil && !bkt.IsObjNotFoundErr(err) {
   100  		return errors.Wrap(err, "delete bucket index")
   101  	}
   102  	return nil
   103  }