Repository: cockroachdb/pebble
Pull request #3783 (2024-07-20 23:55:06)

keyspan: add contexts to keyspan iterators
Informs #3728

#### db: convert TestTracing to datadriven

Rewrite this test to use `datadriven`. This will make it much easier
to update and extend.

#### keyspan: add contexts to keyspan iterators

This change adds a `SetContext()` method to `FragmentIterator` which
updates any context stored in an iterator and its children. The
iterator context is now plumbed correctly through to the table range
key iterators.

cockroach-teamcity: 

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/cockroachdb/pebble/3783)
<!-- Reviewable:end -->


Patch:
diff --git a/compaction.go b/compaction.go
index 71186e86..44cd103d 100644
--- a/compaction.go
+++ b/compaction.go
@@ -807,9 +807,8 @@ func (c *compaction) newInputIters(
 				}
 			}
 			if hasRangeKeys {
-				li := &keyspanimpl.LevelIter{}
-				newRangeKeyIterWrapper := func(file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
-					rangeKeyIter, err := newRangeKeyIter(file, iterOptions)
+				newRangeKeyIterWrapper := func(ctx context.Context, file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
+					rangeKeyIter, err := newRangeKeyIter(ctx, file, iterOptions)
 					if err != nil {
 						return nil, err
 					} else if rangeKeyIter == nil {
@@ -829,7 +828,10 @@ func (c *compaction) newInputIters(
 					// in this sstable must wholly lie within the file's bounds.
 					return noCloseIter, err
 				}
-				li.Init(keyspan.SpanIterOptions{}, c.cmp, newRangeKeyIterWrapper, level.files.Iter(), l, manifest.KeyTypeRange)
+				li := keyspanimpl.NewLevelIter(
+					context.Background(), keyspan.SpanIterOptions{}, c.cmp,
+					newRangeKeyIterWrapper, level.files.Iter(), l, manifest.KeyTypeRange,
+				)
 				rangeKeyIters = append(rangeKeyIters, li)
 			}
 			return nil
diff --git a/error_iter.go b/error_iter.go
index 61911c9f..3c741f86 100644
--- a/error_iter.go
+++ b/error_iter.go
@@ -90,6 +90,7 @@ func (i *errorKeyspanIter) First() (*keyspan.Span, error)            { return ni
 func (i *errorKeyspanIter) Last() (*keyspan.Span, error)             { return nil, i.err }
 func (i *errorKeyspanIter) Next() (*keyspan.Span, error)             { return nil, i.err }
 func (i *errorKeyspanIter) Prev() (*keyspan.Span, error)             { return nil, i.err }
+func (i *errorKeyspanIter) SetContext(ctx context.Context)           {}
 func (i *errorKeyspanIter) Close()                                   {}
 func (*errorKeyspanIter) String() string                             { return "error" }
 func (*errorKeyspanIter) WrapChildren(wrap keyspan.WrapFn)           {}
diff --git a/flushable.go b/flushable.go
index bc2915c2..b06f7012 100644
--- a/flushable.go
+++ b/flushable.go
@@ -233,9 +233,9 @@ func (s *ingestedFlushable) newFlushIter(*IterOptions) internalIterator {
 }
 
 func (s *ingestedFlushable) constructRangeDelIter(
-	file *manifest.FileMetadata, _ keyspan.SpanIterOptions,
+	ctx context.Context, file *manifest.FileMetadata, _ keyspan.SpanIterOptions,
 ) (keyspan.FragmentIterator, error) {
-	iters, err := s.newIters(context.Background(), file, nil, internalIterOpts{}, iterRangeDeletions)
+	iters, err := s.newIters(ctx, file, nil, internalIterOpts{}, iterRangeDeletions)
 	if err != nil {
 		return nil, err
 	}
@@ -250,6 +250,7 @@ func (s *ingestedFlushable) constructRangeDelIter(
 // the point iterator in constructRangeDeIter is not tracked.
 func (s *ingestedFlushable) newRangeDelIter(_ *IterOptions) keyspan.FragmentIterator {
 	return keyspanimpl.NewLevelIter(
+		context.TODO(),
 		keyspan.SpanIterOptions{}, s.comparer.Compare,
 		s.constructRangeDelIter, s.slice.Iter(), manifest.Level(0),
 		manifest.KeyTypePoint,
@@ -263,6 +264,7 @@ func (s *ingestedFlushable) newRangeKeyIter(o *IterOptions) keyspan.FragmentIter
 	}
 
 	return keyspanimpl.NewLevelIter(
+		context.TODO(),
 		keyspan.SpanIterOptions{}, s.comparer.Compare, s.newRangeKeyIters,
 		s.slice.Iter(), manifest.Level(0), manifest.KeyTypeRange,
 	)
diff --git a/internal/keyspan/assert_iter.go b/internal/keyspan/assert_iter.go
index 42f76bb9..ded00db9 100644
--- a/internal/keyspan/assert_iter.go
+++ b/internal/keyspan/assert_iter.go
@@ -5,6 +5,7 @@
 package keyspan
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/cockroachdb/errors"
@@ -165,6 +166,11 @@ func (i *assertIter) Prev() (*Span, error) {
 	return span, err
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *assertIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 // Close implements FragmentIterator.
 func (i *assertIter) Close() {
 	i.iter.Close()
diff --git a/internal/keyspan/bounded.go b/internal/keyspan/bounded.go
index 580cb74e..379676b7 100644
--- a/internal/keyspan/bounded.go
+++ b/internal/keyspan/bounded.go
@@ -5,6 +5,8 @@
 package keyspan
 
 import (
+	"context"
+
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/treeprinter"
 )
@@ -209,6 +211,11 @@ func (i *BoundedIter) Close() {
 	i.iter.Close()
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *BoundedIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 // SetBounds modifies the FragmentIterator's bounds.
 func (i *BoundedIter) SetBounds(lower, upper []byte) {
 	i.lower, i.upper = lower, upper
diff --git a/internal/keyspan/defragment.go b/internal/keyspan/defragment.go
index eb9917d2..35dda16d 100644
--- a/internal/keyspan/defragment.go
+++ b/internal/keyspan/defragment.go
@@ -6,6 +6,7 @@ package keyspan
 
 import (
 	"bytes"
+	"context"
 
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/bytealloc"
@@ -196,6 +197,11 @@ func (i *DefragmentingIter) Init(
 	}
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *DefragmentingIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 // Close closes the underlying iterators.
 func (i *DefragmentingIter) Close() {
 	i.iter.Close()
diff --git a/internal/keyspan/filter.go b/internal/keyspan/filter.go
index efba323d..8b09b2fd 100644
--- a/internal/keyspan/filter.go
+++ b/internal/keyspan/filter.go
@@ -5,6 +5,8 @@
 package keyspan
 
 import (
+	"context"
+
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/treeprinter"
 )
@@ -93,6 +95,11 @@ func (i *filteringIter) Prev() (*Span, error) {
 	return i.filter(s, -1)
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *filteringIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 // Close implements FragmentIterator.
 func (i *filteringIter) Close() {
 	i.iter.Close()
diff --git a/internal/keyspan/interleaving_iter.go b/internal/keyspan/interleaving_iter.go
index 5c02ed49..39a3759f 100644
--- a/internal/keyspan/interleaving_iter.go
+++ b/internal/keyspan/interleaving_iter.go
@@ -1140,6 +1140,7 @@ func (i *InterleavingIter) SetBounds(lower, upper []byte) {
 // SetContext implements (base.InternalIterator).SetContext.
 func (i *InterleavingIter) SetContext(ctx context.Context) {
 	i.pointIter.SetContext(ctx)
+	i.keyspanIter.SetContext(ctx)
 }
 
 // DebugTree is part of the InternalIterator interface.
diff --git a/internal/keyspan/iter.go b/internal/keyspan/iter.go
index 06ff261a..5d13e065 100644
--- a/internal/keyspan/iter.go
+++ b/internal/keyspan/iter.go
@@ -5,6 +5,8 @@
 package keyspan
 
 import (
+	"context"
+
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/treeprinter"
 )
@@ -69,6 +71,10 @@ type FragmentIterator interface {
 	// stack. Used only for debug logging.
 	WrapChildren(wrap WrapFn)
 
+	// SetContext replaces the context provided at iterator creation, or the last
+	// one provided by SetContext.
+	SetContext(ctx context.Context)
+
 	base.IteratorDebug
 }
 
@@ -211,6 +217,9 @@ func (i *Iter) Prev() (*Span, error) {
 	return &i.spans[i.index], nil
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *Iter) SetContext(ctx context.Context) {}
+
 // Close implements FragmentIterator.Close.
 func (i *Iter) Close() {}
 
diff --git a/internal/keyspan/keyspanimpl/level_iter.go b/internal/keyspan/keyspanimpl/level_iter.go
index e4b266e2..2fc5b836 100644
--- a/internal/keyspan/keyspanimpl/level_iter.go
+++ b/internal/keyspan/keyspanimpl/level_iter.go
@@ -5,6 +5,7 @@
 package keyspanimpl
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/cockroachdb/pebble/internal/base"
@@ -17,7 +18,7 @@ import (
 // TableNewSpanIter creates a new iterator for range key spans for the given
 // file.
 type TableNewSpanIter func(
-	file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions,
+	ctx context.Context, file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions,
 ) (keyspan.FragmentIterator, error)
 
 // LevelIter provides a merged view of spans from sstables in an L1+ level or an
@@ -46,6 +47,8 @@ type LevelIter struct {
 	// newIter creates a range del iterator if keyType is KeyTypePoint or a range
 	// key iterator if keyType is KeyTypeRange.
 	newIter TableNewSpanIter
+	// ctx is passed to TableNewSpanIter.
+	ctx context.Context
 
 	// The options that were passed in.
 	tableOpts keyspan.SpanIterOptions
@@ -79,6 +82,7 @@ var _ keyspan.FragmentIterator = (*LevelIter)(nil)
 // newIter must create a range del iterator for the given file if keyType is
 // KeyTypePoint or a range key iterator if keyType is KeyTypeRange.
 func NewLevelIter(
+	ctx context.Context,
 	opts keyspan.SpanIterOptions,
 	cmp base.Compare,
 	newIter TableNewSpanIter,
@@ -87,7 +91,7 @@ func NewLevelIter(
 	keyType manifest.KeyType,
 ) *LevelIter {
 	l := &LevelIter{}
-	l.Init(opts, cmp, newIter, files, level, keyType)
+	l.Init(ctx, opts, cmp, newIter, files, level, keyType)
 	return l
 }
 
@@ -96,6 +100,7 @@ func NewLevelIter(
 // newIter must create a range del iterator for the given file if keyType is
 // KeyTypePoint or a range key iterator if keyType is KeyTypeRange.
 func (l *LevelIter) Init(
+	ctx context.Context,
 	opts keyspan.SpanIterOptions,
 	cmp base.Compare,
 	newIter TableNewSpanIter,
@@ -111,6 +116,7 @@ func (l *LevelIter) Init(
 		keyType:   keyType,
 		level:     level,
 		newIter:   newIter,
+		ctx:       ctx,
 		tableOpts: opts,
 		files:     files.Filter(keyType),
 	}
@@ -387,6 +393,14 @@ func (l *LevelIter) moveToPrevFile() (*keyspan.Span, error) {
 	}
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (l *LevelIter) SetContext(ctx context.Context) {
+	l.ctx = ctx
+	if l.lastIter != nil {
+		l.lastIter.SetContext(ctx)
+	}
+}
+
 // Close implements keyspan.FragmentIterator.
 func (l *LevelIter) Close() {
 	l.file = nil
@@ -442,7 +456,7 @@ func (l *LevelIter) setPosAtFile(f *manifest.FileMetadata) error {
 			l.lastIter = nil
 			l.lastIterFile = nil
 		}
-		iter, err := l.newIter(l.file, l.tableOpts)
+		iter, err := l.newIter(l.ctx, l.file, l.tableOpts)
 		if err != nil {
 			return err
 		}
diff --git a/internal/keyspan/keyspanimpl/merging_iter.go b/internal/keyspan/keyspanimpl/merging_iter.go
index 5cb07254..9fd87960 100644
--- a/internal/keyspan/keyspanimpl/merging_iter.go
+++ b/internal/keyspan/keyspanimpl/merging_iter.go
@@ -7,6 +7,7 @@ package keyspanimpl
 import (
 	"bytes"
 	"cmp"
+	"context"
 	"fmt"
 	"slices"
 
@@ -676,6 +677,13 @@ func (m *MergingIter) Prev() (*keyspan.Span, error) {
 	return m.findPrevFragmentSet()
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (m *MergingIter) SetContext(ctx context.Context) {
+	for i := range m.levels {
+		m.levels[i].iter.SetContext(ctx)
+	}
+}
+
 // Close closes the iterator, releasing all acquired resources.
 func (m *MergingIter) Close() {
 	for i := range m.levels {
diff --git a/internal/keyspan/logging_iter.go b/internal/keyspan/logging_iter.go
index a0a4b1ab..250d6691 100644
--- a/internal/keyspan/logging_iter.go
+++ b/internal/keyspan/logging_iter.go
@@ -5,6 +5,7 @@
 package keyspan
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/cockroachdb/pebble/internal/base"
@@ -132,6 +133,11 @@ func (i *loggingIter) Prev() (*Span, error) {
 	return span, err
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *loggingIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 // Close implements FragmentIterator.
 func (i *loggingIter) Close() {
 	opEnd := i.opStartf("Close()")
diff --git a/internal/keyspan/truncate.go b/internal/keyspan/truncate.go
index 20ea12da..721d514f 100644
--- a/internal/keyspan/truncate.go
+++ b/internal/keyspan/truncate.go
@@ -5,6 +5,8 @@
 package keyspan
 
 import (
+	"context"
+
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/invariants"
 	"github.com/cockroachdb/pebble/internal/treeprinter"
@@ -112,6 +114,11 @@ func (i *truncatingIter) Prev() (*Span, error) {
 	return span, err
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *truncatingIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 // Close implements FragmentIterator.
 func (i *truncatingIter) Close() {
 	i.iter.Close()
diff --git a/open.go b/open.go
index c52e80d6..cb77a52b 100644
--- a/open.go
+++ b/open.go
@@ -406,7 +406,7 @@ func Open(dirname string, opts *Options) (db *DB, err error) {
 		opts.TableCache, d.cacheID, d.objProvider, d.opts, tableCacheSize,
 		&sstable.CategoryStatsCollector{})
 	d.newIters = d.tableCache.newIters
-	d.tableNewRangeKeyIter = tableNewRangeKeyIter(context.TODO(), d.newIters)
+	d.tableNewRangeKeyIter = tableNewRangeKeyIter(d.newIters)
 
 	var previousOptionsFileNum base.DiskFileNum
 	var previousOptionsFilename string
diff --git a/range_keys.go b/range_keys.go
index 8a91dbbb..52640a41 100644
--- a/range_keys.go
+++ b/range_keys.go
@@ -92,6 +92,7 @@ func (i *Iterator) constructRangeKeyIter() {
 
 				li := i.rangeKey.iterConfig.NewLevelIter()
 				li.Init(
+					i.ctx,
 					i.opts.SpanIterOptions(),
 					i.cmp,
 					i.newIterRangeKey,
@@ -110,7 +111,7 @@ func (i *Iterator) constructRangeKeyIter() {
 			}
 			li := i.rangeKey.iterConfig.NewLevelIter()
 			spanIterOpts := i.opts.SpanIterOptions()
-			li.Init(spanIterOpts, i.cmp, i.newIterRangeKey, current.RangeKeyLevels[level].Iter(),
+			li.Init(i.ctx, spanIterOpts, i.cmp, i.newIterRangeKey, current.RangeKeyLevels[level].Iter(),
 				manifest.Level(level), manifest.KeyTypeRange)
 			i.rangeKey.iterConfig.AddLevel(li)
 		}
diff --git a/scan_internal.go b/scan_internal.go
index 7f7ca327..6081c172 100644
--- a/scan_internal.go
+++ b/scan_internal.go
@@ -908,8 +908,8 @@ func (i *scanInternalIterator) constructPointIter(
 			i.ctx, i.opts.IterOptions, i.comparer, i.newIters, files, level,
 			internalIterOpts{})
 		mlevels[mlevelsIndex].iter = li
-		rli.Init(keyspan.SpanIterOptions{RangeKeyFilters: i.opts.RangeKeyFilters},
-			i.comparer.Compare, tableNewRangeDelIter(i.ctx, i.newIters), files, level,
+		rli.Init(i.ctx, keyspan.SpanIterOptions{RangeKeyFilters: i.opts.RangeKeyFilters},
+			i.comparer.Compare, tableNewRangeDelIter(i.newIters), files, level,
 			manifest.KeyTypePoint)
 		rangeDelIters = append(rangeDelIters, rli)
 
@@ -1030,7 +1030,7 @@ func (i *scanInternalIterator) constructRangeKeyIter() error {
 	// around Key InternalKeyTrailer order.
 	iter := current.RangeKeyLevels[0].Iter()
 	for f := iter.Last(); f != nil; f = iter.Prev() {
-		spanIter, err := i.newIterRangeKey(f, i.opts.SpanIterOptions())
+		spanIter, err := i.newIterRangeKey(i.ctx, f, i.opts.SpanIterOptions())
 		if err != nil {
 			return err
 		}
@@ -1065,7 +1065,7 @@ func (i *scanInternalIterator) constructRangeKeyIter() error {
 			levSlice := manifest.NewLevelSliceKeySorted(i.db.cmp, nonRemoteFiles)
 			levIter = levSlice.Iter()
 		}
-		li.Init(spanIterOpts, i.comparer.Compare, i.newIterRangeKey, levIter,
+		li.Init(i.ctx, spanIterOpts, i.comparer.Compare, i.newIterRangeKey, levIter,
 			manifest.Level(level), manifest.KeyTypeRange)
 		i.rangeKey.iterConfig.AddLevel(li)
 	}
diff --git a/sstable/rowblk/rowblk_fragment_iter.go b/sstable/rowblk/rowblk_fragment_iter.go
index b4648b5e..2504e74e 100644
--- a/sstable/rowblk/rowblk_fragment_iter.go
+++ b/sstable/rowblk/rowblk_fragment_iter.go
@@ -6,6 +6,7 @@ package rowblk
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"os"
 	"sync"
@@ -258,6 +259,9 @@ func (i *fragmentIter) gatherBackward(kv *base.InternalKV) (*keyspan.Span, error
 	return &i.span, nil
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *fragmentIter) SetContext(ctx context.Context) {}
+
 // Close implements (keyspan.FragmentIterator).Close.
 func (i *fragmentIter) Close() {
 	i.blockIter.Close()
diff --git a/table_cache.go b/table_cache.go
index aee22ac3..dba3c4b5 100644
--- a/table_cache.go
+++ b/table_cache.go
@@ -62,10 +62,8 @@ type tableNewIters func(
 
 // tableNewRangeDelIter takes a tableNewIters and returns a TableNewSpanIter
 // for the rangedel iterator returned by tableNewIters.
-func tableNewRangeDelIter(
-	ctx context.Context, newIters tableNewIters,
-) keyspanimpl.TableNewSpanIter {
-	return func(file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
+func tableNewRangeDelIter(newIters tableNewIters) keyspanimpl.TableNewSpanIter {
+	return func(ctx context.Context, file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
 		iters, err := newIters(ctx, file, nil, internalIterOpts{}, iterRangeDeletions)
 		if err != nil {
 			return nil, err
@@ -76,10 +74,8 @@ func tableNewRangeDelIter(
 
 // tableNewRangeKeyIter takes a tableNewIters and returns a TableNewSpanIter
 // for the range key iterator returned by tableNewIters.
-func tableNewRangeKeyIter(
-	ctx context.Context, newIters tableNewIters,
-) keyspanimpl.TableNewSpanIter {
-	return func(file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
+func tableNewRangeKeyIter(newIters tableNewIters) keyspanimpl.TableNewSpanIter {
+	return func(ctx context.Context, file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
 		iters, err := newIters(ctx, file, nil, internalIterOpts{}, iterRangeKeys)
 		if err != nil {
 			return nil, err


Test patch:
diff --git a/db_test.go b/db_test.go
index 69217ec1..5fd308b3 100644
--- a/db_test.go
+++ b/db_test.go
@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"io"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"slices"
 	"strconv"
@@ -24,6 +25,8 @@ import (
 	"github.com/cockroachdb/fifo"
 	"github.com/cockroachdb/pebble/internal/base"
 	"github.com/cockroachdb/pebble/internal/cache"
+	"github.com/cockroachdb/pebble/internal/testkeys"
+	"github.com/cockroachdb/pebble/internal/testutils"
 	"github.com/cockroachdb/pebble/objstorage/objstorageprovider"
 	"github.com/cockroachdb/pebble/sstable"
 	"github.com/cockroachdb/pebble/vfs"
@@ -1405,8 +1408,11 @@ func (t *testTracer) Eventf(ctx context.Context, format string, args ...interfac
 	if t.enabledOnlyForNonBackgroundContext && ctx == context.Background() {
 		return
 	}
-	fmt.Fprintf(&t.buf, format, args...)
-	fmt.Fprint(&t.buf, "\n")
+	str := fmt.Sprintf(format, args...)
+	// Redact known strings that depend on source code line numbers.
+	str = regexp.MustCompile(`\(fileNum=[^)]+\)$`).ReplaceAllString(str, "(<redacted>)")
+	t.buf.WriteString(str)
+	t.buf.WriteString("\n")
 }
 
 func (t *testTracer) IsTracingEnabled(ctx context.Context) bool {
@@ -1418,68 +1424,76 @@ func (t *testTracer) IsTracingEnabled(ctx context.Context) bool {
 
 func TestTracing(t *testing.T) {
 	defer sstable.DeterministicReadBlockDurationForTesting()()
+
 	var tracer testTracer
-	c := NewCache(0)
-	defer c.Unref()
-	d, err := Open("", &Options{
-		FS:              vfs.NewMem(),
-		Cache:           c,
-		LoggerAndTracer: &tracer,
-	})
-	require.NoError(t, err)
+	buf := &tracer.buf
+	var db *DB
 	defer func() {
-		require.NoError(t, d.Close())
+		if db != nil {
+			db.Close()
+		}
 	}()
-
-	// Create a sstable.
-	require.NoError(t, d.Set([]byte("hello"), nil, nil))
-	require.NoError(t, d.Flush())
-	_, closer, err := d.Get([]byte("hello"))
-	require.NoError(t, err)
-	closer.Close()
-	readerInitRegexp := `reading footer of 53 bytes took 5ms
-reading block of 37 bytes took 5ms \([^)]*\)
-reading block of 419 bytes took 5ms \([^)]*\)
-`
-	iterTraceRegexp :=
-		`reading block of 27 bytes took 5ms \([^)]*\)
-reading block of 29 bytes took 5ms \([^)]*\)
-`
-
-	require.Regexp(t, readerInitRegexp+iterTraceRegexp, tracer.buf.String())
-
-	// Get again, but since it currently uses context.Background(), no trace
-	// output is produced.
-	tracer.buf.Reset()
-	tracer.enabledOnlyForNonBackgroundContext = true
-	_, closer, err = d.Get([]byte("hello"))
-	require.NoError(t, err)
-	closer.Close()
-	require.Equal(t, "", tracer.buf.String())
-
+	cache := NewCache(0)
+	defer cache.Unref()
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
-	iter, _ := d.NewIterWithContext(ctx, nil)
-	iter.SeekGE([]byte("hello"))
-	iter.Close()
-	require.Regexp(t, iterTraceRegexp, tracer.buf.String())
-
-	tracer.buf.Reset()
-	snap := d.NewSnapshot()
-	iter, _ = snap.NewIterWithContext(ctx, nil)
-	iter.SeekGE([]byte("hello"))
-	iter.Close()
-	require.Regexp(t, iterTraceRegexp, tracer.buf.String())
-	snap.Close()
-
-	tracer.buf.Reset()
-	b := d.NewIndexedBatch()
-	iter, err = b.NewIterWithContext(ctx, nil)
-	require.NoError(t, err)
-	iter.SeekGE([]byte("hello"))
-	iter.Close()
-	require.Regexp(t, iterTraceRegexp, tracer.buf.String())
-	b.Close()
+	datadriven.RunTest(t, "testdata/tracing", func(t *testing.T, td *datadriven.TestData) string {
+		buf.Reset()
+		tracer.enabledOnlyForNonBackgroundContext = td.HasArg("disable-background-tracing")
+		ctx := ctx
+		switch td.Cmd {
+		case "build":
+			if db != nil {
+				db.Close()
+			}
+			db = testutils.CheckErr(Open("", &Options{
+				FS:              vfs.NewMem(),
+				Comparer:        testkeys.Comparer,
+				Cache:           cache,
+				LoggerAndTracer: &tracer,
+			}))
+
+			b := db.NewBatch()
+			require.NoError(t, runBatchDefineCmd(td, b))
+			require.NoError(t, b.Commit(nil))
+			require.NoError(t, db.Flush())
+			return ""
+
+		case "get":
+			for _, key := range strings.Split(td.Input, "\n") {
+				v, closer, err := db.Get([]byte(key))
+				require.NoError(t, err)
+				fmt.Fprintf(buf, "%s:%s\n", key, v)
+				closer.Close()
+			}
+
+		case "iter":
+			iter, _ := db.NewIterWithContext(ctx, &IterOptions{
+				KeyTypes: IterKeyTypePointsAndRanges,
+			})
+			buf.WriteString(runIterCmd(td, iter, true))
+
+		case "snapshot-iter":
+			snap := db.NewSnapshot()
+			defer snap.Close()
+			iter, _ := snap.NewIterWithContext(ctx, &IterOptions{
+				KeyTypes: IterKeyTypePointsAndRanges,
+			})
+			buf.WriteString(runIterCmd(td, iter, true))
+
+		case "indexed-batch-iter":
+			b := db.NewIndexedBatch()
+			defer b.Close()
+			iter, _ := b.NewIterWithContext(ctx, &IterOptions{
+				KeyTypes: IterKeyTypePointsAndRanges,
+			})
+			buf.WriteString(runIterCmd(td, iter, true))
+
+		default:
+			td.Fatalf(t, "unknown command: %s", td.Cmd)
+		}
+		return buf.String()
+	})
 }
 
 func TestMemtableIngestInversion(t *testing.T) {
diff --git a/internal/keyspan/keyspanimpl/level_iter_test.go b/internal/keyspan/keyspanimpl/level_iter_test.go
index 955fe5c2..3ef71a12 100644
--- a/internal/keyspan/keyspanimpl/level_iter_test.go
+++ b/internal/keyspan/keyspanimpl/level_iter_test.go
@@ -5,6 +5,7 @@
 package keyspanimpl
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"testing"
@@ -301,7 +302,7 @@ func TestLevelIterEquivalence(t *testing.T) {
 				metas = append(metas, meta)
 			}
 
-			tableNewIters := func(file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
+			tableNewIters := func(ctx context.Context, file *manifest.FileMetadata, iterOptions keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
 				return keyspan.NewIter(base.DefaultComparer.Compare, tc.levels[j][file.FileNum-1]), nil
 			}
 			// Add all the fileMetadatas to L6.
@@ -314,6 +315,7 @@ func TestLevelIterEquivalence(t *testing.T) {
 			v, err := b.Apply(nil, base.DefaultComparer, 0, 0)
 			require.NoError(t, err)
 			levelIter.Init(
+				context.Background(),
 				keyspan.SpanIterOptions{}, base.DefaultComparer.Compare, tableNewIters,
 				v.Levels[6].Iter(), 0, manifest.KeyTypeRange,
 			)
@@ -414,7 +416,7 @@ func TestLevelIter(t *testing.T) {
 					keyType = manifest.KeyTypePoint
 				}
 			}
-			tableNewIters := func(file *manifest.FileMetadata, _ keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
+			tableNewIters := func(ctx context.Context, file *manifest.FileMetadata, _ keyspan.SpanIterOptions) (keyspan.FragmentIterator, error) {
 				f := files[file.FileNum-1]
 				if keyType == manifest.KeyTypePoint {
 					return keyspan.NewIter(cmp, f.rangeDels), nil
@@ -426,7 +428,7 @@ func TestLevelIter(t *testing.T) {
 				metas[i] = files[i].meta
 			}
 			lm := manifest.MakeLevelMetadata(cmp, 6, metas)
-			iter := NewLevelIter(keyspan.SpanIterOptions{}, cmp, tableNewIters, lm.Iter(), 6, keyType)
+			iter := NewLevelIter(context.Background(), keyspan.SpanIterOptions{}, cmp, tableNewIters, lm.Iter(), 6, keyType)
 			extraInfo := func() string {
 				return iter.String()
 			}
diff --git a/internal/keyspan/test_utils.go b/internal/keyspan/test_utils.go
index bbff49a7..060aadd3 100644
--- a/internal/keyspan/test_utils.go
+++ b/internal/keyspan/test_utils.go
@@ -6,6 +6,7 @@ package keyspan
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"go/token"
 	"io"
@@ -365,6 +366,11 @@ func (p *probeIterator) Prev() (*Span, error) {
 	return p.handleOp(op)
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (p *probeIterator) SetContext(ctx context.Context) {
+	p.iter.SetContext(ctx)
+}
+
 func (p *probeIterator) Close() {
 	op := op{Kind: OpClose}
 	if p.iter != nil {
@@ -561,6 +567,11 @@ func (i *invalidatingIter) Last() (*Span, error)             { return i.invalida
 func (i *invalidatingIter) Next() (*Span, error)             { return i.invalidate(i.iter.Next()) }
 func (i *invalidatingIter) Prev() (*Span, error)             { return i.invalidate(i.iter.Prev()) }
 
+// SetContext is part of the FragmentIterator interface.
+func (i *invalidatingIter) SetContext(ctx context.Context) {
+	i.iter.SetContext(ctx)
+}
+
 func (i *invalidatingIter) Close() {
 	_, _ = i.invalidate(nil, nil)
 	i.iter.Close()
diff --git a/keyspan_probe_test.go b/keyspan_probe_test.go
index 759ca3b2..afee56a4 100644
--- a/keyspan_probe_test.go
+++ b/keyspan_probe_test.go
@@ -5,6 +5,7 @@
 package pebble
 
 import (
+	"context"
 	"fmt"
 	"go/token"
 	"io"
@@ -377,6 +378,11 @@ func (p *probeKeyspanIterator) DebugTree(tp treeprinter.Node) {
 	}
 }
 
+// SetContext is part of the FragmentIterator interface.
+func (p *probeKeyspanIterator) SetContext(ctx context.Context) {
+	p.iter.SetContext(ctx)
+}
+
 func (p *probeKeyspanIterator) Close() {
 	op := keyspanOp{Kind: opSpanClose}
 	if p.iter != nil {
diff --git a/testdata/tracing b/testdata/tracing
new file mode 100644
index 00000000..b24f8aef
--- /dev/null
+++ b/testdata/tracing
@@ -0,0 +1,88 @@
+build
+set hello foo
+----
+
+get
+hello
+----
+reading footer of 53 bytes took 5ms
+reading block of 37 bytes took 5ms (<redacted>)
+reading block of 417 bytes took 5ms (<redacted>)
+reading block of 31 bytes took 5ms (<redacted>)
+reading block of 32 bytes took 5ms (<redacted>)
+hello:foo
+
+# If we disable background tracing, we should see no traces with Get.
+get disable-background-tracing
+hello
+----
+hello:foo
+
+iter
+seek-ge hello
+----
+reading block of 31 bytes took 5ms (<redacted>)
+reading block of 32 bytes took 5ms (<redacted>)
+hello: (foo, .)
+
+snapshot-iter
+seek-ge hello
+----
+reading block of 31 bytes took 5ms (<redacted>)
+reading block of 32 bytes took 5ms (<redacted>)
+hello: (foo, .)
+
+indexed-batch-iter
+seek-ge hello
+----
+reading block of 31 bytes took 5ms (<redacted>)
+reading block of 32 bytes took 5ms (<redacted>)
+hello: (foo, .)
+
+build
+set hello@2 foo1
+range-key-set a z @4 foo2
+----
+
+# This operation should not read the range key block.
+get
+hello@2
+----
+reading footer of 53 bytes took 5ms
+reading block of 62 bytes took 5ms (<redacted>)
+reading block of 504 bytes took 5ms (<redacted>)
+reading block of 27 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+hello@2:foo1
+
+# These operations should read the range key block, which shows up as a third
+# log.
+iter
+seek-ge hello
+next
+----
+reading block of 27 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+hello: (., [a-z) @4=foo2 UPDATED)
+hello@2: (foo1, [a-z) @4=foo2)
+
+snapshot-iter
+seek-ge hello
+next
+----
+reading block of 27 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+hello: (., [a-z) @4=foo2 UPDATED)
+hello@2: (foo1, [a-z) @4=foo2)
+
+indexed-batch-iter
+seek-ge hello
+next
+----
+reading block of 27 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+reading block of 35 bytes took 5ms (<redacted>)
+hello: (., [a-z) @4=foo2 UPDATED)
+hello@2: (foo1, [a-z) @4=foo2)