sstable: reimplement lazy load the index block in two level iterator

This commit reimplements lazy loading for two-level sstable iterators,
fixing critical issues from the previous attempt (fe43e7b) which was
reverted due to data corruption in stress tests.

Benchmark tests are added to validate lazy loading behavior under
various scenarios.

Implements #3248

Author: xxmplus

sstable/reader_iter_two_lvl.go

@@ -32,6 +32,10 @@ type twoLevelIterator[I any, PI indexBlockIterator[I], D any, PD dataBlockIterat
 	// false - any filtering happens at the top level.
 	useFilterBlock         bool
 	lastBloomFilterMatched bool
+
+	// topLevelIndexLoaded is set to true if the top-level index block load
+	// operation completed successfully.
+	topLevelIndexLoaded bool
 }
 
 var _ Iterator = (*twoLevelIteratorRowBlocks)(nil)
@@ -45,6 +49,12 @@ func (i *twoLevelIterator[I, PI, D, PD]) loadSecondLevelIndexBlock(dir int8) loa
 	// the index fails.
 	PD(&i.secondLevel.data).Invalidate()
 	PI(&i.secondLevel.index).Invalidate()
+
+	// Ensure top-level index is loaded before accessing it
+	if !i.ensureTopLevelIndexLoaded() {
+		return loadBlockFailed
+	}
+
 	if !PI(&i.topLevelIndex).Valid() {
 		return loadBlockFailed
 	}
@@ -87,6 +97,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) loadSecondLevelIndexBlock(dir int8) loa
 // appropriate bound, depending on the iteration direction, and returns either
 // `blockIntersects` or `blockExcluded`.
 func (i *twoLevelIterator[I, PI, D, PD]) resolveMaybeExcluded(dir int8) intersectsResult {
+	if invariants.Enabled && !i.topLevelIndexLoaded {
+		panic("pebble: resolveMaybeExcluded called without loaded top-level index")
+	}
+
 	// This iterator is configured with a bound-limited block property filter.
 	// The bpf determined this entire index block could be excluded from
 	// iteration based on the property encoded in the block handle. However, we
@@ -162,6 +176,7 @@ func newColumnBlockTwoLevelIterator(
 	}
 	i := twoLevelIterColumnBlockPool.Get().(*twoLevelIteratorColumnBlocks)
 	i.secondLevel.init(ctx, r, opts)
+	i.secondLevel.indexLoaded = true
 	// Only check the bloom filter at the top level.
 	i.useFilterBlock = i.secondLevel.useFilterBlock
 	i.secondLevel.useFilterBlock = false
@@ -181,14 +196,7 @@ func newColumnBlockTwoLevelIterator(
 			objstorage.NoReadBefore, &i.secondLevel.vbRHPrealloc)
 	}
 	i.secondLevel.data.InitOnce(r.keySchema, r.Comparer, &i.secondLevel.internalValueConstructor)
-	topLevelIndexH, err := r.readTopLevelIndexBlock(ctx, i.secondLevel.readEnv.Block, i.secondLevel.indexFilterRH)
-	if err == nil {
-		err = i.topLevelIndex.InitHandle(r.Comparer, topLevelIndexH, opts.Transforms)
-	}
-	if err != nil {
-		_ = i.Close()
-		return nil, err
-	}
+
 	return i, nil
 }
 
@@ -210,6 +218,7 @@ func newRowBlockTwoLevelIterator(
 	}
 	i := twoLevelIterRowBlockPool.Get().(*twoLevelIteratorRowBlocks)
 	i.secondLevel.init(ctx, r, opts)
+	i.secondLevel.indexLoaded = true
 	// Only check the bloom filter at the top level.
 	i.useFilterBlock = i.secondLevel.useFilterBlock
 	i.secondLevel.useFilterBlock = false
@@ -235,14 +244,6 @@ func newRowBlockTwoLevelIterator(
 		i.secondLevel.data.SetHasValuePrefix(true)
 	}
 
-	topLevelIndexH, err := r.readTopLevelIndexBlock(ctx, i.secondLevel.readEnv.Block, i.secondLevel.indexFilterRH)
-	if err == nil {
-		err = i.topLevelIndex.InitHandle(r.Comparer, topLevelIndexH, opts.Transforms)
-	}
-	if err != nil {
-		_ = i.Close()
-		return nil, err
-	}
 	return i, nil
 }
 
@@ -275,6 +276,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) SeekGE(
 	err := i.secondLevel.err
 	i.secondLevel.err = nil // clear cached iteration error
 
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	// The twoLevelIterator could be already exhausted. Utilize that when
 	// trySeekUsingNext is true. See the comment about data-exhausted, PGDE, and
 	// bounds-exhausted near the top of the file.
@@ -417,6 +422,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) SeekPrefixGE(
 	err := i.secondLevel.err
 	i.secondLevel.err = nil // clear cached iteration error
 
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	// The twoLevelIterator could be already exhausted. Utilize that when
 	// trySeekUsingNext is true. See the comment about data-exhausted, PGDE, and
 	// bounds-exhausted near the top of the file.
@@ -584,6 +593,11 @@ func (i *twoLevelIterator[I, PI, D, PD]) virtualLastSeekLE() *base.InternalKV {
 		panic("unexpected virtualLastSeekLE with exclusive upper bounds")
 	}
 	key := i.secondLevel.upper
+
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	// Need to position the topLevelIndex.
 	//
 	// The previous exhausted state of singleLevelIterator is no longer
@@ -641,6 +655,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) SeekLT(
 	// Seek optimization only applies until iterator is first positioned after SetBounds.
 	i.secondLevel.boundsCmp = 0
 
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	var result loadBlockResult
 	// NB: Unlike SeekGE, we don't have a fast-path here since we don't know
 	// whether the topLevelIndex is positioned after the position that would
@@ -714,6 +732,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) First() *base.InternalKV {
 	// Seek optimization only applies until iterator is first positioned after SetBounds.
 	i.secondLevel.boundsCmp = 0
 
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	if !PI(&i.topLevelIndex).First() {
 		return nil
 	}
@@ -763,6 +785,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) Last() *base.InternalKV {
 	// Seek optimization only applies until iterator is first positioned after SetBounds.
 	i.secondLevel.boundsCmp = 0
 
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	if !PI(&i.topLevelIndex).Last() {
 		return nil
 	}
@@ -830,6 +856,11 @@ func (i *twoLevelIterator[I, PI, D, PD]) NextPrefix(succKey []byte) *base.Intern
 
 	// Did not find prefix in the existing second-level index block. This is the
 	// slow-path where we seek the iterator.
+
+	if !i.ensureTopLevelIndexLoaded() {
+		return nil
+	}
+
 	if !PI(&i.topLevelIndex).SeekGE(succKey) {
 		PD(&i.secondLevel.data).Invalidate()
 		PI(&i.secondLevel.index).Invalidate()
@@ -877,6 +908,10 @@ func (i *twoLevelIterator[I, PI, D, PD]) skipForward() *base.InternalKV {
 			return nil
 		}
 
+		if !i.ensureTopLevelIndexLoaded() {
+			return nil
+		}
+
 		// It is possible that skipBackward went too far and the virtual table lower
 		// bound is after the first key in the block we are about to load, in which
 		// case we must use SeekGE below. The keys in the block we are about to load
@@ -954,6 +989,11 @@ func (i *twoLevelIterator[I, PI, D, PD]) skipBackward() *base.InternalKV {
 		if i.secondLevel.err != nil || i.secondLevel.exhaustedBounds < 0 {
 			return nil
 		}
+
+		if !i.ensureTopLevelIndexLoaded() {
+			return nil
+		}
+
 		i.secondLevel.exhaustedBounds = 0
 		if !PI(&i.topLevelIndex).Prev() {
 			PD(&i.secondLevel.data).Invalidate()
@@ -1007,8 +1047,39 @@ func (i *twoLevelIterator[I, PI, D, PD]) SetupForCompaction() {
 	i.secondLevel.SetupForCompaction()
 }
 
-// Close implements internalIterator.Close, as documented in the pebble
-// package.
+func (i *twoLevelIterator[I, PI, D, PD]) ensureTopLevelIndexLoaded() bool {
+	if i.topLevelIndexLoaded && i.secondLevel.err == nil {
+		return true
+	}
+
+	// Perform the deferred top-level index loading calls
+	topLevelIndexH, err := i.secondLevel.reader.readTopLevelIndexBlock(
+		i.secondLevel.ctx,
+		i.secondLevel.readEnv.Block,
+		i.secondLevel.indexFilterRH,
+	)
+	if err != nil {
+		i.secondLevel.err = err
+		i.topLevelIndexLoaded = false
+		return false
+	}
+
+	err = PI(&i.topLevelIndex).InitHandle(
+		i.secondLevel.reader.Comparer,
+		topLevelIndexH,
+		i.secondLevel.transforms,
+	)
+	if err != nil {
+		i.secondLevel.err = err
+		i.topLevelIndexLoaded = false
+		return false
+	}
+
+	i.topLevelIndexLoaded = true
+	return true
+}
+
+// Close implements internalIterator.Close, as documented in the pebble package.
 func (i *twoLevelIterator[I, PI, D, PD]) Close() error {
 	if invariants.Enabled && i.secondLevel.pool != nil {
 		panic("twoLevelIterator's singleLevelIterator has its own non-nil pool")
@@ -1019,6 +1090,7 @@ func (i *twoLevelIterator[I, PI, D, PD]) Close() error {
 	err = firstError(err, PI(&i.topLevelIndex).Close())
 	i.useFilterBlock = false
 	i.lastBloomFilterMatched = false
+	i.topLevelIndexLoaded = false
 	if pool != nil {
 		pool.Put(i)
 	}