db: fix incorrect SeekPrefixGE call in merging iterator

We also add assertions checking the prefix at various levels.
cockroachdb · Feb 21, 2024 · 6e92bbc · 6e92bbc
1 parent 21f5b6f
commit 6e92bbc
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/level_iter.go b/level_iter.go
@@ -5,6 +5,7 @@
 package pebble
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"runtime/debug"
@@ -714,6 +715,9 @@ func (l *levelIter) SeekGE(key []byte, flags base.SeekGEFlags) (*InternalKey, ba
 func (l *levelIter) SeekPrefixGE(
 	prefix, key []byte, flags base.SeekGEFlags,
 ) (*base.InternalKey, base.LazyValue) {
+	if invariants.Enabled && (!bytes.HasPrefix(key, prefix) || len(prefix) != l.comparer.Split(key)) {
+		panic(fmt.Sprintf("invalid SeekPrefixGE prefix %q for key %q", prefix, key))
+	}
 	l.err = nil // clear cached iteration error
 	if l.boundaryContext != nil {
 		l.boundaryContext.isSyntheticIterBoundsKey = false

diff --git a/merging_iter.go b/merging_iter.go
@@ -1098,6 +1098,9 @@ func (m *mergingIter) seekGE(key []byte, level int, flags base.SeekGEFlags) erro
 				// was greater than or equal to m.lower, the new key will
 				// continue to be greater than or equal to m.lower.
 				key = l.tombstone.End
+				if m.prefix != nil && !bytes.Equal(m.prefix, key[:m.split(key)]) {
+					m.prefix = nil
+				}
 			}
 		}
 	}
@@ -1126,6 +1129,9 @@ func (m *mergingIter) SeekGE(key []byte, flags base.SeekGEFlags) (*InternalKey,
 func (m *mergingIter) SeekPrefixGE(
 	prefix, key []byte, flags base.SeekGEFlags,
 ) (*base.InternalKey, base.LazyValue) {
+	if invariants.Enabled && (!bytes.HasPrefix(key, prefix) || len(prefix) != m.split(key)) {
+		panic(fmt.Sprintf("invalid SeekPrefixGE prefix %q for key %q", prefix, key))
+	}
 	m.prefix = prefix
 	m.err = m.seekGE(key, 0 /* start level */, flags)
 	if m.err != nil {

diff --git a/sstable/reader_iter_single_lvl.go b/sstable/reader_iter_single_lvl.go
@@ -810,6 +810,9 @@ func (i *singleLevelIterator) seekGEHelper(
 func (i *singleLevelIterator) SeekPrefixGE(
 	prefix, key []byte, flags base.SeekGEFlags,
 ) (*base.InternalKey, base.LazyValue) {
+	if invariants.Enabled && !bytes.HasPrefix(key, prefix) {
+		panic(fmt.Sprintf("invalid SeekPrefixGE prefix %q for key %q", prefix, key))
+	}
 	if i.vState != nil {
 		// Callers of SeekPrefixGE aren't aware of virtual sstable bounds, so
 		// we may have to internally restrict the bounds.