From d989f68492c3310d115193fdf515425890fbffc0 Mon Sep 17 00:00:00 2001
From: Jackson Owens <jackson@cockroachlabs.com>
Date: Sun, 5 May 2024 21:34:50 -0400
Subject: [PATCH] internal/keyspan: respect bounds when switching directions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

A SeekGE(k) on an iterator with an upper bound ≤ k is permitted, and may be
issued during normal iteration due to cascading range deletion seek keys. The
symmetric dynamic exists in reverse. When an interleaving iterator's point
iterator is exhausted and switching directions, respect the bounds.
---
 internal/keyspan/interleaving_iter.go | 36 ++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/internal/keyspan/interleaving_iter.go b/internal/keyspan/interleaving_iter.go
index 418358b5003..63a556da0a7 100644
--- a/internal/keyspan/interleaving_iter.go
+++ b/internal/keyspan/interleaving_iter.go
@@ -465,12 +465,12 @@ func (i *InterleavingIter) Next() *base.InternalKV {
 			// Since we're positioned on a Span, the pointIter is positioned
 			// entirely behind the current iterator position. Reposition it
 			// ahead of the current iterator position.
-			i.savePoint(i.pointIter.Next())
+			i.switchPointIteratorIntoForward()
 		case posKeyspanEnd:
 			// Since we're positioned on a Span, the pointIter is positioned
 			// entirely behind of the current iterator position. Reposition it
 			// ahead the current iterator position.
-			i.savePoint(i.pointIter.Next())
+			i.switchPointIteratorIntoForward()
 		}
 		// Fallthrough to calling i.nextPos.
 	}
@@ -543,7 +543,7 @@ func (i *InterleavingIter) Prev() *base.InternalKV {
 			// Since we're positioned on a Span, the pointIter is positioned
 			// entirely ahead of the current iterator position. Reposition it
 			// behind the current iterator position.
-			i.savePoint(i.pointIter.Prev())
+			i.switchPointIteratorIntoReverse()
 			// Without considering truncation of spans to seek keys, the keyspan
 			// iterator is already in the right place. But consider span [a, z)
 			// and this sequence of iterator calls:
@@ -565,7 +565,7 @@ func (i *InterleavingIter) Prev() *base.InternalKV {
 			// Since we're positioned on a Span, the pointIter is positioned
 			// entirely ahead of the current iterator position. Reposition it
 			// behind the current iterator position.
-			i.savePoint(i.pointIter.Prev())
+			i.switchPointIteratorIntoReverse()
 		}
 
 		if i.spanMarkerTruncated {
@@ -637,7 +637,7 @@ func (i *InterleavingIter) nextPos() {
 
 	switch i.pos {
 	case posExhausted:
-		i.savePoint(i.pointIter.Next())
+		i.switchPointIteratorIntoForward()
 		i.saveSpanForward(i.keyspanIter.Next())
 		i.savedKeyspan()
 		i.computeSmallestPos()
@@ -709,7 +709,7 @@ func (i *InterleavingIter) prevPos() {
 
 	switch i.pos {
 	case posExhausted:
-		i.savePoint(i.pointIter.Prev())
+		i.switchPointIteratorIntoReverse()
 		i.saveSpanBackward(i.keyspanIter.Prev())
 		i.savedKeyspan()
 		i.computeLargestPos()
@@ -848,6 +848,30 @@ func (i *InterleavingIter) keyspanSeekLT(k []byte) {
 	i.savedKeyspan()
 }
 
+// switchPointIteratorIntoReverse switches the direction of the point iterator
+// into reverse, stepping to the previous point key. If the point iterator is
+// exhausted in the forward direction and there's an upper bound present, it's
+// re-seeked to ensure the iterator obeys the upper bound.
+func (i *InterleavingIter) switchPointIteratorIntoReverse() {
+	if i.pointKV == nil && i.opts.UpperBound != nil {
+		i.savePoint(i.pointIter.SeekLT(i.opts.UpperBound, base.SeekLTFlagsNone))
+		return
+	}
+	i.savePoint(i.pointIter.Prev())
+}
+
+// switchPointIteratorIntoForward switches the direction of the point iterator
+// into the forward direction, stepping to the next point key. If the point
+// iterator is exhausted in the reverse direction and there's a lower bound
+// present, it's re-seeked to ensure the iterator obeys the lower bound.
+func (i *InterleavingIter) switchPointIteratorIntoForward() {
+	if i.pointKV == nil && i.opts.LowerBound != nil {
+		i.savePoint(i.pointIter.SeekGE(i.opts.LowerBound, base.SeekGEFlagsNone))
+		return
+	}
+	i.savePoint(i.pointIter.Next())
+}
+
 func (i *InterleavingIter) saveSpanForward(span *Span, err error) {
 	i.span = span
 	i.err = firstError(i.err, err)