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 cockroachdb#3044.

Informs cockroachdb/cockroach#113973, cockroachdb/cockroach#114056.

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.

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

internal/keyspan/testdata/truncate

@@ -196,3 +196,123 @@ truncate g-g
 truncate g-h
 ----
 3:       gh
+
+# Regression test for https://github.com/cockroachdb/cockroach/issues/113973.
+
+truncate-and-save-iter a-dd
+----
+ok
+
+saved-iter
+first
+next
+next
+next
+----
+b-d:{(#1,RANGEDEL)}
+d-dd:{(#2,RANGEDEL)}
+<nil>
+<nil>
+
+saved-iter
+seek-ge e
+next
+next
+----
+<nil>
+<nil>
+<nil>
+
+saved-iter
+seek-ge e
+prev
+prev
+----
+<nil>
+d-dd:{(#2,RANGEDEL)}
+b-d:{(#1,RANGEDEL)}
+
+saved-iter
+seek-lt e
+prev
+prev
+----
+d-dd:{(#2,RANGEDEL)}
+b-d:{(#1,RANGEDEL)}
+<nil>
+
+saved-iter
+seek-lt e
+next
+next
+----
+d-dd:{(#2,RANGEDEL)}
+<nil>
+<nil>
+
+truncate-and-save-iter ee-h
+----
+ok
+
+saved-iter
+first
+next
+next
+next
+----
+ee-f:{(#2,RANGEDEL)}
+f-h:{(#3,RANGEDEL)}
+<nil>
+<nil>
+
+saved-iter
+seek-ge e
+next
+next
+----
+ee-f:{(#2,RANGEDEL)}
+f-h:{(#3,RANGEDEL)}
+<nil>
+
+saved-iter
+seek-ge e
+prev
+prev
+----
+ee-f:{(#2,RANGEDEL)}
+<nil>
+<nil>
+
+saved-iter
+seek-lt e
+prev
+prev
+----
+<nil>
+<nil>
+<nil>
+
+saved-iter
+seek-lt e
+next
+next
+----
+<nil>
+ee-f:{(#2,RANGEDEL)}
+f-h:{(#3,RANGEDEL)}
+
+
+truncate-and-save-iter a-g
+----
+ok
+
+saved-iter
+seek-ge h
+prev
+seek-lt h
+next
+----
+<nil>
+f-g:{(#3,RANGEDEL)}
+f-g:{(#3,RANGEDEL)}
+<nil>

internal/keyspan/truncate.go

@@ -69,5 +69,5 @@ func Truncate(
 		}
 
 		return !out.Empty() && cmp(out.Start, out.End) < 0
-	})
+	}, cmp)
 }

internal/keyspan/truncate_test.go

@@ -17,15 +17,16 @@ func TestTruncate(t *testing.T) {
 	cmp := base.DefaultComparer.Compare
 	fmtKey := base.DefaultComparer.FormatKey
 	var iter FragmentIterator
+	var savedIter FragmentIterator
+	defer func() {
+		if savedIter != nil {
+			savedIter.Close()
+			savedIter = nil
+		}
+	}()
 
 	datadriven.RunTest(t, "testdata/truncate", func(t *testing.T, d *datadriven.TestData) string {
-		switch d.Cmd {
-		case "build":
-			tombstones := buildSpans(t, cmp, fmtKey, d.Input, base.InternalKeyKindRangeDelete)
-			iter = NewIter(cmp, tombstones)
-			return formatAlphabeticSpans(tombstones)
-
-		case "truncate":
+		doTruncate := func() FragmentIterator {
 			if len(d.Input) > 0 {
 				t.Fatalf("unexpected input: %s", d.Input)
 			}
@@ -55,13 +56,34 @@ func TestTruncate(t *testing.T) {
 			tIter := Truncate(
 				cmp, iter, lower, upper, startKey, endKey, false,
 			)
+			return tIter
+		}
+
+		switch d.Cmd {
+		case "build":
+			tombstones := buildSpans(t, cmp, fmtKey, d.Input, base.InternalKeyKindRangeDelete)
+			iter = NewIter(cmp, tombstones)
+			return formatAlphabeticSpans(tombstones)
+
+		case "truncate":
+			tIter := doTruncate()
 			defer tIter.Close()
 			var truncated []Span
 			for s := tIter.First(); s != nil; s = tIter.Next() {
 				truncated = append(truncated, s.ShallowClone())
 			}
 			return formatAlphabeticSpans(truncated)
 
+		case "truncate-and-save-iter":
+			if savedIter != nil {
+				savedIter.Close()
+			}
+			savedIter = doTruncate()
+			return "ok"
+
+		case "saved-iter":
+			return runIterCmd(t, d, savedIter)
+
 		default:
 			return fmt.Sprintf("unknown command: %s", d.Cmd)
 		}

table_stats.go

@@ -894,7 +894,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