db: fix incorrect SeekPrefixGE call in merging iterator

We also add assertions checking the prefix at various levels.

get_iter.go

@@ -188,11 +188,12 @@ func (g *getIter) Next() (*InternalKey, base.LazyValue) {
 
 				// Compute the key prefix for bloom filtering if split function is
 				// specified, or use the user key as default.
-				prefix := g.key
 				if g.comparer.Split != nil {
-					prefix = g.key[:g.comparer.Split(g.key)]
+					prefix := g.key[:g.comparer.Split(g.key)]
+					g.iterKey, g.iterValue = g.iter.SeekPrefixGE(prefix, g.key, base.SeekGEFlagsNone)
+				} else {
+					g.iterKey, g.iterValue = g.iter.SeekGE(g.key, base.SeekGEFlagsNone)
 				}
-				g.iterKey, g.iterValue = g.iter.SeekPrefixGE(prefix, g.key, base.SeekGEFlagsNone)
 				if err := g.iter.Error(); err != nil {
 					g.err = err
 					return nil, base.LazyValue{}
@@ -230,12 +231,13 @@ func (g *getIter) Next() (*InternalKey, base.LazyValue) {
 		g.iter = &g.levelIter
 
 		// Compute the key prefix for bloom filtering if split function is
-		// specified, or use the user key as default.
-		prefix := g.key
+		// specified.
 		if g.comparer.Split != nil {
-			prefix = g.key[:g.comparer.Split(g.key)]
+			prefix := g.key[:g.comparer.Split(g.key)]
+			g.iterKey, g.iterValue = g.iter.SeekPrefixGE(prefix, g.key, base.SeekGEFlagsNone)
+		} else {
+			g.iterKey, g.iterValue = g.iter.SeekGE(g.key, base.SeekGEFlagsNone)
 		}
-		g.iterKey, g.iterValue = g.iter.SeekPrefixGE(prefix, g.key, base.SeekGEFlagsNone)
 		if err := g.iter.Error(); err != nil {
 			g.err = err
 			return nil, base.LazyValue{}

iterator_test.go

@@ -910,6 +910,9 @@ func TestIteratorSeekOpt(t *testing.T) {
 					return &minSeqNumPropertyCollector{}
 				},
 			}
+			comparer := *base.DefaultComparer
+			comparer.Split = func(a []byte) int { return len(a) }
+			opts.Comparer = &comparer
 
 			var err error
 			if d, err = runDBDefineCmd(td, opts); err != nil {
@@ -944,7 +947,6 @@ func TestIteratorSeekOpt(t *testing.T) {
 				}
 				iter, _ = snap.NewIter(nil)
 				iter.readSampling.forceReadSampling = true
-				iter.comparer.Split = func(a []byte) int { return len(a) }
 				iter.forceEnableSeekOpt = true
 				iter.merging.forceEnableSeekOpt = true
 			}

level_iter.go

@@ -5,6 +5,7 @@
 package pebble
 
 import (
+	"bytes"
 	"context"
 	"fmt"
 	"runtime/debug"
@@ -714,6 +715,11 @@ 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 {
+		if !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

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,11 @@ 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 {
+		if !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 {

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.

testdata/merging_iter

@@ -590,7 +590,7 @@ seek-prefix-ge d true
 ----
 a#10,1:a10
 .
-.
+d#10,1:d10
 d#10,1:d10
 d#10,1:d10
 
@@ -603,7 +603,7 @@ next
 ----
 a#10,1:a10
 .
-.
+d#10,1:d10
 d#10,1:d10
 .