mirror of
https://source.quilibrium.com/quilibrium/ceremonyclient.git
synced 2024-11-14 04:05:17 +00:00
292 lines
7.9 KiB
Go
292 lines
7.9 KiB
Go
|
// 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) {}
|