ceremonyclient/pebble/internal/keyspan/interleaving_iter_test.go

292 lines
7.9 KiB
Go
Raw Permalink Normal View History

2024-01-03 07:31:42 +00:00
// Copyright 2021 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.
package keyspan
import (
"bytes"
"context"
"fmt"
"io"
"sort"
"strings"
"testing"
"github.com/cockroachdb/datadriven"
"github.com/cockroachdb/pebble/internal/base"
"github.com/cockroachdb/pebble/internal/testkeys"
"github.com/stretchr/testify/require"
)
func TestInterleavingIter(t *testing.T) {
runInterleavingIterTest(t, "testdata/interleaving_iter")
}
func TestInterleavingIter_Masking(t *testing.T) {
runInterleavingIterTest(t, "testdata/interleaving_iter_masking")
}
type maskingHooks struct {
log io.Writer
cmp base.Compare
split base.Split
threshold []byte
maskSuffix []byte
}
func (m *maskingHooks) SpanChanged(s *Span) {
if m.log != nil {
if s == nil {
fmt.Fprintln(m.log, "-- SpanChanged(nil)")
} else {
fmt.Fprintf(m.log, "-- SpanChanged(%s)\n", s)
}
}
// Find the smallest suffix of a key contained within the Span, excluding
// suffixes less than m.threshold.
m.maskSuffix = nil
if s == nil || m.threshold == nil || len(s.Keys) == 0 {
return
}
for i := range s.Keys {
if s.Keys[i].Suffix == nil {
continue
}
if m.cmp(s.Keys[i].Suffix, m.threshold) < 0 {
continue
}
if m.maskSuffix == nil || m.cmp(m.maskSuffix, s.Keys[i].Suffix) > 0 {
m.maskSuffix = s.Keys[i].Suffix
}
}
}
func (m *maskingHooks) SkipPoint(userKey []byte) bool {
pointSuffix := userKey[m.split(userKey):]
return m.maskSuffix != nil && len(pointSuffix) > 0 && m.cmp(m.maskSuffix, pointSuffix) < 0
}
func runInterleavingIterTest(t *testing.T, filename string) {
cmp := testkeys.Comparer.Compare
var keyspanIter MergingIter
var pointIter pointIterator
var iter InterleavingIter
var buf bytes.Buffer
hooks := maskingHooks{
log: &buf,
cmp: testkeys.Comparer.Compare,
split: testkeys.Comparer.Split,
}
var prevKey *base.InternalKey
formatKey := func(k *base.InternalKey, _ base.LazyValue) {
if k == nil {
fmt.Fprint(&buf, ".")
return
}
prevKey = k
s := iter.Span()
fmt.Fprintf(&buf, "PointKey: %s\n", k.String())
if s != nil {
fmt.Fprintf(&buf, "Span: %s\n-", s)
} else {
fmt.Fprintf(&buf, "Span: %s\n-", Span{})
}
}
datadriven.RunTest(t, filename, func(t *testing.T, td *datadriven.TestData) string {
buf.Reset()
switch td.Cmd {
case "set-masking-threshold":
hooks.threshold = []byte(strings.TrimSpace(td.Input))
return "OK"
case "define-rangekeys":
var spans []Span
lines := strings.Split(strings.TrimSpace(td.Input), "\n")
for _, line := range lines {
spans = append(spans, ParseSpan(line))
}
keyspanIter.Init(cmp, noopTransform, new(MergingBuffers), NewIter(cmp, spans))
hooks.maskSuffix = nil
iter.Init(testkeys.Comparer, &pointIter, &keyspanIter,
InterleavingIterOpts{Mask: &hooks})
return "OK"
case "define-pointkeys":
var points []base.InternalKey
lines := strings.Split(strings.TrimSpace(td.Input), "\n")
for _, line := range lines {
points = append(points, base.ParseInternalKey(line))
}
pointIter = pointIterator{cmp: cmp, keys: points}
hooks.maskSuffix = nil
iter.Init(testkeys.Comparer, &pointIter, &keyspanIter,
InterleavingIterOpts{Mask: &hooks})
return "OK"
case "iter":
buf.Reset()
// Clear any previous bounds.
iter.SetBounds(nil, nil)
prevKey = nil
lines := strings.Split(strings.TrimSpace(td.Input), "\n")
for _, line := range lines {
bufLen := buf.Len()
line = strings.TrimSpace(line)
i := strings.IndexByte(line, ' ')
iterCmd := line
if i > 0 {
iterCmd = string(line[:i])
}
switch iterCmd {
case "first":
formatKey(iter.First())
case "last":
formatKey(iter.Last())
case "next":
formatKey(iter.Next())
case "next-prefix":
succKey := testkeys.Comparer.ImmediateSuccessor(nil, prevKey.UserKey[:testkeys.Comparer.Split(prevKey.UserKey)])
formatKey(iter.NextPrefix(succKey))
case "prev":
formatKey(iter.Prev())
case "seek-ge":
formatKey(iter.SeekGE([]byte(strings.TrimSpace(line[i:])), base.SeekGEFlagsNone))
case "seek-prefix-ge":
key := []byte(strings.TrimSpace(line[i:]))
prefix := key[:testkeys.Comparer.Split(key)]
formatKey(iter.SeekPrefixGE(prefix, key, base.SeekGEFlagsNone))
case "seek-lt":
formatKey(iter.SeekLT([]byte(strings.TrimSpace(line[i:])), base.SeekLTFlagsNone))
case "set-bounds":
bounds := strings.Fields(line[i:])
if len(bounds) != 2 {
return fmt.Sprintf("set-bounds expects 2 bounds, got %d", len(bounds))
}
l, u := []byte(bounds[0]), []byte(bounds[1])
if bounds[0] == "." {
l = nil
}
if bounds[1] == "." {
u = nil
}
iter.SetBounds(l, u)
default:
return fmt.Sprintf("unrecognized iter command %q", iterCmd)
}
require.NoError(t, iter.Error())
if buf.Len() > bufLen {
fmt.Fprintln(&buf)
}
}
return strings.TrimSpace(buf.String())
default:
return fmt.Sprintf("unrecognized command %q", td.Cmd)
}
})
require.NoError(t, iter.Close())
}
type pointIterator struct {
cmp base.Compare
keys []base.InternalKey
lower []byte
upper []byte
index int
}
var _ base.InternalIterator = &pointIterator{}
func (i *pointIterator) SeekGE(
key []byte, flags base.SeekGEFlags,
) (*base.InternalKey, base.LazyValue) {
i.index = sort.Search(len(i.keys), func(j int) bool {
return i.cmp(i.keys[j].UserKey, key) >= 0
})
if i.index < 0 || i.index >= len(i.keys) {
return nil, base.LazyValue{}
}
if i.upper != nil && i.cmp(i.keys[i.index].UserKey, i.upper) >= 0 {
return nil, base.LazyValue{}
}
return &i.keys[i.index], base.LazyValue{}
}
func (i *pointIterator) SeekPrefixGE(
prefix, key []byte, flags base.SeekGEFlags,
) (*base.InternalKey, base.LazyValue) {
return i.SeekGE(key, flags)
}
func (i *pointIterator) SeekLT(
key []byte, flags base.SeekLTFlags,
) (*base.InternalKey, base.LazyValue) {
i.index = sort.Search(len(i.keys), func(j int) bool {
return i.cmp(i.keys[j].UserKey, key) >= 0
})
i.index--
if i.index < 0 || i.index >= len(i.keys) {
return nil, base.LazyValue{}
}
if i.lower != nil && i.cmp(i.keys[i.index].UserKey, i.lower) < 0 {
return nil, base.LazyValue{}
}
return &i.keys[i.index], base.LazyValue{}
}
func (i *pointIterator) First() (*base.InternalKey, base.LazyValue) {
i.index = 0
if i.index < 0 || i.index >= len(i.keys) {
return nil, base.LazyValue{}
}
if i.upper != nil && i.cmp(i.keys[i.index].UserKey, i.upper) >= 0 {
return nil, base.LazyValue{}
}
return &i.keys[i.index], base.LazyValue{}
}
func (i *pointIterator) Last() (*base.InternalKey, base.LazyValue) {
i.index = len(i.keys) - 1
if i.index < 0 || i.index >= len(i.keys) {
return nil, base.LazyValue{}
}
if i.lower != nil && i.cmp(i.keys[i.index].UserKey, i.lower) < 0 {
return nil, base.LazyValue{}
}
return &i.keys[i.index], base.LazyValue{}
}
func (i *pointIterator) Next() (*base.InternalKey, base.LazyValue) {
i.index++
if i.index < 0 || i.index >= len(i.keys) {
return nil, base.LazyValue{}
}
if i.upper != nil && i.cmp(i.keys[i.index].UserKey, i.upper) >= 0 {
return nil, base.LazyValue{}
}
return &i.keys[i.index], base.LazyValue{}
}
func (i *pointIterator) NextPrefix(succKey []byte) (*base.InternalKey, base.LazyValue) {
return i.SeekGE(succKey, base.SeekGEFlagsNone)
}
func (i *pointIterator) Prev() (*base.InternalKey, base.LazyValue) {
i.index--
if i.index < 0 || i.index >= len(i.keys) {
return nil, base.LazyValue{}
}
if i.lower != nil && i.cmp(i.keys[i.index].UserKey, i.lower) < 0 {
return nil, base.LazyValue{}
}
return &i.keys[i.index], base.LazyValue{}
}
func (i *pointIterator) Close() error { return nil }
func (i *pointIterator) Error() error { return nil }
func (i *pointIterator) String() string { return "test-point-iterator" }
func (i *pointIterator) SetBounds(lower, upper []byte) {
i.lower, i.upper = lower, upper
}
func (i *pointIterator) SetContext(_ context.Context) {}