From 6e92bbc59b8cc382cb88adeb47fb177611acb73c Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Tue, 20 Feb 2024 17:50:39 -0800
Subject: [PATCH] db: fix incorrect SeekPrefixGE call in merging iterator

We also add assertions checking the prefix at various levels.
---
 level_iter.go                     | 4 ++++
 merging_iter.go                   | 6 ++++++
 sstable/reader_iter_single_lvl.go | 3 +++
 3 files changed, 13 insertions(+)

diff --git a/level_iter.go b/level_iter.go
index adde41d3be..29c3058042 100644
--- 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
index d1028eb95b..ddd4d3bf8a 100644
--- 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
index 2cec8be64f..00b27aa39d 100644
--- 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.