compact: avoid cleaner deadlock on cancellation

Signed-off-by: yuchen-db <yw3642@gmail.com>

Author: yw3642

pkg/compact/blocks_cleaner.go

@@ -20,6 +20,8 @@ import (
 	"github.com/thanos-io/thanos/pkg/errutil"
 )
 
+const deleteMarkedBlocksConcurrency = 32
+
 // BlocksCleaner is a struct that deletes blocks from bucket which are marked for deletion.
 type BlocksCleaner struct {
 	logger                   log.Logger
@@ -45,8 +47,6 @@ func NewBlocksCleaner(logger log.Logger, bkt objstore.Bucket, ignoreDeletionMark
 // DeleteMarkedBlocks uses ignoreDeletionMarkFilter to gather the blocks that are marked for deletion and deletes those
 // if older than given deleteDelay.
 func (s *BlocksCleaner) DeleteMarkedBlocks(ctx context.Context) (map[ulid.ULID]struct{}, error) {
-	const conc = 32
-
 	level.Info(s.logger).Log("msg", "started cleaning of blocks marked for deletion")
 
 	var (
@@ -55,10 +55,10 @@ func (s *BlocksCleaner) DeleteMarkedBlocks(ctx context.Context) (map[ulid.ULID]s
 		deletedBlocks    = make(map[ulid.ULID]struct{}, 0)
 		deletionMarkMap  = s.ignoreDeletionMarkFilter.DeletionMarkBlocks()
 		wg               sync.WaitGroup
-		dm               = make(chan *metadata.DeletionMark, conc)
+		dm               = make(chan *metadata.DeletionMark, deleteMarkedBlocksConcurrency)
 	)
 
-	for range conc {
+	for range deleteMarkedBlocksConcurrency {
 		wg.Go(func() {
 			for deletionMark := range dm {
 				if ctx.Err() != nil {
@@ -82,8 +82,13 @@ func (s *BlocksCleaner) DeleteMarkedBlocks(ctx context.Context) (map[ulid.ULID]s
 		})
 	}
 
+enqueue:
 	for _, deletionMark := range deletionMarkMap {
-		dm <- deletionMark
+		select {
+		case dm <- deletionMark:
+		case <-ctx.Done():
+			break enqueue
+		}
 	}
 	close(dm)
 	wg.Wait()

pkg/compact/compact_test.go

@@ -749,6 +749,54 @@ func TestGarbageCollect_FilterRace(t *testing.T) {
 	}
 }
 
+func TestBlocksCleanerCancellationWhileEnqueuing(t *testing.T) {
+	bkt := objstore.WithNoopInstr(objstore.NewInMemBucket())
+	deletionMarkFilter := block.NewIgnoreDeletionMarkFilter(log.NewNopLogger(), bkt, 0, 1)
+	deletionMarkCount := 2*deleteMarkedBlocksConcurrency + 1
+	metas := make(map[ulid.ULID]*metadata.Meta, deletionMarkCount)
+
+	for i := range deletionMarkCount {
+		id := ulid.MustNew(uint64(i+1), bytes.NewReader(make([]byte, 10)))
+		metas[id] = &metadata.Meta{}
+
+		var buf bytes.Buffer
+		testutil.Ok(t, json.NewEncoder(&buf).Encode(&metadata.DeletionMark{
+			ID:           id,
+			DeletionTime: 0,
+			Version:      1,
+		}))
+		testutil.Ok(t, bkt.Upload(context.Background(), path.Join(id.String(), metadata.DeletionMarkFilename), &buf))
+	}
+
+	synced := promauto.With(nil).NewGaugeVec(prometheus.GaugeOpts{Name: "test_blocks_cleaner_synced"}, []string{"state"})
+	testutil.Ok(t, deletionMarkFilter.Filter(context.Background(), metas, synced, nil))
+
+	cleaner := NewBlocksCleaner(
+		log.NewNopLogger(),
+		bkt,
+		deletionMarkFilter,
+		0,
+		promauto.With(nil).NewCounter(prometheus.CounterOpts{Name: "test_blocks_cleaner_cleaned"}),
+		promauto.With(nil).NewCounter(prometheus.CounterOpts{Name: "test_blocks_cleaner_failures"}),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	done := make(chan error, 1)
+	go func() {
+		_, err := cleaner.DeleteMarkedBlocks(ctx)
+		done <- err
+	}()
+
+	select {
+	case err := <-done:
+		testutil.NotOk(t, err)
+		testutil.Assert(t, errors.Is(err, context.Canceled), "expected context cancellation, got %v", err)
+	case <-time.After(time.Second):
+		t.Fatal("DeleteMarkedBlocks blocked after cancellation")
+	}
+}
+
 func TestCompactExtensions(t *testing.T) {
 	t.Parallel()