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

     1  package store
     2  
     3  import (
     4  	"encoding/binary"
     5  	"errors"
     6  	"fmt"
     7  
     8  	"go.etcd.io/bbolt"
     9  
    10  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    11  	"github.com/grafana/pyroscope/pkg/iter"
    12  	"github.com/grafana/pyroscope/pkg/metastore/store"
    13  )
    14  
    15  var ErrInvalidTombstoneEntry = errors.New("invalid tombstone entry")
    16  
    17  var tombstoneBucketName = []byte("tombstones")
    18  
    19  type TombstoneEntry struct {
    20  	// Key the entry was stored under. This is needed for backward
    21  	// compatibility: if the logic for generating keys changes, we
    22  	// will still be able to delete old entries.
    23  	key        []byte
    24  	Index      uint64
    25  	AppendedAt int64
    26  	*metastorev1.Tombstones
    27  }
    28  
    29  func (e TombstoneEntry) Name() string {
    30  	switch {
    31  	case e.Blocks != nil:
    32  		return e.Blocks.Name
    33  	case e.Shard != nil:
    34  		return e.Shard.Name
    35  	default:
    36  		return ""
    37  	}
    38  }
    39  
    40  type TombstoneStore struct{ bucketName []byte }
    41  
    42  func NewTombstoneStore() *TombstoneStore {
    43  	return &TombstoneStore{bucketName: tombstoneBucketName}
    44  }
    45  
    46  func (s *TombstoneStore) CreateBuckets(tx *bbolt.Tx) error {
    47  	_, err := tx.CreateBucketIfNotExists(s.bucketName)
    48  	return err
    49  }
    50  
    51  func (s *TombstoneStore) StoreTombstones(tx *bbolt.Tx, entry TombstoneEntry) error {
    52  	kv := marshalTombstoneEntry(entry)
    53  	return tx.Bucket(s.bucketName).Put(kv.Key, kv.Value)
    54  }
    55  
    56  func (s *TombstoneStore) DeleteTombstones(tx *bbolt.Tx, entry TombstoneEntry) error {
    57  	return tx.Bucket(s.bucketName).Delete(marshalTombstoneEntryKey(entry))
    58  }
    59  
    60  func (s *TombstoneStore) ListEntries(tx *bbolt.Tx) iter.Iterator[TombstoneEntry] {
    61  	return newTombstoneEntriesIterator(tx.Bucket(s.bucketName))
    62  }
    63  
    64  type tombstoneEntriesIterator struct {
    65  	iter *store.CursorIterator
    66  	cur  TombstoneEntry
    67  	err  error
    68  }
    69  
    70  func newTombstoneEntriesIterator(bucket *bbolt.Bucket) *tombstoneEntriesIterator {
    71  	return &tombstoneEntriesIterator{iter: store.NewCursorIter(bucket.Cursor())}
    72  }
    73  
    74  func (x *tombstoneEntriesIterator) Next() bool {
    75  	if x.err != nil || !x.iter.Next() {
    76  		return false
    77  	}
    78  	x.err = unmarshalTombstoneEntry(&x.cur, x.iter.At())
    79  	return x.err == nil
    80  }
    81  
    82  func (x *tombstoneEntriesIterator) At() TombstoneEntry { return x.cur }
    83  
    84  func (x *tombstoneEntriesIterator) Close() error { return x.iter.Close() }
    85  
    86  func (x *tombstoneEntriesIterator) Err() error {
    87  	if err := x.iter.Err(); err != nil {
    88  		return err
    89  	}
    90  	return x.err
    91  }
    92  
    93  func marshalTombstoneEntry(e TombstoneEntry) store.KV {
    94  	k := marshalTombstoneEntryKey(e)
    95  	b := make([]byte, e.SizeVT())
    96  	_, _ = e.MarshalToSizedBufferVT(b)
    97  	return store.KV{Key: k, Value: b}
    98  }
    99  
   100  func marshalTombstoneEntryKey(e TombstoneEntry) []byte {
   101  	if e.key != nil {
   102  		b := make([]byte, len(e.key))
   103  		copy(b, e.key)
   104  		return b
   105  	}
   106  	name := e.Name()
   107  	b := make([]byte, 16+len(name))
   108  	binary.BigEndian.PutUint64(b[0:8], e.Index)
   109  	binary.BigEndian.PutUint64(b[8:16], uint64(e.AppendedAt))
   110  	copy(b[16:], name)
   111  	return b
   112  }
   113  
   114  func unmarshalTombstoneEntry(e *TombstoneEntry, kv store.KV) error {
   115  	if len(kv.Key) < 16 {
   116  		return ErrInvalidTombstoneEntry
   117  	}
   118  	e.key = make([]byte, len(kv.Key))
   119  	copy(e.key, kv.Key)
   120  	e.Index = binary.BigEndian.Uint64(kv.Key[0:8])
   121  	e.AppendedAt = int64(binary.BigEndian.Uint64(kv.Key[8:16]))
   122  	e.Tombstones = new(metastorev1.Tombstones)
   123  	if err := e.UnmarshalVT(kv.Value); err != nil {
   124  		return fmt.Errorf("%w: %w", ErrInvalidTombstoneEntry, err)
   125  	}
   126  	return nil
   127  }