iterv2: fix invalid TrySeekUsingNext in mergingIterV2 batch-refresh path

`mergingIterV2.seekGEAfterBatchRefresh` handles the
`SeekGE(TrySeekUsingNext|BatchJustRefreshed)` case by pre-seeking the
batch level to decide whether the merging iterator must move back.
Previously, when the pre-seek showed that the iterator did not need to
move back, the code fell through to a full slab rebuild
(`m.seekGE(key, flags)`) with `TrySeekUsingNext` still enabled.

That is unsound when `seekKey < prevTopKey`: lower levels may have been
advanced past `seekKey` across one or more slab boundaries (see the
analogous case in `SeekGE` at the non-`BatchJustRefreshed` fast path),
so re-seeking them at `seekKey` with TSUN is illegal and can return
incorrect results.

This commit handles the `seekKey < prevTopKey` case directly: if the
batch level's span keys are unchanged (so the existing slab state —
per-level `minSeqNum`/`maxSeqNum`/`parked` — remains valid) the heap
is simply rebuilt in place and the iterator returns its current top.
Otherwise we fall through with TSUN disabled. The early exit at the
`SeekGE` call site is routed through `findNextEntry` so any
`SpanBoundary` or visibility-filtered KV at the new heap top is
handled correctly. A small additional fast path skips the pre-seek
entirely when the batch level is already at or past `seekKey`.

`TestSetOptionsBatchRefreshRand` is extended to exercise the affected
paths: RANGEDELs are added to the flushed sstables and to the batch
(the batch gets an inert RANGEDEL up front so `SetOptions` doesn't
tear down `pointIter`), RANGEDEL bounds may be either bare prefixes or
suffixed keys so boundary user keys can collide with point keys, and
the iteration count is raised.

Author: RaduBerinde

iterator_test.go

@@ -1302,32 +1302,79 @@ func TestSetOptionsBatchRefreshRand(t *testing.T) {
 	rng := rand.New(rand.NewPCG(seed, seed))
 
 	mem := vfs.NewMem()
-	d, err := Open("", &Options{FS: mem, Comparer: testkeys.Comparer})
+	d, err := Open("", &Options{
+		FS:                          mem,
+		Comparer:                    testkeys.Comparer,
+		DisableAutomaticCompactions: true,
+	})
 	require.NoError(t, err)
 	defer func() { require.NoError(t, d.Close()) }()
 
-	// Generate keys via iterv2.KeyGenConfig. A small prefix length and suffix
-	// range encourage collisions between batch and LSM keys.
-	keyCfg := iterv2.KeyGenConfig{MaxPrefixLen: 2, MaxSuffix: 3, MinSeqNum: 1, MaxSeqNum: 1}
-	prefixCfg := iterv2.KeyGenConfig{MaxPrefixLen: 2, MinSeqNum: 1, MaxSeqNum: 1}
+	// Generate keys via iterv2.KeyGenConfig. A moderate prefix length keeps the
+	// key space sparse enough that a seek to one key often lands on a
+	// significantly different key (so prevTopKey > newSeekKey can hold), while
+	// still producing collisions between batch and LSM keys.
+	keyCfg := iterv2.KeyGenConfig{MaxPrefixLen: 4, MaxSuffix: 3, MinSeqNum: 1, MaxSeqNum: 1}
+	prefixCfg := iterv2.KeyGenConfig{MaxPrefixLen: 4, MinSeqNum: 1, MaxSeqNum: 1}
 	randKey := func() []byte { return keyCfg.RandKey(rng) }
 	randPrefix := func() []byte { return prefixCfg.RandKey(rng) }
 
+	// addRangeDel writes a random RANGEDEL to the given writer. Each bound is
+	// independently either a bare prefix or a full suffixed key, so RANGEDEL
+	// boundary user keys can collide with point-key user keys generated by
+	// randKey() — this is required to exercise the merging iterator's heap
+	// tiebreak path when the batch level lands at the same user key as the
+	// previous top.
+	randBound := func() []byte {
+		if rng.IntN(2) == 0 {
+			return slices.Clone(randPrefix())
+		}
+		return slices.Clone(randKey())
+	}
+	addRangeDel := func(w interface {
+		DeleteRange([]byte, []byte, *WriteOptions) error
+	}) {
+		a := randBound()
+		b := randBound()
+		if bytes.Compare(a, b) > 0 {
+			a, b = b, a
+		}
+		if bytes.Equal(a, b) {
+			// Extend the prefix portion of b so it remains a valid testkeys
+			// key (possibly bare-prefix) that is strictly greater than a.
+			// Appending to a suffixed key directly would yield an invalid
+			// testkeys encoding.
+			b = append(slices.Clip(testkeys.Comparer.Prefix(b)), byte('a')+byte(rng.IntN(8)))
+		}
+		require.NoError(t, w.DeleteRange(a, b, nil))
+	}
+
 	// Seed the LSM with multiple flushed sstables so the merging iterator has
 	// several non-batch levels beneath the batch. Each flush writes a random
-	// scattering of keys; the final memtable layer adds another level.
+	// scattering of keys plus a couple of RANGEDELs so the lower levels have
+	// non-trivial slab boundaries; the final memtable layer adds another level.
 	for f := 0; f < 3; f++ {
 		for k := 0; k < 8; k++ {
 			require.NoError(t, d.Set(randKey(), randValue(4, rng), nil))
 		}
+		for r := 0; r < 1+rng.IntN(2); r++ {
+			addRangeDel(d)
+		}
 		require.NoError(t, d.Flush())
 	}
 	for k := 0; k < 8; k++ {
 		require.NoError(t, d.Set(randKey(), randValue(4, rng), nil))
 	}
+	addRangeDel(d)
 
 	batch := d.NewIndexedBatch()
 	defer batch.Close()
+	// Add an inert RANGEDEL so the batch's range-del iterator is wired into the
+	// InterleavingIter at iterator construction time. Without this, the first
+	// batch RANGEDEL added during a refresh would cause SetOptions to close and
+	// rebuild pointIter (see iterator.go SetOptions), bypassing the
+	// BatchJustRefreshed+TSUN path under test.
+	require.NoError(t, batch.DeleteRange([]byte("\xff\xfe"), []byte("\xff\xff"), nil))
 
 	iter, _ := batch.NewIter(nil)
 	defer iter.Close()
@@ -1376,17 +1423,7 @@ func TestSetOptionsBatchRefreshRand(t *testing.T) {
 			case 0:
 				require.NoError(t, batch.Set(randKey(), randValue(4, rng), nil))
 			case 1:
-				// Range del bounds use bare prefixes so we don't have to deal
-				// with testkeys suffix encoding when extending the upper bound.
-				a := slices.Clone(randPrefix())
-				b := slices.Clone(randPrefix())
-				if bytes.Compare(a, b) > 0 {
-					a, b = b, a
-				}
-				if bytes.Equal(a, b) {
-					b = append(b, byte('a')+byte(rng.IntN(8)))
-				}
-				require.NoError(t, batch.DeleteRange(a, b, nil))
+				addRangeDel(batch)
 			}
 		}
 

merging_iter_v2.go

@@ -437,7 +437,10 @@ func (m *mergingIterV2) SeekGE(key []byte, flags base.SeekGEFlags) (kv *base.Int
 		if flags.BatchJustRefreshed() {
 			var earlyExit bool
 			if flags, earlyExit = m.seekGEAfterBatchRefresh(key, flags); earlyExit {
-				return nil
+				if m.heap.Len() == 0 {
+					return nil
+				}
+				return m.findNextEntry()
 			}
 			// We fall back to rebuilding the entire slab, with the updated flags.
 			// TODO(radu): the batch level is already in the right place; we could
@@ -509,13 +512,19 @@ func (m *mergingIterV2) SeekGE(key []byte, flags base.SeekGEFlags) (kv *base.Int
 // a backward move is required, and BatchJustRefreshed is always cleared
 // (the batch level was just seeked).
 func (m *mergingIterV2) seekGEAfterBatchRefresh(
-	key []byte, flags base.SeekGEFlags,
+	seekKey []byte, flags base.SeekGEFlags,
 ) (newFlags base.SeekGEFlags, earlyExit bool) {
-	newFlags = flags.DisableBatchJustRefreshed()
 	if invariants.Enabled && m.slab.batchSnapshot == 0 {
 		panic(errors.AssertionFailedf("BatchJustRefreshed with no batch iterator"))
 	}
 	level := &m.levels[m.slab.batchLevelIdx]
+
+	if level.iterKV != nil && m.heap.cmp(level.iterKV.K.UserKey, seekKey) <= 0 {
+		// Fast path: we are seeking past the current batch iterator key anyway, so
+		// no special handling needed. We must leave BatchJustRefreshed set.
+		return flags, false
+	}
+
 	var prevTopKey []byte
 	// Whether the heap is empty or non-empty, we need to seek the batch level to
 	// see if it causes the mergingIter to move back. If the heap is non-empty, we
@@ -530,24 +539,44 @@ func (m *mergingIterV2) seekGEAfterBatchRefresh(
 			prevTopKey = m.keyBuf
 		}
 	}
+	spanKeysDetector, _ := makeSpanKeysChangeDetector(level.span.Keys)
 	m.prepareForLevelOp(level)
-	level.iterKV = level.iter.SeekGE(key, flags)
+	level.iterKV = level.iter.SeekGE(seekKey, flags)
+	newFlags = flags.DisableBatchJustRefreshed()
 	switch {
 	case level.iterKV == nil:
 		// The entire operation is done if there was an error or the merging
 		// iterator is still exhausted.
 		if m.levelHasError(level) || prevTopKey == nil {
-			return newFlags, true
+			return 0, true
 		}
-		// The heap is not moving back, we can use TrySeekUsingNext on other levels.
 	case prevTopKey == nil || m.heap.cmp(level.iterKV.K.UserKey, prevTopKey) < 0:
 		// Either the heap was empty (so other levels are exhausted and need
 		// absolute re-seeks) or the slab moves back. Fall back to the general
 		// seek path.
-		newFlags = newFlags.DisableTrySeekUsingNext()
-	default:
-		// The heap is not moving back, we can use TrySeekUsingNext on other levels.
-	}
+		return newFlags.DisableTrySeekUsingNext(), false
+	}
+	// The heap is not moving back. Check if the iterator is already at the right
+	// position. We have to handle this case because seeking all levels at seekKey
+	// with TrySeekUsingNext might not be legal. See the corresponding case in
+	// SeekGE for an example.
+	if m.heap.cmp(seekKey, prevTopKey) < 0 {
+		if spanKeysDetector.MayHaveChanged(level.span.Keys) {
+			// The batch level's span may have changed. This could affect the
+			// per-level visibility (minSeqNum) and parked status of lower levels,
+			// requiring a proper slab rebuild. Fall back to the general seek path
+			// with TSUN disabled.
+			return newFlags.DisableTrySeekUsingNext(), false
+		}
+		// Rebuild the heap, since level.iterKV may have changed.
+		level.maxSeqNum = m.slab.batchSnapshot
+		m.initHeap(+1)
+		// The batch level's span boundary may have changed, so the slab's
+		// nextBoundary is potentially stale. Recompute it.
+		m.slab.calcNextBoundary(+1)
+		return 0, true
+	}
+	// We can use TrySeekUsingNext on other levels.
 	return newFlags, false
 }