diff --git a/get_iter.go b/get_iter.go index b90fa8ca3f..8c5c742d94 100644 --- a/get_iter.go +++ b/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{} diff --git a/iterator_test.go b/iterator_test.go index 8e77c8012b..64d885b0e3 100644 --- a/iterator_test.go +++ b/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 } diff --git a/level_iter.go b/level_iter.go index adde41d3be..37888d8c72 100644 --- a/level_iter.go +++ b/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 diff --git a/merging_iter.go b/merging_iter.go index d1028eb95b..4caf6815f3 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,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 { 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. diff --git a/testdata/merging_iter b/testdata/merging_iter index 928f13ece1..625c2cf892 100644 --- a/testdata/merging_iter +++ b/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 .