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

201 lines
5.7 KiB
Go

package replay
import (
"bytes"
"crypto/rand"
"fmt"
"io"
"strconv"
"strings"
"testing"
"time"
"unicode"
"github.com/cockroachdb/datadriven"
"github.com/cockroachdb/pebble"
"github.com/cockroachdb/pebble/internal/base"
"github.com/cockroachdb/pebble/internal/humanize"
"github.com/cockroachdb/pebble/vfs"
"github.com/stretchr/testify/require"
)
func TestWorkloadCollector(t *testing.T) {
const srcDir = `src`
const destDir = `dst`
datadriven.Walk(t, "testdata/collect", func(t *testing.T, path string) {
fs := vfs.NewMem()
require.NoError(t, fs.MkdirAll(srcDir, 0755))
require.NoError(t, fs.MkdirAll(destDir, 0755))
c := NewWorkloadCollector(srcDir)
o := &pebble.Options{FS: fs}
c.Attach(o)
var currentManifest vfs.File
var buf bytes.Buffer
defer func() {
if currentManifest != nil {
currentManifest.Close()
}
}()
datadriven.RunTest(t, path, func(t *testing.T, td *datadriven.TestData) string {
buf.Reset()
switch td.Cmd {
case "cmp-files":
if len(td.CmdArgs) != 2 {
return fmt.Sprintf("expected exactly 2 args, received %d", len(td.CmdArgs))
}
b1 := readFile(t, fs, td.CmdArgs[0].String())
b2 := readFile(t, fs, td.CmdArgs[1].String())
if !bytes.Equal(b1, b2) {
return fmt.Sprintf("files are unequal: %s (%s) and %s (%s)",
td.CmdArgs[0].String(), humanize.Bytes.Uint64(uint64(len(b1))),
td.CmdArgs[1].String(), humanize.Bytes.Uint64(uint64(len(b2))))
}
return "equal"
case "clean":
for _, path := range strings.Fields(td.Input) {
typ, _, ok := base.ParseFilename(fs, path)
require.True(t, ok)
require.NoError(t, o.Cleaner.Clean(fs, typ, path))
}
return ""
case "create-manifest":
if currentManifest != nil {
require.NoError(t, currentManifest.Close())
}
var fileNum uint64
var err error
td.ScanArgs(t, "filenum", &fileNum)
path := base.MakeFilepath(fs, srcDir, base.FileTypeManifest, base.DiskFileNum(fileNum))
currentManifest, err = fs.Create(path)
require.NoError(t, err)
_, err = currentManifest.Write(randData(100))
require.NoError(t, err)
c.onManifestCreated(pebble.ManifestCreateInfo{
Path: path,
FileNum: base.DiskFileNum(fileNum),
})
return ""
case "flush":
flushInfo := pebble.FlushInfo{
Done: true,
Input: 1,
Duration: 100 * time.Millisecond,
TotalDuration: 100 * time.Millisecond,
}
for _, line := range strings.Split(td.Input, "\n") {
if line == "" {
continue
}
parts := strings.FieldsFunc(line, func(r rune) bool { return unicode.IsSpace(r) || r == ':' })
tableInfo := pebble.TableInfo{Size: 10 << 10}
fileNum, err := strconv.ParseUint(parts[0], 10, 64)
require.NoError(t, err)
tableInfo.FileNum = base.FileNum(fileNum)
p := writeFile(t, fs, srcDir, base.FileTypeTable, tableInfo.FileNum.DiskFileNum(), randData(int(tableInfo.Size)))
fmt.Fprintf(&buf, "created %s\n", p)
flushInfo.Output = append(flushInfo.Output, tableInfo)
// Simulate a version edit applied to the current manifest.
_, err = currentManifest.Write(randData(25))
require.NoError(t, err)
}
flushInfo.InputBytes = 100 // Determinism
fmt.Fprint(&buf, flushInfo.String())
c.onFlushEnd(flushInfo)
return buf.String()
case "ingest":
ingestInfo := pebble.TableIngestInfo{}
for _, line := range strings.Split(td.Input, "\n") {
if line == "" {
continue
}
parts := strings.FieldsFunc(line, func(r rune) bool { return unicode.IsSpace(r) || r == ':' })
tableInfo := pebble.TableInfo{Size: 10 << 10}
fileNum, err := strconv.ParseUint(parts[0], 10, 64)
require.NoError(t, err)
tableInfo.FileNum = base.FileNum(fileNum)
p := writeFile(t, fs, srcDir, base.FileTypeTable, tableInfo.FileNum.DiskFileNum(), randData(int(tableInfo.Size)))
fmt.Fprintf(&buf, "created %s\n", p)
ingestInfo.Tables = append(ingestInfo.Tables, struct {
pebble.TableInfo
Level int
}{Level: 0, TableInfo: tableInfo})
// Simulate a version edit applied to the current manifest.
_, err = currentManifest.Write(randData(25))
require.NoError(t, err)
}
fmt.Fprint(&buf, ingestInfo.String())
c.onTableIngest(ingestInfo)
return buf.String()
case "ls":
return runListFiles(t, fs, td)
case "start":
c.Start(fs, destDir)
return ""
case "stat":
var buf bytes.Buffer
for _, arg := range td.CmdArgs {
fi, err := fs.Stat(arg.String())
if err != nil {
fmt.Fprintf(&buf, "%s: %s\n", arg.String(), err)
continue
}
fmt.Fprintf(&buf, "%s:\n", arg.String())
fmt.Fprintf(&buf, " size: %d\n", fi.Size())
}
return buf.String()
case "stop":
c.Stop()
return ""
case "wait":
// Wait until all pending sstables have been copied, then list
// the files in the destination directory.
c.mu.Lock()
for c.mu.tablesEnqueued != c.mu.tablesCopied {
c.mu.copyCond.Wait()
}
c.mu.Unlock()
listFiles(t, fs, &buf, destDir)
return buf.String()
default:
return fmt.Sprintf("unrecognized command %q", td.Cmd)
}
})
})
}
func randData(byteCount int) []byte {
b := make([]byte, byteCount)
rand.Read(b)
return b
}
func writeFile(
t *testing.T, fs vfs.FS, dir string, typ base.FileType, fileNum base.DiskFileNum, data []byte,
) string {
path := base.MakeFilepath(fs, dir, typ, fileNum)
f, err := fs.Create(path)
require.NoError(t, err)
_, err = f.Write(data)
require.NoError(t, err)
require.NoError(t, f.Close())
return path
}
func readFile(t *testing.T, fs vfs.FS, path string) []byte {
r, err := fs.Open(path)
require.NoError(t, err)
b, err := io.ReadAll(r)
require.NoError(t, err)
require.NoError(t, r.Close())
return b
}