db: consolidate the file size annotators

We currently have three file size annotators which are used for
`EstimateDiskUsage`. Partially overlapping remote files are inspected
twice and external remote files three times. It also means 3x more
annotation-related objects per B-Tree node.

This change consolidates the three aggregators into a single one. We
generalize `SumAggregator` to allow arbitrary value types.

Author: RaduBerinde

db.go

@@ -505,17 +505,7 @@ type DB struct {
 			// validating is set to true when validation is running.
 			validating bool
 		}
-
-		// annotators contains various instances of manifest.TableAnnotator which
-		// should be protected from concurrent access.
-		annotators struct {
-			// totalFileSize is the sum of the size of all files in the
-			// database. This includes local, remote, and external sstables --
-			// along with blob files.
-			totalFileSize *manifest.TableAnnotator[uint64]
-			remoteSize    *manifest.TableAnnotator[uint64]
-			externalSize  *manifest.TableAnnotator[uint64]
-		}
+		fileSizeAnnotator *manifest.TableAnnotator[fileSizeByBacking]
 	}
 
 	// problemSpans keeps track of spans of keys within LSM levels where

disk_usage.go

@@ -48,40 +48,67 @@ func (d *DB) EstimateDiskUsageByBackingType(
 	readState := d.loadReadState()
 	defer readState.unref()
 
-	totalSize = *d.mu.annotators.totalFileSize.VersionRangeAnnotation(readState.current, bounds)
-	remoteSize = *d.mu.annotators.remoteSize.VersionRangeAnnotation(readState.current, bounds)
-	externalSize = *d.mu.annotators.externalSize.VersionRangeAnnotation(readState.current, bounds)
+	sizes := d.mu.fileSizeAnnotator.VersionRangeAnnotation(readState.current, bounds)
+	return sizes.totalSize, sizes.remoteSize, sizes.externalSize, nil
+}
 
-	return
+// fileSizeByBacking contains the estimated file size for LSM data within some
+// bounds. It is broken down by backing type. The file size refers to both the
+// sstable size and an estimate of the referenced blob sizes.
+type fileSizeByBacking struct {
+	// totalSize is the estimated size of all files for the given bounds.
+	totalSize uint64
+	// remoteSize is the estimated size of remote files for the given bounds.
+	remoteSize uint64
+	// externalSize is the estimated size of external files for the given bounds.
+	externalSize uint64
 }
 
-// makeFileSizeAnnotator returns an annotator that computes the total
-// storage size of files that meet some criteria defined by filter. When
-// applicable, this includes both the sstable size and the size of any
-// referenced blob files.
-func (d *DB) makeFileSizeAnnotator(
-	filter func(f *manifest.TableMetadata) bool,
-) *manifest.TableAnnotator[uint64] {
-	return manifest.NewTableAnnotator[uint64](manifest.SumAggregator{
-		AccumulateFunc: func(f *manifest.TableMetadata) (uint64, bool) {
-			if filter(f) {
-				return f.Size + f.EstimatedReferenceSize(), true
-			}
-			return 0, true
+func (d *DB) singleFileSizeByBacking(
+	fileSize uint64, t *manifest.TableMetadata,
+) (_ fileSizeByBacking, ok bool) {
+	res := fileSizeByBacking{
+		totalSize: fileSize,
+	}
+
+	objMeta, err := d.objProvider.Lookup(base.FileTypeTable, t.TableBacking.DiskFileNum)
+	if err != nil {
+		return res, false
+	}
+	if objMeta.IsRemote() {
+		res.remoteSize += fileSize
+		if objMeta.IsExternal() {
+			res.externalSize += fileSize
+		}
+	}
+	return res, true
+}
+
+// makeFileSizeAnnotator returns an annotator that computes the storage size of
+// files. When applicable, this includes both the sstable size and the size of
+// any referenced blob files.
+func (d *DB) makeFileSizeAnnotator() *manifest.TableAnnotator[fileSizeByBacking] {
+	return manifest.NewTableAnnotator[fileSizeByBacking](manifest.SumAggregator[fileSizeByBacking]{
+		AddFunc: func(src, dst *fileSizeByBacking) {
+			dst.totalSize += src.totalSize
+			dst.remoteSize += src.remoteSize
+			dst.externalSize += src.externalSize
+		},
+		AccumulateFunc: func(f *manifest.TableMetadata) (v fileSizeByBacking, cacheOK bool) {
+			return d.singleFileSizeByBacking(f.Size+f.EstimatedReferenceSize(), f)
 		},
-		AccumulatePartialOverlapFunc: func(f *manifest.TableMetadata, bounds base.UserKeyBounds) uint64 {
-			if filter(f) {
-				overlappingFileSize, err := d.fileCache.estimateSize(f, bounds.Start, bounds.End.Key)
-				if err != nil {
-					return 0
-				}
-				overlapFraction := float64(overlappingFileSize) / float64(f.Size)
-				// Scale the blob reference size proportionally to the file
-				// overlap from the bounds to approximate only the blob
-				// references that overlap with the requested bounds.
-				return overlappingFileSize + uint64(float64(f.EstimatedReferenceSize())*overlapFraction)
+		AccumulatePartialOverlapFunc: func(f *manifest.TableMetadata, bounds base.UserKeyBounds) fileSizeByBacking {
+			overlappingFileSize, err := d.fileCache.estimateSize(f, bounds.Start, bounds.End.Key)
+			if err != nil {
+				return fileSizeByBacking{}
 			}
-			return 0
+			overlapFraction := float64(overlappingFileSize) / float64(f.Size)
+			// Scale the blob reference size proportionally to the file
+			// overlap from the bounds to approximate only the blob
+			// references that overlap with the requested bounds.
+			size := overlappingFileSize + uint64(float64(f.EstimatedReferenceSize())*overlapFraction)
+			res, _ := d.singleFileSizeByBacking(size, f)
+			return res
 		},
 	})
 }

internal/manifest/annotator.go

@@ -333,44 +333,46 @@ func (a *BlobFileAnnotator[T]) Annotation(blobFiles *BlobFileSet) *T {
 	return v
 }
 
-// SumAggregator defines an Aggregator which sums together a uint64 value
+// SumAggregator defines an Aggregator which sums together a T value
 // across files.
-type SumAggregator struct {
-	AccumulateFunc               func(f *TableMetadata) (v uint64, cacheOK bool)
-	AccumulatePartialOverlapFunc func(f *TableMetadata, bounds base.UserKeyBounds) uint64
+type SumAggregator[T any] struct {
+	AddFunc                      func(src, dst *T)
+	AccumulateFunc               func(f *TableMetadata) (v T, cacheOK bool)
+	AccumulatePartialOverlapFunc func(f *TableMetadata, bounds base.UserKeyBounds) T
 }
 
 // Zero implements AnnotationAggregator.Zero, returning a new uint64 set to 0.
-func (sa SumAggregator) Zero() *uint64 {
-	return new(uint64)
+func (sa SumAggregator[T]) Zero() *T {
+	return new(T)
 }
 
 // Accumulate implements AnnotationAggregator.Accumulate, accumulating a single
 // file's uint64 value.
-func (sa SumAggregator) Accumulate(f *TableMetadata, dst *uint64) (v *uint64, cacheOK bool) {
+func (sa SumAggregator[T]) Accumulate(f *TableMetadata, dst *T) (v *T, cacheOK bool) {
 	accumulated, ok := sa.AccumulateFunc(f)
-	*dst += accumulated
+	sa.AddFunc(&accumulated, dst)
 	return dst, ok
 }
 
 // AccumulatePartialOverlap implements
 // PartialOverlapAnnotationAggregator.AccumulatePartialOverlap, accumulating a
 // single file's uint64 value for a file which only partially overlaps with the
 // range defined by bounds.
-func (sa SumAggregator) AccumulatePartialOverlap(
-	f *TableMetadata, dst *uint64, bounds base.UserKeyBounds,
-) *uint64 {
+func (sa SumAggregator[T]) AccumulatePartialOverlap(
+	f *TableMetadata, dst *T, bounds base.UserKeyBounds,
+) *T {
 	if sa.AccumulatePartialOverlapFunc == nil {
 		v, _ := sa.Accumulate(f, dst)
 		return v
 	}
-	*dst += sa.AccumulatePartialOverlapFunc(f, bounds)
+	accumulated := sa.AccumulatePartialOverlapFunc(f, bounds)
+	sa.AddFunc(&accumulated, dst)
 	return dst
 }
 
 // Merge implements AnnotationAggregator.Merge by summing two uint64 values.
-func (sa SumAggregator) Merge(src *uint64, dst *uint64) *uint64 {
-	*dst += *src
+func (sa SumAggregator[T]) Merge(src *T, dst *T) *T {
+	sa.AddFunc(src, dst)
 	return dst
 }
 
@@ -380,7 +382,10 @@ func (sa SumAggregator) Merge(src *uint64, dst *uint64) *uint64 {
 func SumAnnotator(
 	accumulate func(f *TableMetadata) (v uint64, cacheOK bool),
 ) *TableAnnotator[uint64] {
-	return NewTableAnnotator[uint64](SumAggregator{AccumulateFunc: accumulate})
+	return NewTableAnnotator[uint64](SumAggregator[uint64]{
+		AddFunc:        func(src, dst *uint64) { *dst += *src },
+		AccumulateFunc: accumulate,
+	})
 }
 
 // NumFilesAnnotator is an TableAnnotator which computes an annotation value

open.go

@@ -442,23 +442,7 @@ func Open(dirname string, opts *Options) (db *DB, err error) {
 	d.newIters = d.fileCache.newIters
 	d.tableNewRangeKeyIter = tableNewRangeKeyIter(d.newIters)
 
-	d.mu.annotators.totalFileSize = d.makeFileSizeAnnotator(func(f *manifest.TableMetadata) bool {
-		return true
-	})
-	d.mu.annotators.remoteSize = d.makeFileSizeAnnotator(func(f *manifest.TableMetadata) bool {
-		meta, err := d.objProvider.Lookup(base.FileTypeTable, f.TableBacking.DiskFileNum)
-		if err != nil {
-			return false
-		}
-		return meta.IsRemote()
-	})
-	d.mu.annotators.externalSize = d.makeFileSizeAnnotator(func(f *manifest.TableMetadata) bool {
-		meta, err := d.objProvider.Lookup(base.FileTypeTable, f.TableBacking.DiskFileNum)
-		if err != nil {
-			return false
-		}
-		return meta.IsRemote() && meta.Remote.CleanupMethod == objstorage.SharedNoCleanup
-	})
+	d.mu.fileSizeAnnotator = d.makeFileSizeAnnotator()
 
 	var previousOptionsFileNum base.DiskFileNum
 	var previousOptionsFilename string