From 1c6b496298b534970dd40533d9dc5f457bbb547e Mon Sep 17 00:00:00 2001
From: Bilal Akhtar <bilal@cockroachlabs.com>
Date: Wed, 8 Nov 2023 17:10:20 -0500
Subject: [PATCH] internal/keyspan: obey Seek invariants in filteringIter

Previously if filteringIter's FilterFunc mutated the passed-in span
to no longer be a valid return value for a SeekLT or SeekGE call,
we would still return that span even though it could be >=  the seek
key (for SeekLT), or less than it (for SeekGE).

This change updates filteringIter to guard for this case before
returning from a seek call.

Found with #3044.

Informs cockroachdb/cockroach#113973, cockroachdb/cockroach#114056.
---
 internal/keyspan/filter.go      | 25 +++++++++++++++++++++----
 internal/keyspan/filter_test.go |  2 +-
 internal/keyspan/truncate.go    |  6 ++++--
 table_stats.go                  |  2 +-
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/internal/keyspan/filter.go b/internal/keyspan/filter.go
index 5332d90e235..a63a43ca8fe 100644
--- a/internal/keyspan/filter.go
+++ b/internal/keyspan/filter.go
@@ -4,6 +4,8 @@
 
 package keyspan
 
+import "github.com/cockroachdb/pebble/internal/base"
+
 // FilterFunc defines a transform from the input Span into the output Span. The
 // function returns true if the Span should be returned by the iterator, and
 // false if the Span should be skipped. The FilterFunc is permitted to mutate
@@ -21,6 +23,7 @@ type FilterFunc func(in *Span, out *Span) (keep bool)
 type filteringIter struct {
 	iter     FragmentIterator
 	filterFn FilterFunc
+	cmp      base.Compare
 
 	// span is a mutable Span passed to the filterFn. The filterFn is free to
 	// mutate this Span. The slice of Keys in the Span is reused with every call
@@ -32,18 +35,32 @@ var _ FragmentIterator = (*filteringIter)(nil)
 
 // Filter returns a new filteringIter that will filter the Spans from the
 // provided child iterator using the provided FilterFunc.
-func Filter(iter FragmentIterator, filter FilterFunc) FragmentIterator {
-	return &filteringIter{iter: iter, filterFn: filter}
+func Filter(iter FragmentIterator, filter FilterFunc, cmp base.Compare) FragmentIterator {
+	return &filteringIter{iter: iter, filterFn: filter, cmp: cmp}
 }
 
 // SeekGE implements FragmentIterator.
 func (i *filteringIter) SeekGE(key []byte) *Span {
-	return i.filter(i.iter.SeekGE(key), +1)
+	span := i.filter(i.iter.SeekGE(key), +1)
+	// i.filter could return a span that's less than key, _if_ the filterFunc
+	// (which has no knowledge of the seek key) mutated the span to end at a key
+	// less than or equal to `key`. Detect this case and next/invalidate the iter.
+	if span != nil && i.cmp(span.End, key) <= 0 {
+		return i.Next()
+	}
+	return span
 }
 
 // SeekLT implements FragmentIterator.
 func (i *filteringIter) SeekLT(key []byte) *Span {
-	return i.filter(i.iter.SeekLT(key), -1)
+	span := i.filter(i.iter.SeekLT(key), -1)
+	// i.filter could return a span that's >= key, _if_ the filterFunc (which has
+	// no knowledge of the seek key) mutated the span to start at a key greater
+	// than or equal to `key`. Detect this case and prev/invalidate the iter.
+	if span != nil && i.cmp(span.Start, key) >= 0 {
+		return i.Prev()
+	}
+	return span
 }
 
 // First implements FragmentIterator.
diff --git a/internal/keyspan/filter_test.go b/internal/keyspan/filter_test.go
index 052d3ffcf2d..beb4de8cd74 100644
--- a/internal/keyspan/filter_test.go
+++ b/internal/keyspan/filter_test.go
@@ -67,7 +67,7 @@ func TestFilteringIter(t *testing.T) {
 				}
 			}
 			innerIter := NewIter(cmp, spans)
-			iter := Filter(innerIter, filter)
+			iter := Filter(innerIter, filter, cmp)
 			defer iter.Close()
 			s := runFragmentIteratorCmd(iter, td.Input, nil)
 			return s
diff --git a/internal/keyspan/truncate.go b/internal/keyspan/truncate.go
index 97b44809954..c3fd92ae81f 100644
--- a/internal/keyspan/truncate.go
+++ b/internal/keyspan/truncate.go
@@ -4,7 +4,9 @@
 
 package keyspan
 
-import "github.com/cockroachdb/pebble/internal/base"
+import (
+	"github.com/cockroachdb/pebble/internal/base"
+)
 
 // Truncate creates a new iterator where every span in the supplied iterator is
 // truncated to be contained within the range [lower, upper). If start and end
@@ -69,5 +71,5 @@ func Truncate(
 		}
 
 		return !out.Empty() && cmp(out.Start, out.End) < 0
-	})
+	}, cmp)
 }
diff --git a/table_stats.go b/table_stats.go
index 64f87c76cdf..8b3e02bc934 100644
--- a/table_stats.go
+++ b/table_stats.go
@@ -931,7 +931,7 @@ func newCombinedDeletionKeyspanIter(
 				out.Keys = append(out.Keys, k)
 			}
 			return len(out.Keys) > 0
-		})
+		}, comparer.Compare)
 		dIter := &keyspan.DefragmentingIter{}
 		dIter.Init(comparer, iter, equal, reducer, new(keyspan.DefragmentingBuffers))
 		iter = dIter