ceremonyclient/pebble/sstable/suffix_rewriter_test.go

279 lines
9.1 KiB
Go
Raw Normal View History

2024-01-03 07:31:42 +00:00
package sstable
import (
"bytes"
"encoding/binary"
"fmt"
"math/rand"
"strconv"
"testing"
"github.com/cockroachdb/pebble/bloom"
"github.com/cockroachdb/pebble/internal/base"
"github.com/cockroachdb/pebble/objstorage"
"github.com/stretchr/testify/require"
)
func TestRewriteSuffixProps(t *testing.T) {
from, to := []byte("_212"), []byte("_646")
for format := TableFormatPebblev2; format <= TableFormatMax; format++ {
t.Run(format.String(), func(t *testing.T) {
wOpts := WriterOptions{
FilterPolicy: bloom.FilterPolicy(10),
Comparer: test4bSuffixComparer,
TablePropertyCollectors: []func() TablePropertyCollector{
intSuffixTablePropCollectorFn("ts3", 3), intSuffixTablePropCollectorFn("ts2", 2),
},
BlockPropertyCollectors: []func() BlockPropertyCollector{
keyCountCollectorFn("count"),
intSuffixIntervalCollectorFn("bp3", 3),
intSuffixIntervalCollectorFn("bp2", 2),
intSuffixIntervalCollectorFn("bp1", 1),
},
TableFormat: format,
}
if format >= TableFormatPebblev4 {
wOpts.IsStrictObsolete = true
}
const keyCount = 1e5
const rangeKeyCount = 100
// Setup our test SST.
sst := make4bSuffixTestSST(t, wOpts, []byte(from), keyCount, rangeKeyCount)
expectedProps := make(map[string]string)
expectedProps["ts2.min"] = "46"
expectedProps["ts2.max"] = "46"
expectedProps["ts3.min"] = "646"
expectedProps["ts3.max"] = "646"
// Also expect to see the aggregated block properties with their updated value
// at the correct (new) shortIDs. Seeing the rolled up value here is almost an
// end-to-end test since we only fed them each block during rewrite.
expectedProps["count"] = string(append([]byte{1}, strconv.Itoa(keyCount+rangeKeyCount)...))
expectedProps["bp2"] = string(interval{46, 47}.encode([]byte{2}))
expectedProps["bp3"] = string(interval{646, 647}.encode([]byte{0}))
// Swap the order of two of the props so they have new shortIDs, and remove
// one. rwOpts inherits the IsStrictObsolete value from wOpts.
rwOpts := wOpts
if rand.Intn(2) != 0 {
rwOpts.TableFormat = TableFormatPebblev2
rwOpts.IsStrictObsolete = false
t.Log("table format set to TableFormatPebblev2")
}
fmt.Printf("from format %s, to format %s\n", format.String(), rwOpts.TableFormat.String())
rwOpts.BlockPropertyCollectors = rwOpts.BlockPropertyCollectors[:3]
rwOpts.BlockPropertyCollectors[0], rwOpts.BlockPropertyCollectors[1] = rwOpts.BlockPropertyCollectors[1], rwOpts.BlockPropertyCollectors[0]
// Rewrite the SST using updated options and check the returned props.
readerOpts := ReaderOptions{
Comparer: test4bSuffixComparer,
Filters: map[string]base.FilterPolicy{wOpts.FilterPolicy.Name(): wOpts.FilterPolicy},
}
r, err := NewMemReader(sst, readerOpts)
require.NoError(t, err)
defer r.Close()
var sstBytes [2][]byte
adjustPropsForEffectiveFormat := func(effectiveFormat TableFormat) {
if effectiveFormat == TableFormatPebblev4 {
expectedProps["obsolete-key"] = string([]byte{3})
} else {
delete(expectedProps, "obsolete-key")
}
}
for i, byBlocks := range []bool{false, true} {
t.Run(fmt.Sprintf("byBlocks=%v", byBlocks), func(t *testing.T) {
rewrittenSST := &memFile{}
if byBlocks {
_, rewriteFormat, err := rewriteKeySuffixesInBlocks(
r, rewrittenSST, rwOpts, from, to, 8)
// rewriteFormat is equal to the original format, since
// rwOpts.TableFormat is ignored.
require.Equal(t, wOpts.TableFormat, rewriteFormat)
require.NoError(t, err)
adjustPropsForEffectiveFormat(rewriteFormat)
} else {
_, err := RewriteKeySuffixesViaWriter(r, rewrittenSST, rwOpts, from, to)
require.NoError(t, err)
adjustPropsForEffectiveFormat(rwOpts.TableFormat)
}
sstBytes[i] = rewrittenSST.Data()
// Check that a reader on the rewritten STT has the expected props.
rRewritten, err := NewMemReader(rewrittenSST.Data(), readerOpts)
require.NoError(t, err)
defer rRewritten.Close()
require.Equal(t, expectedProps, rRewritten.Properties.UserProperties)
require.False(t, rRewritten.Properties.IsStrictObsolete)
// Compare the block level props from the data blocks in the layout,
// only if we did not do a rewrite from one format to another. If the
// format changes, the block boundaries change slightly.
if !byBlocks && wOpts.TableFormat != rwOpts.TableFormat {
return
}
layout, err := r.Layout()
require.NoError(t, err)
newLayout, err := rRewritten.Layout()
require.NoError(t, err)
ival := interval{}
for i := range layout.Data {
oldProps := make([][]byte, len(wOpts.BlockPropertyCollectors))
oldDecoder := blockPropertiesDecoder{layout.Data[i].Props}
for !oldDecoder.done() {
id, val, err := oldDecoder.next()
require.NoError(t, err)
oldProps[id] = val
}
newProps := make([][]byte, len(rwOpts.BlockPropertyCollectors))
newDecoder := blockPropertiesDecoder{newLayout.Data[i].Props}
for !newDecoder.done() {
id, val, err := newDecoder.next()
require.NoError(t, err)
if int(id) < len(newProps) {
newProps[id] = val
}
}
require.Equal(t, oldProps[0], newProps[1])
ival.decode(newProps[0])
require.Equal(t, interval{646, 647}, ival)
ival.decode(newProps[2])
require.Equal(t, interval{46, 47}, ival)
}
})
}
if wOpts.TableFormat == rwOpts.TableFormat {
// Both methods of rewriting should produce the same result.
require.Equal(t, sstBytes[0], sstBytes[1])
}
})
}
}
// memFile is a file-like struct that buffers all data written to it in memory.
// Implements the objstorage.Writable interface.
type memFile struct {
buf bytes.Buffer
}
var _ objstorage.Writable = (*memFile)(nil)
// Finish is part of the objstorage.Writable interface.
func (*memFile) Finish() error {
return nil
}
// Abort is part of the objstorage.Writable interface.
func (*memFile) Abort() {}
// Write is part of the objstorage.Writable interface.
func (f *memFile) Write(p []byte) error {
_, err := f.buf.Write(p)
return err
}
// Data returns the in-memory buffer behind this MemFile.
func (f *memFile) Data() []byte {
return f.buf.Bytes()
}
func make4bSuffixTestSST(
t testing.TB, writerOpts WriterOptions, suffix []byte, keys int, rangeKeys int,
) []byte {
key := make([]byte, 28)
endKey := make([]byte, 24)
copy(key[24:], suffix)
f := &memFile{}
w := NewWriter(f, writerOpts)
for i := 0; i < keys; i++ {
binary.BigEndian.PutUint64(key[:8], 123) // 16-byte shared prefix
binary.BigEndian.PutUint64(key[8:16], 456)
binary.BigEndian.PutUint64(key[16:], uint64(i))
err := w.AddWithForceObsolete(
base.MakeInternalKey(key, 0, InternalKeyKindSet), key, false)
if err != nil {
t.Fatal(err)
}
}
for i := 0; i < rangeKeys; i++ {
binary.BigEndian.PutUint64(key[:8], 123) // 16-byte shared prefix
binary.BigEndian.PutUint64(key[8:16], 456)
binary.BigEndian.PutUint64(key[16:], uint64(i))
binary.BigEndian.PutUint64(endKey[:8], 123) // 16-byte shared prefix
binary.BigEndian.PutUint64(endKey[8:16], 456)
binary.BigEndian.PutUint64(endKey[16:], uint64(i+1))
if err := w.RangeKeySet(key[:24], endKey[:24], suffix, key); err != nil {
t.Fatal(err)
}
}
if err := w.Close(); err != nil {
t.Fatal(err)
}
return f.buf.Bytes()
}
func BenchmarkRewriteSST(b *testing.B) {
from, to := []byte("_123"), []byte("_456")
writerOpts := WriterOptions{
FilterPolicy: bloom.FilterPolicy(10),
Comparer: test4bSuffixComparer,
TableFormat: TableFormatPebblev2,
}
sizes := []int{100, 10000, 1e6}
compressions := []Compression{NoCompression, SnappyCompression}
files := make([][]*Reader, len(compressions))
for comp := range compressions {
files[comp] = make([]*Reader, len(sizes))
for size := range sizes {
writerOpts.Compression = compressions[comp]
sst := make4bSuffixTestSST(b, writerOpts, from, sizes[size], 0 /* rangeKeys */)
r, err := NewMemReader(sst, ReaderOptions{
Comparer: test4bSuffixComparer,
Filters: map[string]base.FilterPolicy{writerOpts.FilterPolicy.Name(): writerOpts.FilterPolicy},
})
if err != nil {
b.Fatal(err)
}
files[comp][size] = r
}
}
b.ResetTimer()
for comp := range compressions {
b.Run(compressions[comp].String(), func(b *testing.B) {
for sz := range sizes {
r := files[comp][sz]
b.Run(fmt.Sprintf("keys=%d", sizes[sz]), func(b *testing.B) {
b.Run("ReaderWriterLoop", func(b *testing.B) {
b.SetBytes(r.readable.Size())
for i := 0; i < b.N; i++ {
if _, err := RewriteKeySuffixesViaWriter(r, &discardFile{}, writerOpts, from, to); err != nil {
b.Fatal(err)
}
}
})
for _, concurrency := range []int{1, 2, 4, 8, 16} {
b.Run(fmt.Sprintf("RewriteKeySuffixes,concurrency=%d", concurrency), func(b *testing.B) {
b.SetBytes(r.readable.Size())
for i := 0; i < b.N; i++ {
if _, _, err := rewriteKeySuffixesInBlocks(r, &discardFile{}, writerOpts, []byte("_123"), []byte("_456"), concurrency); err != nil {
b.Fatal(err)
}
}
})
}
})
}
})
}
}