github.com/grafana/pyroscope@v1.18.0/pkg/metastore/index/store/index_store.go (about)

     1  package store
     2  
     3  import (
     4  	"encoding/binary"
     5  	"errors"
     6  	"fmt"
     7  	goiter "iter"
     8  
     9  	"go.etcd.io/bbolt"
    10  	bbolterrors "go.etcd.io/bbolt/errors"
    11  )
    12  
    13  const (
    14  	partitionBucketName   = "partition"
    15  	emptyTenantBucketName = "-"
    16  )
    17  
    18  var (
    19  	partitionBucketNameBytes   = []byte(partitionBucketName)
    20  	emptyTenantBucketNameBytes = []byte(emptyTenantBucketName)
    21  )
    22  
    23  type IndexStore struct{}
    24  
    25  func tenantBucketName(tenant string) []byte {
    26  	if tenant == "" {
    27  		return emptyTenantBucketNameBytes
    28  	}
    29  	return []byte(tenant)
    30  }
    31  
    32  func getPartitionsBucket(tx *bbolt.Tx) *bbolt.Bucket {
    33  	return tx.Bucket(partitionBucketNameBytes)
    34  }
    35  
    36  func getOrCreateSubBucket(parent *bbolt.Bucket, name []byte) (*bbolt.Bucket, error) {
    37  	bucket := parent.Bucket(name)
    38  	if bucket == nil {
    39  		return parent.CreateBucket(name)
    40  	}
    41  	return bucket, nil
    42  }
    43  
    44  func NewIndexStore() *IndexStore {
    45  	return &IndexStore{}
    46  }
    47  
    48  func (m *IndexStore) CreateBuckets(tx *bbolt.Tx) error {
    49  	_, err := tx.CreateBucketIfNotExists(partitionBucketNameBytes)
    50  	return err
    51  }
    52  
    53  func (m *IndexStore) Partitions(tx *bbolt.Tx) goiter.Seq[Partition] {
    54  	root := getPartitionsBucket(tx)
    55  	if root == nil {
    56  		return func(func(Partition) bool) {}
    57  	}
    58  	return func(yield func(Partition) bool) {
    59  		cursor := root.Cursor()
    60  		for partitionKey, _ := cursor.First(); partitionKey != nil; partitionKey, _ = cursor.Next() {
    61  			p := Partition{}
    62  			if err := p.UnmarshalBinary(partitionKey); err != nil {
    63  				continue
    64  			}
    65  			if !yield(p) {
    66  				return
    67  			}
    68  		}
    69  	}
    70  }
    71  
    72  func (m *IndexStore) LoadShard(tx *bbolt.Tx, p Partition, tenant string, shard uint32) (*Shard, error) {
    73  	s, err := loadTenantShard(tx, p, tenant, shard)
    74  	if err != nil {
    75  		return nil, fmt.Errorf("error loading tenant shard %s/%d partition %q: %w", tenant, shard, p, err)
    76  	}
    77  	return s, nil
    78  }
    79  
    80  func (m *IndexStore) DeleteShard(tx *bbolt.Tx, p Partition, tenant string, shard uint32) error {
    81  	partitions := getPartitionsBucket(tx)
    82  	partitionKey := p.Bytes()
    83  	if partition := partitions.Bucket(partitionKey); partition != nil {
    84  		tenantKey := tenantBucketName(tenant)
    85  		if shards := partition.Bucket(tenantKey); shards != nil {
    86  			if err := shards.DeleteBucket(binary.BigEndian.AppendUint32(nil, shard)); err != nil {
    87  				if !errors.Is(err, bbolterrors.ErrBucketNotFound) {
    88  					return err
    89  				}
    90  			}
    91  			if isBucketEmpty(shards) {
    92  				if err := partition.DeleteBucket(tenantKey); err != nil {
    93  					if !errors.Is(err, bbolterrors.ErrBucketNotFound) {
    94  						return err
    95  					}
    96  				}
    97  			}
    98  		}
    99  		if isBucketEmpty(partition) {
   100  			if err := partitions.DeleteBucket(partitionKey); err != nil {
   101  				if !errors.Is(err, bbolterrors.ErrBucketNotFound) {
   102  					return err
   103  				}
   104  			}
   105  		}
   106  	}
   107  	return nil
   108  }
   109  
   110  func isBucketEmpty(bucket *bbolt.Bucket) bool {
   111  	if bucket == nil {
   112  		return true
   113  	}
   114  	c := bucket.Cursor()
   115  	k, _ := c.First()
   116  	return k == nil
   117  }