sstable: treesteps support for sstable iterators

Author: RaduBerinde

sstable/blockiter/block_iter.go

@@ -6,6 +6,7 @@ package blockiter
 
 import (
 	"github.com/cockroachdb/pebble/internal/base"
+	"github.com/cockroachdb/pebble/internal/treesteps"
 	"github.com/cockroachdb/pebble/sstable/block"
 )
 
@@ -119,4 +120,6 @@ type Index interface {
 	// the iterator must be reset such that it could be reused after a call to
 	// Init or InitHandle.
 	Close() error
+
+	treesteps.Node
 }

sstable/colblk/data_block.go

@@ -1460,7 +1460,13 @@ func (i *DataBlockIter) String() string {
 
 // TreeStepsNode is part of the InternalIterator interface.
 func (i *DataBlockIter) TreeStepsNode() treesteps.NodeInfo {
-	return treesteps.NodeInfof(i, "%T(%p)", i, i)
+	ni := treesteps.NodeInfof(i, "colblk.DataBlockIter")
+	if i.Valid() {
+		ni.AddPropf("at", "%s", i.kv.K.String())
+	} else {
+		ni.AddPropf("not positioned", "")
+	}
+	return ni
 }
 
 // decodeRow decodes i.row into i.kv. If i.row is invalid, it returns nil.

sstable/colblk/index_block.go

@@ -13,6 +13,7 @@ import (
 	"github.com/cockroachdb/pebble/internal/binfmt"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/treeprinter"
+	"github.com/cockroachdb/pebble/internal/treesteps"
 	"github.com/cockroachdb/pebble/sstable/block"
 	"github.com/cockroachdb/pebble/sstable/blockiter"
 )
@@ -389,3 +390,13 @@ func (i *IndexIter) Close() error {
 	i.syntheticPrefixAndSuffix = blockiter.SyntheticPrefixAndSuffix{}
 	return nil
 }
+
+func (i *IndexIter) TreeStepsNode() treesteps.NodeInfo {
+	ni := treesteps.NodeInfof(i, "colblk.IndexIter")
+	if i.Valid() {
+		ni.AddPropf("at", "%s", i.Separator())
+	} else {
+		ni.AddPropf("not positioned", "")
+	}
+	return ni
+}

sstable/reader_iter_single_lvl.go

@@ -516,6 +516,9 @@ func (i *singleLevelIterator[I, PI, P, PD]) loadDataBlock(dir int8) loadBlockRes
 		return loadBlockFailed
 	}
 	i.initBounds()
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		treesteps.NodeUpdated(i, fmt.Sprintf("loadDataBlock(%d) offset=%d length=%d", dir, i.dataBH.Offset, i.dataBH.Length))
+	}
 	return loadBlockOK
 }
 
@@ -649,7 +652,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) trySeekLTUsingPrevWithinBlock(
 // caller to ensure that key is greater than or equal to the lower bound.
 func (i *singleLevelIterator[I, PI, D, PD]) SeekGE(
 	key []byte, flags base.SeekGEFlags,
-) *base.InternalKV {
+) (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "SeekGE(%q, %d)", key, flags)
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// Clear the tracking flag since this is a new absolute positioning operation
 	i.lastOpWasSeekPrefixGE.Set(false)
 	// The synthetic key is no longer relevant and must be cleared.
@@ -697,7 +706,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) SeekGE(
 // seekGEHelper contains the common functionality for SeekGE and SeekPrefixGE.
 func (i *singleLevelIterator[I, PI, D, PD]) seekGEHelper(
 	key []byte, boundsCmp int, flags base.SeekGEFlags,
-) *base.InternalKV {
+) (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "seekGEHelper(%q, %d, %d)", key, boundsCmp, flags)
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	if !i.ensureIndexLoaded() {
 		return nil
 	}
@@ -808,7 +823,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) seekGEHelper(
 // to the caller to ensure that key is greater than or equal to the lower bound.
 func (i *singleLevelIterator[I, PI, D, PD]) SeekPrefixGE(
 	prefix, key []byte, flags base.SeekGEFlags,
-) *base.InternalKV {
+) (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "SeekPrefixGE(%q, %q, %d)", prefix, key, flags)
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	if i.synthetic.atSyntheticKey {
 		// TODO(sachin) : We have to disable the optimization to avoid false data
 		// invalidation if there are back to back SeekPrefixGE calls. Currently
@@ -921,12 +942,16 @@ func (i *singleLevelIterator[I, PI, D, PD]) seekPrefixGE(
 			return nil
 		}
 		if !mayContain {
+			if treesteps.Enabled && treesteps.IsRecording(i) {
+				treesteps.UpdateLastOpf(i, "pass bloom filter did not match")
+			}
 			// In the no-error bloom filter miss case, the key is definitely not in table.
 			// We can avoid invalidating the already loaded block since the caller is
 			// not allowed to call Next when SeekPrefixGE returns nil.
 			i.lastOpWasSeekPrefixGE.Set(true)
 			return nil
 		}
+		treesteps.UpdateLastOpf(i, "bloom filter matched")
 		i.lastBloomFilterMatched = true
 	}
 	if flags.TrySeekUsingNext() {
@@ -1093,7 +1118,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) virtualLastSeekLE() *base.InternalKV
 // caller to ensure that key is less than or equal to the upper bound.
 func (i *singleLevelIterator[I, PI, D, PD]) SeekLT(
 	key []byte, flags base.SeekLTFlags,
-) *base.InternalKV {
+) (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "SeekLT(%q, %d)", key, flags)
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// Clear the tracking flag since this is a new absolute positioning operation
 	i.lastOpWasSeekPrefixGE.Set(false)
 	// The synthetic key is no longer relevant and must be cleared.
@@ -1208,7 +1239,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) SeekLT(
 // package. Note that First only checks the upper bound. It is up to the caller
 // to ensure that key is greater than or equal to the lower bound (e.g. via a
 // call to SeekGE(lower)).
-func (i *singleLevelIterator[I, PI, D, PD]) First() *base.InternalKV {
+func (i *singleLevelIterator[I, PI, D, PD]) First() (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "First()")
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// Clear the tracking flag since this is a new absolute positioning operation
 	i.lastOpWasSeekPrefixGE.Set(false)
 	// The synthetic key is no longer relevant and must be cleared.
@@ -1282,7 +1319,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) firstInternal() *base.InternalKV {
 // package. Note that Last only checks the lower bound. It is up to the caller
 // to ensure that key is less than the upper bound (e.g. via a call to
 // SeekLT(upper))
-func (i *singleLevelIterator[I, PI, D, PD]) Last() *base.InternalKV {
+func (i *singleLevelIterator[I, PI, D, PD]) Last() (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "Last()")
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// Clear the tracking flag since this is a new absolute positioning operation
 	i.lastOpWasSeekPrefixGE.Set(false)
 	// The synthetic key is no longer relevant and must be cleared.
@@ -1349,7 +1392,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) lastInternal() *base.InternalKV {
 // package.
 // Note: compactionIterator.Next mirrors the implementation of Iterator.Next
 // due to performance. Keep the two in sync.
-func (i *singleLevelIterator[I, PI, D, PD]) Next() *base.InternalKV {
+func (i *singleLevelIterator[I, PI, D, PD]) Next() (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "Next()")
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// The SeekPrefixGE might have returned a synthetic key with latest suffix
 	// contained in the sstable. If the caller is calling Next(), that means
 	// they want to move past the synthetic key and Next() is responsible for
@@ -1398,7 +1447,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) Next() *base.InternalKV {
 }
 
 // NextPrefix implements (base.InternalIterator).NextPrefix.
-func (i *singleLevelIterator[I, PI, D, PD]) NextPrefix(succKey []byte) *base.InternalKV {
+func (i *singleLevelIterator[I, PI, D, PD]) NextPrefix(succKey []byte) (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "NextPrefix(%q)", succKey)
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// Clear the tracking flag since this is a relative positioning operation
 	i.lastOpWasSeekPrefixGE.Set(false)
 	if i.exhaustedBounds == +1 {
@@ -1477,7 +1532,13 @@ func (i *singleLevelIterator[I, PI, D, PD]) NextPrefix(succKey []byte) *base.Int
 
 // Prev implements internalIterator.Prev, as documented in the pebble
 // package.
-func (i *singleLevelIterator[I, PI, D, PD]) Prev() *base.InternalKV {
+func (i *singleLevelIterator[I, PI, D, PD]) Prev() (kv *base.InternalKV) {
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		op := treesteps.StartOpf(i, "Prev()")
+		defer func() {
+			op.Finishf("= %s", kv.String())
+		}()
+	}
 	// Clear the tracking flag since this is a relative positioning operation
 	i.lastOpWasSeekPrefixGE.Set(false)
 	if i.exhaustedBounds == -1 {
@@ -1710,9 +1771,20 @@ func (i *singleLevelIterator[I, PI, D, PD]) String() string {
 
 // TreeStepsNode is part of the InternalIterator interface.
 func (i *singleLevelIterator[I, PI, D, PD]) TreeStepsNode() treesteps.NodeInfo {
-	ni := treesteps.NodeInfof(i, "%T(%p)", i, i)
-	ni.AddPropf("fileNum", "%s", i)
-	return ni
+	info := treesteps.NodeInfof(i, "sstable.singleLevelIterator")
+	if PD(&i.data).Valid() {
+		info.AddPropf("at", "%s", PD(&i.data).KV().K.String())
+	} else {
+		info.AddPropf("not positioned", "")
+	}
+	if i.indexLoaded && !PI(&i.index).IsDataInvalidated() {
+		info.AddChildren(PI(&i.index))
+	}
+	if !PD(&i.data).IsDataInvalidated() {
+		info.AddChildren(PD(&i.data))
+	}
+	info.AddPropf("file", "%s", i.String())
+	return info
 }
 
 func (i *singleLevelIterator[I, PI, D, PD]) ensureIndexLoaded() bool {
@@ -1734,5 +1806,8 @@ func (i *singleLevelIterator[I, PI, D, PD]) ensureIndexLoaded() bool {
 	}
 
 	i.indexLoaded = true
+	if treesteps.Enabled && treesteps.IsRecording(i) {
+		treesteps.NodeUpdated(i, "index block loaded")
+	}
 	return true
 }

sstable/reader_iter_treesteps_test.go

@@ -0,0 +1,80 @@
+// Copyright 2025 The LevelDB-Go and Pebble Authors. All rights reserved. Use
+// of this source code is governed by a BSD-style license that can be found in
+// the LICENSE file.
+
+package sstable
+
+import (
+	"testing"
+
+	"github.com/cockroachdb/datadriven"
+	"github.com/cockroachdb/pebble/internal/testkeys"
+	"github.com/cockroachdb/pebble/internal/treesteps"
+	"github.com/stretchr/testify/require"
+)
+
+// TestTreeSteps tests the treesteps recording for singleLevelIterator and
+// twoLevelIterator operations, generating visualization URLs showing iterator
+// behavior.
+func TestIterTreeSteps(t *testing.T) {
+	if !treesteps.Enabled {
+		t.Skip("treesteps not available in this build")
+	}
+
+	t.Run("single-level", func(t *testing.T) {
+		testIterTreeSteps(t, "testdata/treesteps_single_level_iter")
+	})
+
+	t.Run("two-level", func(t *testing.T) {
+		testIterTreeSteps(t, "testdata/treesteps_two_level_iter")
+	})
+}
+
+func testIterTreeSteps(t *testing.T, testdataPath string) {
+	var r *Reader
+	defer func() {
+		if r != nil {
+			_ = r.Close()
+			r = nil
+		}
+	}()
+	datadriven.RunTest(t, testdataPath, func(t *testing.T, td *datadriven.TestData) string {
+		switch td.Cmd {
+		case "build":
+			if r != nil {
+				_ = r.Close()
+				r = nil
+			}
+			opts := &WriterOptions{
+				Comparer:    testkeys.Comparer,
+				TableFormat: TableFormatMax,
+			}
+			var err error
+			_, r, err = runBuildCmd(td, opts, nil /* cacheHandle */)
+			if err != nil {
+				t.Fatalf("%s", err)
+			}
+			return ""
+
+		case "iter-treesteps":
+			return runIterTreeStepsCmd(t, r, td)
+
+		default:
+			return "unknown command"
+		}
+	})
+}
+
+func runIterTreeStepsCmd(t *testing.T, r *Reader, td *datadriven.TestData) string {
+	require.NotNil(t, r, "must run 'build' before 'iter-treesteps'")
+
+	iter, err := r.NewIter(NoTransforms, nil /* lower */, nil /* upper */, TableBlobContext{})
+	require.NoError(t, err)
+
+	rec := treesteps.StartRecording(iter, td.Pos)
+	runIterCmd(td, iter, false)
+
+	steps := rec.Finish()
+	url := steps.URL()
+	return url.String()
+}