manifest: replace marked for compaction annotator with a set

The marked-for-compaction annotator is an overcomplicated way to
implement an ordered set. It does not fit neatly in the annotator
pattern - unlike all other annotators, it can become invalidated. It
also requires separate logic to maintain a count of marked files.

We remove this annotator and replace it with an ordered set (using the
google btree) stored in the Version. This is a lot cleaner and also
addresses a limitation of the annotator: we were only able to obtain
one candidate (and thus only one rewrite compaction at any time).

Author: RaduBerinde

compaction_picker.go

@@ -1524,7 +1524,7 @@ func (p *compactionPickerByScore) pickAutoNonScore(env compactionEnv) (pc picked
 	// MarkedForCompaction field is persisted in the manifest. That's okay. We
 	// previously would've ignored the designation, whereas now we'll re-compact
 	// the file in place.
-	if p.vers.Stats.MarkedForCompaction > 0 {
+	if p.vers.MarkedForCompaction.Count() > 0 {
 		if pc := p.pickRewriteCompaction(env); pc != nil {
 			return pc
 		}
@@ -1597,19 +1597,6 @@ var elisionOnlyAnnotator = manifest.NewTableAnnotator[manifest.TableMetadata](ma
 	},
 })
 
-// markedForCompactionAnnotator is a manifest.TableAnnotator that annotates B-Tree
-// nodes with the *fileMetadata of a file that is marked for compaction
-// within the subtree. If multiple files meet the criteria, it chooses
-// whichever file has the lowest LargestSeqNum.
-var markedForCompactionAnnotator = manifest.NewTableAnnotator[manifest.TableMetadata](manifest.PickFileAggregator{
-	Filter: func(f *manifest.TableMetadata) (eligible bool, cacheOK bool) {
-		return f.MarkedForCompaction, true
-	},
-	Compare: func(f1 *manifest.TableMetadata, f2 *manifest.TableMetadata) bool {
-		return f1.LargestSeqNum < f2.LargestSeqNum
-	},
-})
-
 // pickedCompactionFromCandidateFile creates a pickedCompaction from a *fileMetadata
 // with various checks to ensure that the file still exists in the expected level
 // and isn't already being compacted.
@@ -1680,17 +1667,8 @@ func (p *compactionPickerByScore) pickElisionOnlyCompaction(
 func (p *compactionPickerByScore) pickRewriteCompaction(
 	env compactionEnv,
 ) (pc *pickedTableCompaction) {
-	if p.vers.Stats.MarkedForCompaction == 0 {
-		return nil
-	}
-	for l := numLevels - 1; l >= 0; l-- {
-		candidate := markedForCompactionAnnotator.LevelAnnotation(p.vers.Levels[l])
-		if candidate == nil {
-			// Try the next level.
-			continue
-		}
-		pc := p.pickedCompactionFromCandidateFile(candidate, env, l, l, compactionKindRewrite)
-		if pc != nil {
+	for candidate, level := range p.vers.MarkedForCompaction.Ascending() {
+		if pc := p.pickedCompactionFromCandidateFile(candidate, env, level, level, compactionKindRewrite); pc != nil {
 			return pc
 		}
 	}

compaction_picker_test.go

@@ -690,9 +690,9 @@ func TestCompactionPickerL0(t *testing.T) {
 					if f.TableNum != base.TableNum(tableNum) {
 						continue
 					}
+					// This code should be identical to the one in DB.markFilesLocked().
 					f.MarkedForCompaction = true
-					picker.vers.Stats.MarkedForCompaction++
-					markedForCompactionAnnotator.InvalidateLevelAnnotation(picker.vers.Levels[l])
+					picker.vers.MarkedForCompaction.Insert(f, l)
 					return fmt.Sprintf("marked L%d.%s", l, f.TableNum)
 				}
 			}

compaction_test.go

@@ -2789,23 +2789,19 @@ func TestMarkedForCompaction(t *testing.T) {
 		},
 		Logger: testutils.Logger{T: t},
 	}
+	opts.Experimental.CompactionScheduler = func() CompactionScheduler {
+		return NewConcurrencyLimitSchedulerWithNoPeriodicGrantingForTest()
+	}
 	opts.WithFSDefaults()
 
-	reset := func() {
-		if d != nil {
-			require.NoError(t, d.Close())
-		}
-		mem = vfs.NewMem()
-		require.NoError(t, mem.MkdirAll("ext", 0755))
-
-		var err error
-		d, err = Open("", opts)
-		require.NoError(t, err)
-	}
 	datadriven.RunTest(t, "testdata/marked_for_compaction", func(t *testing.T, td *datadriven.TestData) string {
+		buf.Reset()
 		switch td.Cmd {
-		case "reset":
-			reset()
+		case "reopen":
+			require.NoError(t, d.Close())
+			var err error
+			d, err = Open("", opts)
+			require.NoError(t, err)
 			return ""
 
 		case "define":
@@ -2829,23 +2825,31 @@ func TestMarkedForCompaction(t *testing.T) {
 			return s
 
 		case "mark-for-compaction":
-			d.mu.Lock()
-			defer d.mu.Unlock()
-			vers := d.mu.versions.currentVersion()
 			var tableNum uint64
 			td.ScanArgs(t, "file", &tableNum)
-			for l, lm := range vers.Levels {
-				for f := range lm.All() {
-					if f.TableNum != base.TableNum(tableNum) {
-						continue
+
+			findFn := func(v *manifest.Version) (found bool, files [numLevels][]*manifest.TableMetadata, _ error) {
+				for l, lm := range v.Levels {
+					for f := range lm.All() {
+						if f.TableNum == base.TableNum(tableNum) {
+							files[l] = append(files[l], f)
+							fmt.Fprintf(&buf, "marked L%d.%s", l, f.TableNum)
+							return true, files, nil
+						}
 					}
-					f.MarkedForCompaction = true
-					vers.Stats.MarkedForCompaction++
-					markedForCompactionAnnotator.InvalidateLevelAnnotation(vers.Levels[l])
-					return fmt.Sprintf("marked L%d.%s", l, f.TableNum)
 				}
+				return false, files, nil
+			}
+			d.mu.Lock()
+			defer d.mu.Unlock()
+			if err := d.markFilesForCompactionLocked(findFn); err != nil {
+				td.Fatalf(t, "markFilesLocked: %s", err)
 			}
-			return "not-found"
+
+			if buf.Len() == 0 {
+				return "not-found"
+			}
+			return buf.String()
 
 		case "maybe-compact":
 			d.mu.Lock()
@@ -2854,11 +2858,11 @@ func TestMarkedForCompaction(t *testing.T) {
 			d.maybeScheduleCompaction()
 			for d.mu.compact.compactingCount > 0 {
 				d.mu.compact.cond.Wait()
+				d.maybeScheduleCompaction()
 			}
 
 			fmt.Fprintln(&buf, d.mu.versions.currentVersion().DebugString())
 			s := strings.TrimSpace(buf.String())
-			buf.Reset()
 			opts.DisableAutomaticCompactions = true
 			return s
 

db.go

@@ -2061,7 +2061,7 @@ func (d *DB) Metrics() *Metrics {
 	metrics.Compact.InProgressBytes = d.mu.versions.atomicInProgressBytes.Load()
 	// TODO(radu): split this to separate the download compactions.
 	metrics.Compact.NumInProgress = int64(d.mu.compact.compactingCount + d.mu.compact.downloadingCount)
-	metrics.Compact.MarkedFiles = vers.Stats.MarkedForCompaction
+	metrics.Compact.MarkedFiles = vers.MarkedForCompaction.Count()
 	metrics.Compact.Duration = d.mu.compact.duration
 	for c := range d.mu.compact.inProgress {
 		if !c.IsFlush() {

format_major_version.go

@@ -536,7 +536,7 @@ func (d *DB) writeFormatVersionMarker(formatVers FormatMajorVersion) error {
 // waiting for compactions to complete (or for slots to free up).
 func (d *DB) compactMarkedFilesLocked() error {
 	curr := d.mu.versions.currentVersion()
-	if curr.Stats.MarkedForCompaction == 0 {
+	if curr.MarkedForCompaction.Count() == 0 {
 		return nil
 	}
 	// Attempt to schedule a compaction to rewrite a file marked for compaction.
@@ -552,7 +552,7 @@ func (d *DB) compactMarkedFilesLocked() error {
 	// compaction.  Or compaction of the file might have already been in
 	// progress. In any scenario, wait until there's some change in the
 	// state of active compactions.
-	for curr.Stats.MarkedForCompaction > 0 {
+	for curr.MarkedForCompaction.Count() > 0 {
 		// Before waiting, check that the database hasn't been closed. Trying to
 		// schedule the compaction may have dropped d.mu while waiting for a
 		// manifest write to complete. In that dropped interim, the database may
@@ -568,7 +568,7 @@ func (d *DB) compactMarkedFilesLocked() error {
 
 		// Only wait on compactions if there are files still marked for compaction.
 		// NB: Waiting on this condition variable drops d.mu while blocked.
-		if curr.Stats.MarkedForCompaction > 0 {
+		if curr.MarkedForCompaction.Count() > 0 {
 			// NB: we cannot assert that d.mu.compact.compactingCount > 0, since
 			// with a CompactionScheduler a DB may not have even one ongoing
 			// compaction (if other competing activities are being preferred by the
@@ -586,13 +586,9 @@ func (d *DB) compactMarkedFilesLocked() error {
 // level.
 type findFilesFunc func(v *manifest.Version) (found bool, files [numLevels][]*manifest.TableMetadata, _ error)
 
-// This method is not used currently, but it will be useful the next time we need
-// to mark files for compaction.
-var _ = (*DB)(nil).markFilesLocked
-
-// markFilesLocked durably marks the files that match the given findFilesFunc for
+// markFilesForCompactionLocked durably marks the files that match the given findFilesFunc for
 // compaction.
-func (d *DB) markFilesLocked(findFn findFilesFunc) error {
+func (d *DB) markFilesForCompactionLocked(findFn findFilesFunc) error {
 	jobID := d.newJobIDLocked()
 
 	// Acquire a read state to have a view of the LSM and a guarantee that none
@@ -634,32 +630,22 @@ func (d *DB) markFilesLocked(findFn findFilesFunc) error {
 	// been re-acquired by the defer within the above anonymous function.
 	_, err = d.mu.versions.UpdateVersionLocked(func() (versionUpdate, error) {
 		vers := d.mu.versions.currentVersion()
-		for l, filesToMark := range files {
-			if len(filesToMark) == 0 {
-				continue
-			}
+		for level, filesToMark := range files {
 			for _, f := range filesToMark {
 				// Ignore files to be marked that have already been compacted or marked.
 				if f.CompactionState == manifest.CompactionStateCompacted ||
 					f.MarkedForCompaction {
 					continue
 				}
 				// Else, mark the file for compaction in this version.
-				vers.Stats.MarkedForCompaction++
 				f.MarkedForCompaction = true
+				// We are modifying the current version in-place (so that the updated
+				// set is reflected in the "base" version in the new manifest). This is
+				// ok because we are holding the DB lock and all code that uses the
+				// MarkedForCompaction set runs under the DB lock.
+				// TODO(radu): find a less sketchy way to do this.
+				vers.MarkedForCompaction.Insert(f, level)
 			}
-			// The compaction picker uses the markedForCompactionAnnotator to
-			// quickly find files marked for compaction, or to quickly determine
-			// that there are no such files marked for compaction within a level.
-			// A b-tree node may be annotated with an annotation recording that
-			// there are no files marked for compaction within the node's subtree,
-			// based on the assumption that it's static.
-			//
-			// Since we're marking files for compaction, these b-tree nodes'
-			// annotations will be out of date. Clear the compaction-picking
-			// annotation, so that it's recomputed the next time the compaction
-			// picker looks for a file marked for compaction.
-			markedForCompactionAnnotator.InvalidateLevelAnnotation(vers.Levels[l])
 		}
 		// The 'marked-for-compaction' bit is persisted in the MANIFEST file
 		// metadata. We've already modified the in-memory table metadata, but the

go.mod

@@ -14,6 +14,7 @@ require (
 	github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06
 	github.com/ghemawat/stream v0.0.0-20171120220530-696b145b53b9
 	github.com/golang/snappy v0.0.5-0.20231225225746-43d5d4cd4e0e
+	github.com/google/btree v1.1.3
 	github.com/guptarohit/asciigraph v0.5.5
 	github.com/klauspost/compress v1.17.11
 	github.com/kr/pretty v0.3.1

go.sum

@@ -108,6 +108,8 @@ github.com/gonum/internal v0.0.0-20181124074243-f884aa714029/go.mod h1:Pu4dmpkhS
 github.com/gonum/lapack v0.0.0-20181123203213-e4cdc5a0bff9/go.mod h1:XA3DeT6rxh2EAE789SSiSJNqxPaC0aE9J8NTOI0Jo/A=
 github.com/gonum/matrix v0.0.0-20181209220409-c518dec07be9/go.mod h1:0EXg4mc1CNP0HCqCz+K4ts155PXIlUywf0wqN+GfPZw=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
+github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=
+github.com/google/btree v1.1.3/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.5.4/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=

internal/manifest/annotator.go

@@ -258,18 +258,6 @@ func (a *TableAnnotator[T]) accumulateRangeAnnotation(
 	return dst
 }
 
-// InvalidateAnnotation removes any existing cached annotations from this
-// annotator from a node's subtree.
-func (a *TableAnnotator[T]) invalidateNodeAnnotation(n *node[*TableMetadata]) {
-	annot := a.findAnnotation(n)
-	annot.valid.Store(false)
-	if !n.isLeaf() {
-		for i := int16(0); i <= n.count; i++ {
-			a.invalidateNodeAnnotation(n.child(i))
-		}
-	}
-}
-
 // LevelAnnotation calculates the annotation defined by this TableAnnotator for all
 // files in the given LevelMetadata. A pointer to the TableAnnotator is used as the
 // key for pre-calculated values, so the same TableAnnotator must be used to avoid
@@ -337,20 +325,6 @@ func (a *TableAnnotator[T]) VersionRangeAnnotation(v *Version, bounds base.UserK
 	return dst
 }
 
-// InvalidateLevelAnnotation clears any cached annotations defined by TableAnnotator.
-// A pointer to the TableAnnotator is used as the key for pre-calculated values, so
-// the same TableAnnotator must be used to clear the appropriate cached annotation.
-// Calls to InvalidateLevelAnnotation are *not* concurrent-safe with any other
-// calls to TableAnnotator methods for the same TableAnnotator (concurrent calls from
-// other annotators are fine). Any calls to this function must have some
-// externally-guaranteed mutual exclusion.
-func (a *TableAnnotator[T]) InvalidateLevelAnnotation(lm LevelMetadata) {
-	if lm.Empty() {
-		return
-	}
-	a.invalidateNodeAnnotation(lm.tree.root)
-}
-
 // Annotation calculates the annotation defined by this BlobFileAnnotator for
 // all blob files in the given set.
 // A pointer to the TableAnnotator is used as the key for pre-calculated values,

internal/manifest/marked_for_compaction.go

@@ -0,0 +1,89 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package manifest
+
+import (
+	"iter"
+
+	"github.com/cockroachdb/errors"
+	"github.com/cockroachdb/pebble/internal/invariants"
+	googlebtree "github.com/google/btree"
+)
+
+// MarkedForCompactionSet implements an ordered set of tables that are marked
+// for compaction. They are ordered by decreasing LSM level, and within each
+// level by increasing LargestSeqNum.
+//
+// Files can be marked for compaction when upgrading a DB to a specific version.
+type MarkedForCompactionSet struct {
+	tree *googlebtree.BTreeG[tableAndLevel]
+}
+
+type tableAndLevel struct {
+	meta  *TableMetadata
+	level int
+}
+
+func markedForCompactionLessFn(a, b tableAndLevel) bool {
+	if a.level != b.level {
+		return a.level > b.level
+	}
+	if a.meta.LargestSeqNum != b.meta.LargestSeqNum {
+		return a.meta.LargestSeqNum < b.meta.LargestSeqNum
+	}
+	return a.meta.TableNum < b.meta.TableNum
+}
+
+// Count returns the number of tables in the set.
+func (s *MarkedForCompactionSet) Count() int {
+	if s.tree == nil {
+		return 0
+	}
+	return s.tree.Len()
+}
+
+// Insert adds a table to the set. The table must not be in the set already.
+func (s *MarkedForCompactionSet) Insert(meta *TableMetadata, level int) {
+	if s.tree == nil {
+		s.tree = googlebtree.NewG[tableAndLevel](8, markedForCompactionLessFn)
+	}
+	_, existed := s.tree.ReplaceOrInsert(tableAndLevel{meta: meta, level: level})
+	if invariants.Enabled && existed {
+		panic(errors.AssertionFailedf("table %s already in MarkedForCompaction", meta.TableNum))
+	}
+}
+
+// Delete removes a table from the set. The table must be in the set.
+func (s *MarkedForCompactionSet) Delete(meta *TableMetadata, level int) {
+	ok := false
+	if s.tree != nil {
+		_, ok = s.tree.Delete(tableAndLevel{meta: meta, level: level})
+	}
+	if invariants.Enabled && !ok {
+		panic(errors.AssertionFailedf("table %s not in MarkedForCompaction", meta.TableNum))
+	}
+}
+
+// Clone the set. The resulting set can be independently modified.
+func (s *MarkedForCompactionSet) Clone() MarkedForCompactionSet {
+	if s.tree == nil {
+		return MarkedForCompactionSet{}
+	}
+	return MarkedForCompactionSet{tree: s.tree.Clone()}
+}
+
+// Ascending returns an iterator which produces the tables marked for compaction
+// (along with the level), ordered by decreasing level and within each level by
+// increasing LargestSeqNum.
+func (s *MarkedForCompactionSet) Ascending() iter.Seq2[*TableMetadata, int] {
+	return func(yield func(*TableMetadata, int) bool) {
+		if s.tree == nil {
+			return
+		}
+		s.tree.Ascend(func(t tableAndLevel) bool {
+			return yield(t.meta, t.level)
+		})
+	}
+}