github.com/grafana/pyroscope@v1.18.0/pkg/metastore/compaction/compactor/compactor.go (about)

     1  package compactor
     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/api/gen/proto/go/metastore/v1/raft_log"
    12  	"github.com/grafana/pyroscope/pkg/iter"
    13  	"github.com/grafana/pyroscope/pkg/metastore/compaction"
    14  	"github.com/grafana/pyroscope/pkg/metastore/compaction/compactor/store"
    15  )
    16  
    17  var (
    18  	_ compaction.Compactor = (*Compactor)(nil)
    19  	_ compaction.Planner   = (*Compactor)(nil)
    20  )
    21  
    22  type Tombstones interface {
    23  	ListTombstones(before time.Time) iter.Iterator[*metastorev1.Tombstones]
    24  }
    25  
    26  type BlockQueueStore interface {
    27  	StoreEntry(*bbolt.Tx, compaction.BlockEntry) error
    28  	DeleteEntry(tx *bbolt.Tx, index uint64, id string) error
    29  	ListEntries(*bbolt.Tx) iter.Iterator[compaction.BlockEntry]
    30  	CreateBuckets(*bbolt.Tx) error
    31  }
    32  
    33  type Compactor struct {
    34  	config     Config
    35  	queue      *compactionQueue
    36  	store      BlockQueueStore
    37  	tombstones Tombstones
    38  }
    39  
    40  func NewCompactor(
    41  	config Config,
    42  	store BlockQueueStore,
    43  	tombstones Tombstones,
    44  	reg prometheus.Registerer,
    45  ) *Compactor {
    46  	queue := newCompactionQueue(config, reg)
    47  	return &Compactor{
    48  		config:     config,
    49  		queue:      queue,
    50  		store:      store,
    51  		tombstones: tombstones,
    52  	}
    53  }
    54  
    55  func NewStore() *store.BlockQueueStore {
    56  	return store.NewBlockQueueStore()
    57  }
    58  
    59  func (c *Compactor) Compact(tx *bbolt.Tx, entry compaction.BlockEntry) error {
    60  	if int(entry.Level) >= len(c.config.Levels) {
    61  		return nil
    62  	}
    63  	if err := c.store.StoreEntry(tx, entry); err != nil {
    64  		return err
    65  	}
    66  	c.enqueue(entry)
    67  	return nil
    68  }
    69  
    70  func (c *Compactor) enqueue(e compaction.BlockEntry) bool {
    71  	return c.queue.push(e)
    72  }
    73  
    74  func (c *Compactor) NewPlan(cmd *raft.Log) compaction.Plan {
    75  	now := cmd.AppendedAt.UnixNano()
    76  	before := cmd.AppendedAt.Add(-c.config.CleanupDelay)
    77  	tombstones := c.tombstones.ListTombstones(before)
    78  	return &plan{
    79  		compactor:  c,
    80  		tombstones: tombstones,
    81  		blocks:     newBlockIter(),
    82  		now:        now,
    83  	}
    84  }
    85  
    86  func (c *Compactor) UpdatePlan(tx *bbolt.Tx, plan *raft_log.CompactionPlanUpdate) error {
    87  	for _, job := range plan.NewJobs {
    88  		// Delete source blocks from the compaction queue.
    89  		k := compactionKey{
    90  			tenant: job.Plan.Tenant,
    91  			shard:  job.Plan.Shard,
    92  			level:  job.Plan.CompactionLevel,
    93  		}
    94  		staged := c.queue.blockQueue(k.level).stagedBlocks(k)
    95  		for _, b := range job.Plan.SourceBlocks {
    96  			e := staged.delete(b)
    97  			if e == zeroBlockEntry {
    98  				continue
    99  			}
   100  			if err := c.store.DeleteEntry(tx, e.index, e.id); err != nil {
   101  				return err
   102  			}
   103  		}
   104  	}
   105  
   106  	return nil
   107  }
   108  
   109  func (c *Compactor) Init(tx *bbolt.Tx) error {
   110  	return c.store.CreateBuckets(tx)
   111  }
   112  
   113  func (c *Compactor) Restore(tx *bbolt.Tx) error {
   114  	// Reset in-memory state before loading entries from the store.
   115  	c.queue.reset()
   116  	entries := c.store.ListEntries(tx)
   117  	defer func() {
   118  		_ = entries.Close()
   119  	}()
   120  	for entries.Next() {
   121  		c.enqueue(entries.At())
   122  	}
   123  	return entries.Err()
   124  }