ceremonyclient/pebble/internal/keyspan/merging_iter_test.go
Cassandra Heart 2e2a1e4789
v1.2.0 (#31)
2024-01-03 01:31:42 -06:00

253 lines
7.6 KiB
Go

// Copyright 2022 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"
"fmt"
"math/rand"
"slices"
"strconv"
"strings"
"testing"
"time"
"github.com/cockroachdb/datadriven"
"github.com/cockroachdb/pebble/internal/base"
"github.com/cockroachdb/pebble/internal/testkeys"
"github.com/stretchr/testify/require"
)
func TestMergingIter(t *testing.T) {
cmp := base.DefaultComparer.Compare
var definedIters []FragmentIterator
var buf bytes.Buffer
datadriven.RunTest(t, "testdata/merging_iter", func(t *testing.T, td *datadriven.TestData) string {
switch td.Cmd {
case "define":
definedIters = definedIters[:0]
lines := strings.Split(strings.TrimSpace(td.Input), "\n")
var spans []Span
for _, line := range lines {
if line == "--" {
definedIters = append(definedIters, &invalidatingIter{iter: NewIter(cmp, spans)})
spans = nil
continue
}
spans = append(spans, ParseSpan(line))
}
if len(spans) > 0 {
definedIters = append(definedIters, &invalidatingIter{iter: NewIter(cmp, spans)})
}
return fmt.Sprintf("%d levels", len(definedIters))
case "iter":
buf.Reset()
pctx := probeContext{log: &buf}
snapshot := base.InternalKeySeqNumMax
iters := slices.Clone(definedIters)
for _, cmdArg := range td.CmdArgs {
switch cmdArg.Key {
case "snapshot":
var err error
snapshot, err = strconv.ParseUint(cmdArg.Vals[0], 10, 64)
require.NoError(t, err)
case "probes":
// The first value indicates which of the merging iterator's
// child iterators is the target.
i, err := strconv.Atoi(cmdArg.Vals[0])
if err != nil {
return err.Error()
}
// The remaining values define probes to attach.
iters[i] = attachProbes(iters[i], pctx, parseProbes(cmdArg.Vals[1:]...)...)
default:
return fmt.Sprintf("unrecognized arg %q", cmdArg.Key)
}
}
var iter MergingIter
iter.Init(cmp, VisibleTransform(snapshot), new(MergingBuffers), iters...)
runIterCmd(t, td, &iter, &buf)
return buf.String()
default:
return fmt.Sprintf("unrecognized command %q", td.Cmd)
}
})
}
// TestMergingIter_FragmenterEquivalence tests for equivalence between the
// fragmentation performed on-the-fly by the MergingIter and the fragmentation
// performed by the Fragmenter.
//
// It does this by producing 1-10 levels of well-formed fragments. Generated
// fragments may overlap other levels arbitrarily, but within their level
// generated fragments may only overlap other fragments that share the same user
// key bounds.
//
// The test then feeds all the fragments, across all levels, into a Fragmenter
// and produces a Iter over those fragments. The test also constructs a
// MergingIter with a separate Iter for each level. It runs a random
// series of operations, applying each operation to both. It asserts that each
// operation has identical results on both iterators.
func TestMergingIter_FragmenterEquivalence(t *testing.T) {
seed := time.Now().UnixNano()
for i := int64(0); i < 10; i++ {
testFragmenterEquivalenceOnce(t, seed+i)
}
}
func TestMergingIter_FragmenterEquivalence_Seed(t *testing.T) {
// This test uses a fixed seed. It's useful to manually edit its seed when
// debugging a test failure of the variable-seed test.
const seed = 1644517830186873000
testFragmenterEquivalenceOnce(t, seed)
}
func testFragmenterEquivalenceOnce(t *testing.T, seed int64) {
cmp := testkeys.Comparer.Compare
rng := rand.New(rand.NewSource(seed))
t.Logf("seed = %d", seed)
// Use a key space of alphanumeric strings, with a random max length between
// 1-3. Repeat keys are more common at the lower max lengths.
ks := testkeys.Alpha(rng.Intn(3) + 1)
// Generate between 1 and 10 levels of fragment iterators.
levels := make([][]Span, rng.Intn(10)+1)
iters := make([]FragmentIterator, len(levels))
var allSpans []Span
var buf bytes.Buffer
for l := 0; l < len(levels); l++ {
fmt.Fprintf(&buf, "level %d: ", l)
for keyspaceStartIdx := int64(0); keyspaceStartIdx < ks.Count(); {
// Generate spans of lengths of up to a third of the keyspace.
spanStartIdx := keyspaceStartIdx + rng.Int63n(ks.Count()/3)
spanEndIdx := spanStartIdx + rng.Int63n(ks.Count()/3) + 1
if spanEndIdx < ks.Count() {
keyCount := uint64(rng.Intn(3) + 1)
s := Span{
Start: testkeys.Key(ks, spanStartIdx),
End: testkeys.Key(ks, spanEndIdx),
Keys: make([]Key, 0, keyCount),
}
for k := keyCount; k > 0; k-- {
seqNum := uint64((len(levels)-l)*3) + k
s.Keys = append(s.Keys, Key{
Trailer: base.MakeTrailer(seqNum, base.InternalKeyKindRangeKeySet),
})
}
if len(levels[l]) > 0 {
fmt.Fprint(&buf, ", ")
}
fmt.Fprintf(&buf, "%s", s)
levels[l] = append(levels[l], s)
allSpans = append(allSpans, s)
}
keyspaceStartIdx = spanEndIdx
}
iters[l] = &invalidatingIter{iter: NewIter(cmp, levels[l])}
fmt.Fprintln(&buf)
}
// Fragment the spans across the levels.
var allFragmented []Span
f := Fragmenter{
Cmp: cmp,
Format: testkeys.Comparer.FormatKey,
Emit: func(span Span) {
allFragmented = append(allFragmented, span)
},
}
Sort(f.Cmp, allSpans)
for _, s := range allSpans {
f.Add(s)
}
f.Finish()
// Log all the levels and their fragments, as well as the fully-fragmented
// spans produced by the Fragmenter.
fmt.Fprintln(&buf, "Fragmenter fragments:")
for i, s := range allFragmented {
if i > 0 {
fmt.Fprint(&buf, ", ")
}
fmt.Fprint(&buf, s)
}
t.Logf("%d levels:\n%s\n", len(levels), buf.String())
fragmenterIter := NewIter(f.Cmp, allFragmented)
mergingIter := &MergingIter{}
mergingIter.Init(f.Cmp, VisibleTransform(base.InternalKeySeqNumMax), new(MergingBuffers), iters...)
// Position both so that it's okay to perform relative positioning
// operations immediately.
mergingIter.First()
fragmenterIter.First()
type opKind struct {
weight int
fn func() (str string, f *Span, m *Span)
}
ops := []opKind{
{weight: 2, fn: func() (string, *Span, *Span) {
return "First()", fragmenterIter.First(), mergingIter.First()
}},
{weight: 2, fn: func() (string, *Span, *Span) {
return "Last()", fragmenterIter.Last(), mergingIter.Last()
}},
{weight: 5, fn: func() (string, *Span, *Span) {
k := testkeys.Key(ks, rng.Int63n(ks.Count()))
return fmt.Sprintf("SeekGE(%q)", k),
fragmenterIter.SeekGE(k),
mergingIter.SeekGE(k)
}},
{weight: 5, fn: func() (string, *Span, *Span) {
k := testkeys.Key(ks, rng.Int63n(ks.Count()))
return fmt.Sprintf("SeekLT(%q)", k),
fragmenterIter.SeekLT(k),
mergingIter.SeekLT(k)
}},
{weight: 50, fn: func() (string, *Span, *Span) {
return "Next()", fragmenterIter.Next(), mergingIter.Next()
}},
{weight: 50, fn: func() (string, *Span, *Span) {
return "Prev()", fragmenterIter.Prev(), mergingIter.Prev()
}},
}
var totalWeight int
for _, op := range ops {
totalWeight += op.weight
}
var fragmenterBuf bytes.Buffer
var mergingBuf bytes.Buffer
opCount := rng.Intn(200) + 50
for i := 0; i < opCount; i++ {
p := rng.Intn(totalWeight)
opIndex := 0
for i, op := range ops {
if p < op.weight {
opIndex = i
break
}
p -= op.weight
}
opString, fs, ms := ops[opIndex].fn()
fragmenterBuf.Reset()
mergingBuf.Reset()
fmt.Fprint(&fragmenterBuf, fs)
fmt.Fprint(&mergingBuf, ms)
if fragmenterBuf.String() != mergingBuf.String() {
t.Fatalf("seed %d, op %d: %s = %s, fragmenter iterator returned %s",
seed, i, opString, mergingBuf.String(), fragmenterBuf.String())
}
t.Logf("op %d: %s = %s", i, opString, fragmenterBuf.String())
}
}