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

     1  package tombstones
     2  
     3  import (
     4  	"time"
     5  
     6  	"github.com/hashicorp/raft"
     7  	"github.com/prometheus/client_golang/prometheus"
     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/index/tombstones/store"
    13  )
    14  
    15  type TombstoneStore interface {
    16  	StoreTombstones(*bbolt.Tx, store.TombstoneEntry) error
    17  	DeleteTombstones(*bbolt.Tx, store.TombstoneEntry) error
    18  	ListEntries(*bbolt.Tx) iter.Iterator[store.TombstoneEntry]
    19  	CreateBuckets(*bbolt.Tx) error
    20  }
    21  
    22  type Tombstones struct {
    23  	metrics    *metrics
    24  	tombstones map[tombstoneKey]*tombstones
    25  	blocks     map[tenantBlockKey]*tenantBlocks
    26  	queue      *tombstoneQueue
    27  	store      TombstoneStore
    28  }
    29  
    30  type tenantBlockKey struct {
    31  	tenant string
    32  	shard  uint32
    33  }
    34  
    35  type tenantBlocks struct {
    36  	blocks map[string]struct{}
    37  }
    38  
    39  func NewTombstones(store TombstoneStore, reg prometheus.Registerer) *Tombstones {
    40  	return &Tombstones{
    41  		metrics:    newMetrics(reg),
    42  		tombstones: make(map[tombstoneKey]*tombstones),
    43  		blocks:     make(map[tenantBlockKey]*tenantBlocks),
    44  		queue:      newTombstoneQueue(),
    45  		store:      store,
    46  	}
    47  }
    48  
    49  func NewStore() *store.TombstoneStore {
    50  	return store.NewTombstoneStore()
    51  }
    52  
    53  func (x *Tombstones) Exists(tenant string, shard uint32, block string) bool {
    54  	t, exists := x.blocks[tenantBlockKey{tenant: tenant, shard: shard}]
    55  	if exists {
    56  		_, exists = t.blocks[block]
    57  	}
    58  	return exists
    59  }
    60  
    61  func (x *Tombstones) ListTombstones(before time.Time) iter.Iterator[*metastorev1.Tombstones] {
    62  	return &tombstoneIter{
    63  		head:   x.queue.head,
    64  		before: before.UnixNano(),
    65  	}
    66  }
    67  
    68  func (x *Tombstones) AddTombstones(tx *bbolt.Tx, cmd *raft.Log, t *metastorev1.Tombstones) error {
    69  	var k tombstoneKey
    70  	if !k.set(t) {
    71  		return nil
    72  	}
    73  	v := store.TombstoneEntry{
    74  		Index:      cmd.Index,
    75  		AppendedAt: cmd.AppendedAt.UnixNano(),
    76  		Tombstones: t,
    77  	}
    78  	if !x.put(k, v) {
    79  		return nil
    80  	}
    81  	return x.store.StoreTombstones(tx, v)
    82  }
    83  
    84  func (x *Tombstones) DeleteTombstones(tx *bbolt.Tx, cmd *raft.Log, tombstones ...*metastorev1.Tombstones) error {
    85  	for _, t := range tombstones {
    86  		if err := x.deleteTombstones(tx, cmd, t); err != nil {
    87  			return err
    88  		}
    89  	}
    90  	return nil
    91  }
    92  
    93  func (x *Tombstones) deleteTombstones(tx *bbolt.Tx, _ *raft.Log, t *metastorev1.Tombstones) error {
    94  	var k tombstoneKey
    95  	if !k.set(t) {
    96  		return nil
    97  	}
    98  	e := x.delete(k)
    99  	if e == nil {
   100  		return nil
   101  	}
   102  	return x.store.DeleteTombstones(tx, e.TombstoneEntry)
   103  }
   104  
   105  func (x *Tombstones) put(k tombstoneKey, v store.TombstoneEntry) bool {
   106  	if _, found := x.tombstones[k]; found {
   107  		return false
   108  	}
   109  	e := &tombstones{TombstoneEntry: v}
   110  	x.tombstones[k] = e
   111  	x.queue.push(e)
   112  	x.metrics.incrementTombstones(v.Tombstones)
   113  	if v.Blocks != nil {
   114  		// Keep track of the blocks we enqueued. This is
   115  		// necessary to answer if the block has already
   116  		// been deleted (within a limited time window).
   117  		x.putBlockTombstones(v.Blocks)
   118  	}
   119  	return true
   120  }
   121  
   122  func (x *Tombstones) delete(k tombstoneKey) (t *tombstones) {
   123  	e, found := x.tombstones[k]
   124  	if !found {
   125  		return nil
   126  	}
   127  	delete(x.tombstones, k)
   128  	x.metrics.decrementTombstones(e.Tombstones)
   129  	if t = x.queue.delete(e); t != nil {
   130  		if t.Blocks != nil {
   131  			x.deleteBlockTombstones(t.Blocks)
   132  		}
   133  	}
   134  	return t
   135  }
   136  
   137  func (x *Tombstones) putBlockTombstones(t *metastorev1.BlockTombstones) {
   138  	bk := tenantBlockKey{
   139  		tenant: t.Tenant,
   140  		shard:  t.Shard,
   141  	}
   142  	m, ok := x.blocks[bk]
   143  	if !ok {
   144  		m = &tenantBlocks{blocks: make(map[string]struct{})}
   145  		x.blocks[bk] = m
   146  	}
   147  	for _, b := range t.Blocks {
   148  		m.blocks[b] = struct{}{}
   149  	}
   150  }
   151  
   152  func (x *Tombstones) deleteBlockTombstones(t *metastorev1.BlockTombstones) {
   153  	bk := tenantBlockKey{
   154  		tenant: t.Tenant,
   155  		shard:  t.Shard,
   156  	}
   157  	m, found := x.blocks[bk]
   158  	if !found {
   159  		return
   160  	}
   161  	for _, b := range t.Blocks {
   162  		delete(m.blocks, b)
   163  	}
   164  }
   165  
   166  func (x *Tombstones) Init(tx *bbolt.Tx) error {
   167  	return x.store.CreateBuckets(tx)
   168  }
   169  
   170  func (x *Tombstones) Restore(tx *bbolt.Tx) error {
   171  	x.queue = newTombstoneQueue()
   172  	clear(x.tombstones)
   173  	clear(x.blocks)
   174  	x.metrics.tombstones.Reset()
   175  	entries := x.store.ListEntries(tx)
   176  	defer func() {
   177  		_ = entries.Close()
   178  	}()
   179  	for entries.Next() {
   180  		var k tombstoneKey
   181  		if v := entries.At(); k.set(v.Tombstones) {
   182  			x.put(k, v)
   183  		}
   184  	}
   185  	return entries.Err()
   186  }