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.

Author: RaduBerinde

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

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)           {}

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,
 	)

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()

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

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()

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()

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.

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() {}
 

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
 		}