From 4f062cb7188edd4343c2487c61f70f474dee2fb5 Mon Sep 17 00:00:00 2001
From: Anish Shanbhag <anish.shanbhag@cockroachlabs.com>
Date: Tue, 30 Jul 2024 11:36:43 -0400
Subject: [PATCH] manifest: fix incorrect Annotator pointer aggregation

cockroachdb#3760 contained a bug which causes Annotator
values to be incorrectly aggregated when pointer values
should be overwritten. This is because the `vtyped`
variable was not being modified due to being on the
stack.

This change fixes this and adds a unit test for
`PickFileAggregator` to catch this issue in the future.

cockroachdb#3759 should already not be affected by
this due to the different way it handles aggregation.
---
 internal/manifest/annotator.go      |  9 ++++-----
 internal/manifest/annotator_test.go | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/internal/manifest/annotator.go b/internal/manifest/annotator.go
index cfbe2f4f16..b9fd80ec54 100644
--- a/internal/manifest/annotator.go
+++ b/internal/manifest/annotator.go
@@ -86,25 +86,24 @@ func (a *Annotator[T]) findAnnotation(n *node) *annotation {
 // annotation is stable and thus cacheable.
 func (a *Annotator[T]) nodeAnnotation(n *node) (v *T, cacheOK bool) {
 	annot := a.findAnnotation(n)
-	vtyped := annot.v.(*T)
 	// If the annotation is already marked as valid, we can return it without
 	// recomputing anything.
 	if annot.valid {
-		return vtyped, true
+		return annot.v.(*T), true
 	}
 
-	annot.v = a.Aggregator.Zero(vtyped)
+	annot.v = a.Aggregator.Zero(annot.v.(*T))
 	annot.valid = true
 
 	for i := int16(0); i <= n.count; i++ {
 		if !n.leaf {
 			v, ok := a.nodeAnnotation(n.children[i])
-			annot.v = a.Aggregator.Merge(v, vtyped)
+			annot.v = a.Aggregator.Merge(v, annot.v.(*T))
 			annot.valid = annot.valid && ok
 		}
 
 		if i < n.count {
-			v, ok := a.Aggregator.Accumulate(n.items[i], vtyped)
+			v, ok := a.Aggregator.Accumulate(n.items[i], annot.v.(*T))
 			annot.v = v
 			annot.valid = annot.valid && ok
 		}
diff --git a/internal/manifest/annotator_test.go b/internal/manifest/annotator_test.go
index 1765188b99..07badf6a38 100644
--- a/internal/manifest/annotator_test.go
+++ b/internal/manifest/annotator_test.go
@@ -56,3 +56,26 @@ func BenchmarkNumFilesAnnotator(b *testing.B) {
 		require.EqualValues(b, uint64(i), numFiles)
 	}
 }
+
+func TestPickFileAggregator(t *testing.T) {
+	const count = 1000
+	a := Annotator[FileMetadata]{
+		Aggregator: PickFileAggregator{
+			Filter: func(f *FileMetadata) (eligible bool, cacheOK bool) {
+				return true, true
+			},
+			Compare: func(f1 *FileMetadata, f2 *FileMetadata) bool {
+				return base.DefaultComparer.Compare(f1.Smallest.UserKey, f2.Smallest.UserKey) < 0
+			},
+		},
+	}
+
+	lm, files := makeTestLevelMetadata(1)
+
+	for i := 1; i <= count; i++ {
+		lm.tree.Insert(newItem(key(i)))
+		pickedFile := a.LevelAnnotation(lm)
+		// The picked file should always be the one with the smallest key.
+		require.Same(t, files[0], pickedFile)
+	}
+}