Fix another case of compaction job state inconsistency

grafana · Jul 17, 2024 · 32a5391 · 32a5391
1 parent d4f7916
commit 32a5391
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 5 deletions.
diff --git a/pkg/metastore/metastore_compaction_planner.go b/pkg/metastore/metastore_compaction_planner.go
@@ -161,7 +161,7 @@ func (m *metastoreState) tryCreateJob(block *metastorev1.BlockMeta) *compactionp
 	return job
 }
 
-func (m *metastoreState) addCompactionJob(job *compactionpb.CompactionJob, compactionLevel uint32) {
+func (m *metastoreState) addCompactionJob(job *compactionpb.CompactionJob) {
 	key := tenantShard{
 		tenant: job.TenantId,
 		shard:  job.Shard,
@@ -171,7 +171,7 @@ func (m *metastoreState) addCompactionJob(job *compactionpb.CompactionJob, compa
 	defer plan.jobsMutex.Unlock()
 
 	plan.jobsByName[job.Name] = job
-	plan.queuedBlocksByLevel[compactionLevel] = plan.queuedBlocksByLevel[compactionLevel][:0]
+	plan.queuedBlocksByLevel[job.CompactionLevel] = plan.queuedBlocksByLevel[job.CompactionLevel][:0]
 }
 
 func (m *metastoreState) addBlockToCompactionJobQueue(block *metastorev1.BlockMeta) {

diff --git a/pkg/metastore/metastore_state_add_block.go b/pkg/metastore/metastore_state_add_block.go
@@ -7,6 +7,7 @@ import (
 	"go.etcd.io/bbolt"
 
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
+	"github.com/grafana/pyroscope/pkg/metastore/compactionpb"
 )
 
 func (m *Metastore) AddBlock(_ context.Context, req *metastorev1.AddBlockRequest) (*metastorev1.AddBlockResponse, error) {
@@ -22,6 +23,10 @@ func (m *metastoreState) applyAddBlock(request *metastorev1.AddBlockRequest) (*m
 	if err != nil {
 		return nil, err
 	}
+
+	var jobToAdd *compactionpb.CompactionJob
+	var blockToAddToQueue *metastorev1.BlockMeta
+
 	err = m.db.boltdb.Update(func(tx *bbolt.Tx) error {
 		err := updateBlockMetadataBucket(tx, name, func(bucket *bbolt.Bucket) error {
 			return bucket.Put(key, value)
@@ -40,9 +45,9 @@ func (m *metastoreState) applyAddBlock(request *metastorev1.AddBlockRequest) (*m
 			if err != nil {
 				return err
 			}
-			m.addCompactionJob(job, request.Block.CompactionLevel)
+			jobToAdd = job
 		} else {
-			m.addBlockToCompactionJobQueue(request.Block)
+			blockToAddToQueue = request.Block
 		}
 		return nil
 	})
@@ -55,5 +60,10 @@ func (m *metastoreState) applyAddBlock(request *metastorev1.AddBlockRequest) (*m
 		return nil, err
 	}
 	m.getOrCreateShard(request.Block.Shard).putSegment(request.Block)
+	if jobToAdd != nil {
+		m.addCompactionJob(jobToAdd)
+	} else if blockToAddToQueue != nil {
+		m.addBlockToCompactionJobQueue(blockToAddToQueue)
+	}
 	return &metastorev1.AddBlockResponse{}, nil
 }
diff --git a/pkg/metastore/metastore_state_poll_compaction_jobs.go b/pkg/metastore/metastore_state_poll_compaction_jobs.go
@@ -95,7 +95,7 @@ func (m *metastoreState) applyPollCompactionJobsStatus(request *compactorv1.Poll
 	}
 
 	for _, j := range jResult.newJobs {
-		m.addCompactionJob(j, j.CompactionLevel)
+		m.addCompactionJob(j)
 	}
 
 	for _, b := range jResult.newQueuedBlocks {