Repository: cockroachdb/pebble
Pull request #3795 (2024-07-26 19:59:42)

db: strictly enforce prefix in [flushable]BatchIter.SeekPrefixGE
Every key we avoid propagating up the iterator stack avoids a key comparison
during initialization of the merging iterator heap.

Informs #3794.

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/3795)
<!-- Reviewable:end -->

jbowens: TFTR!

Patch:
diff --git a/batch.go b/batch.go
index 65b0fe3c..556874e2 100644
--- a/batch.go
+++ b/batch.go
@@ -5,6 +5,7 @@
 package pebble
 
 import (
+	"bytes"
 	"context"
 	"encoding/binary"
 	"fmt"
@@ -270,11 +271,9 @@ type batchInternal struct {
 	// batches. Large batches will set the data field to nil when committed as
 	// the data has been moved to a flushableBatch and inserted into the queue of
 	// memtables.
-	data           []byte
-	cmp            Compare
-	formatKey      base.FormatKey
-	abbreviatedKey AbbreviatedKey
-	opts           batchOptions
+	data     []byte
+	comparer *base.Comparer
+	opts     batchOptions
 
 	// An upper bound on required space to add this batch to a memtable.
 	// Note that although batches are limited to 4 GiB in size, that limit
@@ -464,12 +463,10 @@ func newBatchWithSize(db *DB, size int, opts ...BatchOption) *Batch {
 
 func newIndexedBatch(db *DB, comparer *Comparer) *Batch {
 	i := indexedBatchPool.Get().(*indexedBatch)
-	i.batch.cmp = comparer.Compare
-	i.batch.formatKey = comparer.FormatKey
-	i.batch.abbreviatedKey = comparer.AbbreviatedKey
+	i.batch.comparer = comparer
 	i.batch.db = db
 	i.batch.index = &i.index
-	i.batch.index.Init(&i.batch.data, i.batch.cmp, i.batch.abbreviatedKey)
+	i.batch.index.Init(&i.batch.data, comparer.Compare, comparer.AbbreviatedKey)
 	i.batch.opts.ensureDefaults()
 	return &i.batch
 }
@@ -508,9 +505,7 @@ func (b *Batch) release() {
 	// field. Without using an atomic to clear that field the Go race detector
 	// complains.
 	b.reset()
-	b.cmp = nil
-	b.formatKey = nil
-	b.abbreviatedKey = nil
+	b.comparer = nil
 
 	if b.index == nil {
 		batchPool.Put(b)
@@ -637,14 +632,14 @@ func (b *Batch) Apply(batch *Batch, _ *WriteOptions) error {
 					b.tombstones = nil
 					b.tombstonesSeqNum = 0
 					if b.rangeDelIndex == nil {
-						b.rangeDelIndex = batchskl.NewSkiplist(&b.data, b.cmp, b.abbreviatedKey)
+						b.rangeDelIndex = batchskl.NewSkiplist(&b.data, b.comparer.Compare, b.comparer.AbbreviatedKey)
 					}
 					err = b.rangeDelIndex.Add(uint32(offset))
 				case InternalKeyKindRangeKeySet, InternalKeyKindRangeKeyUnset, InternalKeyKindRangeKeyDelete:
 					b.rangeKeys = nil
 					b.rangeKeysSeqNum = 0
 					if b.rangeKeyIndex == nil {
-						b.rangeKeyIndex = batchskl.NewSkiplist(&b.data, b.cmp, b.abbreviatedKey)
+						b.rangeKeyIndex = batchskl.NewSkiplist(&b.data, b.comparer.Compare, b.comparer.AbbreviatedKey)
 					}
 					err = b.rangeKeyIndex.Add(uint32(offset))
 				default:
@@ -1014,7 +1009,7 @@ func (b *Batch) DeleteRangeDeferred(startLen, endLen int) *DeferredBatchOp {
 		b.tombstonesSeqNum = 0
 		// Range deletions are rare, so we lazily allocate the index for them.
 		if b.rangeDelIndex == nil {
-			b.rangeDelIndex = batchskl.NewSkiplist(&b.data, b.cmp, b.abbreviatedKey)
+			b.rangeDelIndex = batchskl.NewSkiplist(&b.data, b.comparer.Compare, b.comparer.AbbreviatedKey)
 		}
 		b.deferredOp.index = b.rangeDelIndex
 	}
@@ -1069,7 +1064,7 @@ func (b *Batch) incrementRangeKeysCount() {
 		b.rangeKeysSeqNum = 0
 		// Range keys are rare, so we lazily allocate the index for them.
 		if b.rangeKeyIndex == nil {
-			b.rangeKeyIndex = batchskl.NewSkiplist(&b.data, b.cmp, b.abbreviatedKey)
+			b.rangeKeyIndex = batchskl.NewSkiplist(&b.data, b.comparer.Compare, b.comparer.AbbreviatedKey)
 		}
 		b.deferredOp.index = b.rangeKeyIndex
 	}
@@ -1283,7 +1278,6 @@ func (b *Batch) newInternalIter(o *IterOptions) *batchIter {
 
 func (b *Batch) initInternalIter(o *IterOptions, iter *batchIter) {
 	*iter = batchIter{
-		cmp:   b.cmp,
 		batch: b,
 		iter:  b.index.NewIter(o.GetLowerBound(), o.GetUpperBound()),
 		// NB: We explicitly do not propagate the batch snapshot to the point
@@ -1321,7 +1315,7 @@ func (b *Batch) newRangeDelIter(o *IterOptions, batchSnapshot base.SeqNum) *keys
 
 func (b *Batch) initRangeDelIter(_ *IterOptions, iter *keyspan.Iter, batchSnapshot base.SeqNum) {
 	if b.rangeDelIndex == nil {
-		iter.Init(b.cmp, nil)
+		iter.Init(b.comparer.Compare, nil)
 		return
 	}
 
@@ -1335,26 +1329,25 @@ func (b *Batch) initRangeDelIter(_ *IterOptions, iter *keyspan.Iter, batchSnapsh
 	// cleared.
 	nextSeqNum := b.nextSeqNum()
 	if b.tombstones != nil && b.tombstonesSeqNum <= batchSnapshot {
-		iter.Init(b.cmp, b.tombstones)
+		iter.Init(b.comparer.Compare, b.tombstones)
 		return
 	}
 
 	tombstones := make([]keyspan.Span, 0, b.countRangeDels)
 	frag := &keyspan.Fragmenter{
-		Cmp:    b.cmp,
-		Format: b.formatKey,
+		Cmp:    b.comparer.Compare,
+		Format: b.comparer.FormatKey,
 		Emit: func(s keyspan.Span) {
 			tombstones = append(tombstones, s)
 		},
 	}
 	it := &batchIter{
-		cmp:      b.cmp,
 		batch:    b,
 		iter:     b.rangeDelIndex.NewIter(nil, nil),
 		snapshot: batchSnapshot,
 	}
 	fragmentRangeDels(frag, it, int(b.countRangeDels))
-	iter.Init(b.cmp, tombstones)
+	iter.Init(b.comparer.Compare, tombstones)
 
 	// If we just read all the tombstones in the batch (eg, batchSnapshot was
 	// set to b.nextSeqNum()), then cache the tombstones so that a subsequent
@@ -1397,7 +1390,7 @@ func (b *Batch) newRangeKeyIter(o *IterOptions, batchSnapshot base.SeqNum) *keys
 
 func (b *Batch) initRangeKeyIter(_ *IterOptions, iter *keyspan.Iter, batchSnapshot base.SeqNum) {
 	if b.rangeKeyIndex == nil {
-		iter.Init(b.cmp, nil)
+		iter.Init(b.comparer.Compare, nil)
 		return
 	}
 
@@ -1410,26 +1403,25 @@ func (b *Batch) initRangeKeyIter(_ *IterOptions, iter *keyspan.Iter, batchSnapsh
 	// sequence number the cache would've been cleared.
 	nextSeqNum := b.nextSeqNum()
 	if b.rangeKeys != nil && b.rangeKeysSeqNum <= batchSnapshot {
-		iter.Init(b.cmp, b.rangeKeys)
+		iter.Init(b.comparer.Compare, b.rangeKeys)
 		return
 	}
 
 	rangeKeys := make([]keyspan.Span, 0, b.countRangeKeys)
 	frag := &keyspan.Fragmenter{
-		Cmp:    b.cmp,
-		Format: b.formatKey,
+		Cmp:    b.comparer.Compare,
+		Format: b.comparer.FormatKey,
 		Emit: func(s keyspan.Span) {
 			rangeKeys = append(rangeKeys, s)
 		},
 	}
 	it := &batchIter{
-		cmp:      b.cmp,
 		batch:    b,
 		iter:     b.rangeKeyIndex.NewIter(nil, nil),
 		snapshot: batchSnapshot,
 	}
 	fragmentRangeKeys(frag, it, int(b.countRangeKeys))
-	iter.Init(b.cmp, rangeKeys)
+	iter.Init(b.comparer.Compare, rangeKeys)
 
 	// If we just read all the range keys in the batch (eg, batchSnapshot was
 	// set to b.nextSeqNum()), then cache the range keys so that a subsequent
@@ -1553,13 +1545,11 @@ func (b *Batch) reset() {
 	// Zero out the struct, retaining only the fields necessary for manual
 	// reuse.
 	b.batchInternal = batchInternal{
-		data:           b.data,
-		cmp:            b.cmp,
-		formatKey:      b.formatKey,
-		abbreviatedKey: b.abbreviatedKey,
-		opts:           b.opts,
-		index:          b.index,
-		db:             b.db,
+		data:     b.data,
+		comparer: b.comparer,
+		opts:     b.opts,
+		index:    b.index,
+		db:       b.db,
 	}
 	b.applied.Store(false)
 	if b.data != nil {
@@ -1576,7 +1566,7 @@ func (b *Batch) reset() {
 		}
 	}
 	if b.index != nil {
-		b.index.Init(&b.data, b.cmp, b.abbreviatedKey)
+		b.index.Init(&b.data, b.comparer.Compare, b.comparer.AbbreviatedKey)
 	}
 }
 
@@ -1658,7 +1648,6 @@ func (b *Batch) CommitStats() BatchCommitStats {
 // Note: batchIter mirrors the implementation of flushableBatchIter. Keep the
 // two in sync.
 type batchIter struct {
-	cmp   Compare
 	batch *Batch
 	iter  batchskl.Iterator
 	kv    base.InternalKV
@@ -1698,8 +1687,16 @@ func (i *batchIter) SeekGE(key []byte, flags base.SeekGEFlags) *base.InternalKV
 }
 
 func (i *batchIter) SeekPrefixGE(prefix, key []byte, flags base.SeekGEFlags) *base.InternalKV {
-	i.err = nil // clear cached iteration error
-	return i.SeekGE(key, flags)
+	kv := i.SeekGE(key, flags)
+	if kv == nil {
+		return nil
+	}
+	// If the key doesn't have the sought prefix, return nil.
+	if !bytes.Equal(i.batch.comparer.Split.Prefix(kv.K.UserKey), prefix) {
+		i.kv = base.InternalKV{}
+		return nil
+	}
+	return kv
 }
 
 func (i *batchIter) SeekLT(key []byte, flags base.SeekLTFlags) *base.InternalKV {
@@ -1849,9 +1846,9 @@ type flushableBatchEntry struct {
 // flushableBatch wraps an existing batch and provides the interfaces needed
 // for making the batch flushable (i.e. able to mimic a memtable).
 type flushableBatch struct {
-	cmp       Compare
-	formatKey base.FormatKey
-	data      []byte
+	cmp      Compare
+	comparer *base.Comparer
+	data     []byte
 
 	// The base sequence number for the entries in the batch. This is the same
 	// value as Batch.seqNum() and is cached here for performance.
@@ -1883,10 +1880,10 @@ var _ flushable = (*flushableBatch)(nil)
 // of the batch data.
 func newFlushableBatch(batch *Batch, comparer *Comparer) (*flushableBatch, error) {
 	b := &flushableBatch{
-		data:      batch.data,
-		cmp:       comparer.Compare,
-		formatKey: comparer.FormatKey,
-		offsets:   make([]flushableBatchEntry, 0, batch.Count()),
+		data:     batch.data,
+		cmp:      comparer.Compare,
+		comparer: comparer,
+		offsets:  make([]flushableBatchEntry, 0, batch.Count()),
 	}
 	if b.data != nil {
 		// Note that this sequence number is not correct when this batch has not
@@ -1969,7 +1966,7 @@ func newFlushableBatch(batch *Batch, comparer *Comparer) (*flushableBatch, error
 	if len(rangeDelOffsets) > 0 {
 		frag := &keyspan.Fragmenter{
 			Cmp:    b.cmp,
-			Format: b.formatKey,
+			Format: b.comparer.FormatKey,
 			Emit: func(s keyspan.Span) {
 				b.tombstones = append(b.tombstones, s)
 			},
@@ -1986,7 +1983,7 @@ func newFlushableBatch(batch *Batch, comparer *Comparer) (*flushableBatch, error
 	if len(rangeKeyOffsets) > 0 {
 		frag := &keyspan.Fragmenter{
 			Cmp:    b.cmp,
-			Format: b.formatKey,
+			Format: b.comparer.FormatKey,
 			Emit: func(s keyspan.Span) {
 				b.rangeKeys = append(b.rangeKeys, s)
 			},
@@ -2173,7 +2170,15 @@ func (i *flushableBatchIter) SeekGE(key []byte, flags base.SeekGEFlags) *base.In
 func (i *flushableBatchIter) SeekPrefixGE(
 	prefix, key []byte, flags base.SeekGEFlags,
 ) *base.InternalKV {
-	return i.SeekGE(key, flags)
+	kv := i.SeekGE(key, flags)
+	if kv == nil {
+		return nil
+	}
+	// If the key doesn't have the sought prefix, return nil.
+	if !bytes.Equal(i.batch.comparer.Split.Prefix(kv.K.UserKey), prefix) {
+		return nil
+	}
+	return kv
 }
 
 // SeekLT implements internalIterator.SeekLT, as documented in the pebble


Test patch:
diff --git a/batch_test.go b/batch_test.go
index d7130d8f..5d2e3981 100644
--- a/batch_test.go
+++ b/batch_test.go
@@ -528,9 +528,7 @@ func TestIndexedBatchReset(t *testing.T) {
 	require.Equal(t, 1, indexCount(b.index))
 
 	b.Reset()
-	require.NotNil(t, b.cmp)
-	require.NotNil(t, b.formatKey)
-	require.NotNil(t, b.abbreviatedKey)
+	require.NotNil(t, b.comparer)
 	require.NotNil(t, b.index)
 	require.Nil(t, b.rangeDelIndex)
 	require.Nil(t, b.rangeKeyIndex)
diff --git a/data_test.go b/data_test.go
index c5d569f5..59c06a99 100644
--- a/data_test.go
+++ b/data_test.go
@@ -339,7 +339,7 @@ func printIterState(
 			hasPoint, hasRange := iter.HasPointAndRange()
 			fmt.Fprintf(b, "%s:%s (", iter.Key(), validityStateStr)
 			if hasPoint {
-				fmt.Fprintf(b, "%s, ", iter.Value())
+				fmt.Fprintf(b, "%s, ", formatASCIIValue(iter.Value()))
 			} else {
 				fmt.Fprint(b, "., ")
 			}
@@ -388,13 +388,24 @@ func formatASCIIKey(b []byte) string {
 	return string(b)
 }
 
+func formatASCIIValue(b []byte) string {
+	if len(b) > 1<<10 {
+		return fmt.Sprintf("[LARGE VALUE len=%d]", len(b))
+	}
+	if bytes.IndexFunc(b, func(r rune) bool { return r < '!' || r > 'z' }) != -1 {
+		// This key is not just legible ASCII characters. Quote it.
+		return fmt.Sprintf("%q", b)
+	}
+	return string(b)
+}
+
 func writeRangeKeys(b io.Writer, iter *Iterator) {
 	rangeKeys := iter.RangeKeys()
 	for j := 0; j < len(rangeKeys); j++ {
 		if j > 0 {
 			fmt.Fprint(b, ",")
 		}
-		fmt.Fprintf(b, " %s=%s", rangeKeys[j].Suffix, rangeKeys[j].Value)
+		fmt.Fprintf(b, " %s=%s", rangeKeys[j].Suffix, formatASCIIValue(rangeKeys[j].Value))
 	}
 }
 
diff --git a/iterator_histories_test.go b/iterator_histories_test.go
index 89b70760..2e51fa2f 100644
--- a/iterator_histories_test.go
+++ b/iterator_histories_test.go
@@ -192,6 +192,16 @@ func TestIterHistories(t *testing.T) {
 					return err.Error()
 				}
 				return ""
+			case "disable-flushes":
+				d.mu.Lock()
+				d.mu.compact.flushing = true
+				d.mu.Unlock()
+				return ""
+			case "enable-flushes":
+				d.mu.Lock()
+				d.mu.compact.flushing = false
+				d.mu.Unlock()
+				return ""
 			case "get":
 				var reader Reader = d
 				if arg, ok := td.Arg("reader"); ok {
diff --git a/testdata/iter_histories/prefix_iteration b/testdata/iter_histories/prefix_iteration
index 8816b627..19b24c57 100644
--- a/testdata/iter_histories/prefix_iteration
+++ b/testdata/iter_histories/prefix_iteration
@@ -382,3 +382,76 @@ stats
 ----
 .
 stats: seeked 1 times (1 internal); stepped 0 times (0 internal)
+
+# Test that a prefix seek through a batch iterator that enforces the prefix
+# strictly.
+
+reset
+----
+
+batch name=foo
+set b@1 b@1
+set d@9 d@9
+set g@4 g@4
+set e@2 e@2
+----
+wrote 4 keys to batch "foo"
+
+# The stats should indicate only 3 KVs were ever surfaced to the merging iterator.
+
+combined-iter reader=foo name=fooiter
+seek-prefix-ge b@10
+seek-prefix-ge c@10
+seek-prefix-ge d@10
+seek-prefix-ge g@2
+seek-prefix-ge e@2
+stats
+----
+b@1: (b@1, .)
+.
+d@9: (d@9, .)
+.
+e@2: (e@2, .)
+stats: seeked 5 times (5 internal); stepped 0 times (0 internal); blocks: 0B cached; points: 3 (9B keys, 9B values)
+
+# Test the above case but with a large committed batch (which should be a
+# flushableBatchIter).
+
+define memtable-size=65536
+----
+
+# We diable flushes to avoid scheduling a flush that might race with our
+# iterator. If the iterator observed the state after the large batch has been
+# flushed to sstables, we would see nonzero block bytes appear in the iterator
+# stats.
+disable-flushes
+----
+
+batch commit
+set b@1 <rand-bytes=10000>
+set d@9 <rand-bytes=10000>
+set g@4 <rand-bytes=10000>
+set e@2 <rand-bytes=10000>
+----
+committed 4 keys
+
+lsm
+----
+
+combined-iter
+seek-prefix-ge b@10
+seek-prefix-ge c@10
+seek-prefix-ge d@10
+seek-prefix-ge g@2
+seek-prefix-ge e@2
+stats
+----
+b@1: ([LARGE VALUE len=10000], .)
+.
+d@9: ([LARGE VALUE len=10000], .)
+.
+e@2: ([LARGE VALUE len=10000], .)
+stats: seeked 5 times (5 internal); stepped 0 times (0 internal); blocks: 0B cached; points: 3 (9B keys, 29KB values)
+
+enable-flushes
+----