manifest: add range annotations

This change adds a "range annotation" feature to Annotators , which
are computations that aggregate some value over a specific key range within
within a level. Level-wide annotations are now computed internally as a
range annotation with a key range spanning the whole level. Range annotations
use the same B-tree caching behavior as regular annotations, so queries
remain fast even with thousands of tables because they avoid a sequential
iteration over a level's files.

This PR only sets up range annotations without changing any existing
behavior. See #3793 for some potential use cases.

`BenchmarkNumFilesRangeAnnotation` shows that range annotations are
significantly faster than using `version.Overlaps` to aggregate over
a key range:
```
pkg: github.com/cockroachdb/pebble/internal/manifest
BenchmarkNumFilesRangeAnnotation/annotator-10         	  232282	      4716 ns/op	     112 B/op	       7 allocs/op
BenchmarkNumFilesRangeAnnotation/overlaps-10          	    2110	    545482 ns/op	     400 B/op	       9 allocs/op
```

internal/manifest/annotator.go

@@ -4,6 +4,12 @@
 
 package manifest
 
+import (
+	"sort"
+
+	"github.com/cockroachdb/pebble/internal/base"
+)
+
 // The Annotator type defined below is used by other packages to lazily
 // compute a value over a B-Tree. Each node of the B-Tree stores one
 // `annotation` per annotator, containing the result of the computation over
@@ -59,6 +65,11 @@ type annotation struct {
 	// AnnotationAggregator.Accumulate or AnnotationAggregator.Merge.
 	// NB: This is untyped for the same reason as annotator above.
 	v interface{}
+
+	// scratch is used to hold the aggregated annotation value when computing
+	// range annotations in order to avoid overwriting an already cached
+	// annotation, and to also avoid additional allocations.
+	scratch interface{}
 	// valid indicates whether future reads of the annotation may use the
 	// value as-is. If false, v will be zeroed and recalculated.
 	valid bool
@@ -81,39 +92,114 @@ func (a *Annotator[T]) findAnnotation(n *node) *annotation {
 	return &n.annot[len(n.annot)-1]
 }
 
-// nodeAnnotation computes this annotator's annotation of this node across all
-// files in the node's subtree. The second return value indicates whether the
-// annotation is stable and thus cacheable.
-func (a *Annotator[T]) nodeAnnotation(n *node) (_ *T, cacheOK bool) {
+// nodeRangeAnnotation computes this annotator's annotation of a node across
+// all files in the range defined by lowerBound and upperBound. The second
+// return value indicates whether the annotation is stable and thus cacheable.
+func (a *Annotator[T]) nodeRangeAnnotation(
+	n *node,
+	cmp base.Compare,
+	// lowerBound and upperBound may be nil to indicate no lower or upper bound.
+	lowerBound []byte,
+	// upperBound is a UserKeyBoundary that may be inclusive or exclusive.
+	upperBound *base.UserKeyBoundary,
+) (v *T, cacheOK bool) {
 	annot := a.findAnnotation(n)
-	t := annot.v.(*T)
-	// If the annotation is already marked as valid, we can return it without
+	// If the annotation is already marked as valid and this node's
+	// subtree is fully within the bounds, we can return it without
 	// recomputing anything.
-	if annot.valid {
-		return t, true
+	if lowerBound == nil && upperBound == nil && annot.valid {
+		return annot.v.(*T), true
+	}
+
+	// We will accumulate annotations from each item in the end-exclusive
+	// range [leftItem, rightItem).
+	leftItem, rightItem := 0, int(n.count)
+	if lowerBound != nil {
+		// leftItem is the index of the first item that overlaps the lower bound.
+		leftItem = sort.Search(int(n.count), func(i int) bool {
+			return cmp(lowerBound, n.items[i].Largest.UserKey) <= 0
+		})
+	}
+	if upperBound != nil {
+		// rightItem is the index of the first item that does not overlap the
+		// upper bound.
+		rightItem = sort.Search(int(n.count), func(i int) bool {
+			return !upperBound.IsUpperBoundFor(cmp, n.items[i].Smallest.UserKey)
+		})
+	}
+
+	var result *T
+	switch {
+	// If there is no cached annotation, we can directly write to the node's
+	// annotation value.
+	case !annot.valid:
+		result = a.Aggregator.Zero(annot.v.(*T))
+	// Otherwise, use annot.scratch as scratch space to avoid allocations.
+	// The allocation for annot.scratch is performed lazily here instead of
+	// within findAnnotation to avoid an allocation when range annotations
+	// are not used.
+	case annot.scratch == nil:
+		annot.scratch = a.Aggregator.Zero(nil)
+		result = annot.scratch.(*T)
+	default:
+		result = a.Aggregator.Zero(annot.scratch.(*T))
 	}
 
-	t = a.Aggregator.Zero(t)
 	valid := true
+	// Accumulate annotations from every item that overlaps the bounds.
+	for i := leftItem; i < rightItem; i++ {
+		v, ok := a.Aggregator.Accumulate(n.items[i], result)
+		result = v
+		valid = valid && ok
+	}
 
-	for i := int16(0); i <= n.count; i++ {
-		if !n.leaf {
-			v, ok := a.nodeAnnotation(n.children[i])
-			t = a.Aggregator.Merge(v, t)
-			valid = valid && ok
+	if !n.leaf {
+		// We will accumulate annotations from each child in the end-inclusive
+		// range [leftChild, rightChild].
+		leftChild, rightChild := leftItem, rightItem
+		// If the lower bound overlaps with the child at leftItem, there is no
+		// need to accumulate annotations from the child to its left.
+		if leftItem < int(n.count) && cmp(lowerBound, n.items[leftItem].Smallest.UserKey) >= 0 {
+			leftChild++
 		}
+		// If the upper bound spans beyond the child at rightItem, we must also
+		// accumulate annotations from the child to its right.
+		if rightItem < int(n.count) && upperBound.IsUpperBoundFor(cmp, n.items[rightItem].Largest.UserKey) {
+			rightChild++
+		}
+
+		for i := leftChild; i <= rightChild; i++ {
+			newLowerBound, newUpperBound := lowerBound, upperBound
+			// If this child is to the right of leftItem, then its entire
+			// subtree is within the lower bound.
+			if i > leftItem {
+				newLowerBound = nil
+			}
+			// If this child is to the left of rightItem, then its entire
+			// subtree is within the upper bound.
+			if i < rightItem {
+				newUpperBound = nil
+			}
 
-		if i < n.count {
-			var ok bool
-			t, ok = a.Aggregator.Accumulate(n.items[i], t)
+			v, ok := a.nodeRangeAnnotation(
+				n.children[i],
+				cmp,
+				newLowerBound,
+				newUpperBound,
+			)
+			result = a.Aggregator.Merge(v, result)
 			valid = valid && ok
 		}
 	}
 
-	annot.v = t
-	annot.valid = valid
+	// Update this node's cached annotation only if we accumulated from its
+	// entire subtree.
+	if lowerBound == nil && upperBound == nil {
+		annot.v = result
+		annot.valid = valid
+	}
 
-	return t, annot.valid
+	return result, valid
 }
 
 // InvalidateAnnotation removes any existing cached annotations from this
@@ -138,7 +224,7 @@ func (a *Annotator[T]) LevelAnnotation(lm LevelMetadata) *T {
 		return a.Aggregator.Zero(nil)
 	}
 
-	v, _ := a.nodeAnnotation(lm.tree.root)
+	v, _ := a.nodeRangeAnnotation(lm.tree.root, lm.tree.cmp, nil, nil)
 	return v
 }
 
@@ -158,6 +244,21 @@ func (a *Annotator[T]) MultiLevelAnnotation(lms []LevelMetadata) *T {
 	return aggregated
 }
 
+// LevelRangeAnnotation calculates the annotation defined by this Annotator for
+// the files within LevelMetadata which are within the range
+// [lowerBound, upperBound). A pointer to the Annotator is used as the key for
+// pre-calculated values, so the same Annotator must be used to avoid duplicate
+// computation. Annotation must not be called concurrently, and in practice this
+// is achieved by requiring callers to hold DB.mu.
+func (a *Annotator[T]) LevelRangeAnnotation(lm LevelMetadata, bounds *base.UserKeyBounds) *T {
+	if lm.Empty() {
+		return a.Aggregator.Zero(nil)
+	}
+
+	v, _ := a.nodeRangeAnnotation(lm.tree.root, lm.tree.cmp, bounds.Start, &bounds.End)
+	return v
+}
+
 // InvalidateAnnotation clears any cached annotations defined by Annotator. A
 // pointer to the Annotator is used as the key for pre-calculated values, so
 // the same Annotator must be used to clear the appropriate cached annotation.
@@ -206,6 +307,14 @@ func SumAnnotator(accumulate func(f *FileMetadata) (v uint64, cacheOK bool)) *An
 	}
 }
 
+// NumFilesAnnotator is an Annotator which computes an annotation value
+// equal to the number of files included in the annotation. Particularly, it
+// can be used to efficiently calculate the number of files in a given key
+// range using range annotations.
+var NumFilesAnnotator = SumAnnotator(func(f *FileMetadata) (uint64, bool) {
+	return 1, true
+})
+
 // PickFileAggregator implements the AnnotationAggregator interface. It defines
 // an aggregator that picks a single file from a set of eligible files.
 type PickFileAggregator struct {