replay: add option to apply subset of a workload

Add a `--max-writes` flag to the cmd/pebble replay command that limits the
amount of a workload to replay by the configured MB. This allows us to collect
large workloads and apply a smaller subset at replay time if the full workload
is unnecessary.

cmd/pebble/replay.go

@@ -45,6 +45,8 @@ func initReplayCmd() *cobra.Command {
  &c.name, "name", "", "the name of the workload being replayed")
  cmd.Flags().VarPF(
  &c.pacer, "pacer", "p", "the pacer to use: unpaced, reference-ramp, or fixed-ramp=N")
+ cmd.Flags().Uint64Var(
+ &c.maxWritesMB, "max-writes", 0, "the maximum volume of writes (MB) to apply, with 0 denoting unlimited")
  cmd.Flags().StringVar(
  &c.optionsString, "options", "", "Pebble options to override, in the OPTIONS ini format but with any whitespace as field delimiters instead of newlines")
  cmd.Flags().StringVar(
@@ -61,6 +63,7 @@ type replayConfig struct {
  pacer pacerFlag
  runDir string
  count int
+ maxWritesMB uint64
  streamLogs bool
  ignoreCheckpoint bool
  optionsString string
@@ -93,6 +96,9 @@ func (c *replayConfig) runOnce(stdout io.Writer, workloadPath string) error {
  Pacer: c.pacer,
  Opts: &pebble.Options{},
  }
+ if c.maxWritesMB > 0 {
+ r.MaxWriteBytes = c.maxWritesMB * (1 << 20)
+ }
  if err := c.initRunDir(r); err != nil {
  return err
  }

replay/replay.go

@@ -115,6 +115,7 @@ type Metrics struct {
  BytesWeightedByLevel uint64
  }
  TotalWriteAmp float64
+ WriteBytes uint64
  WriteStalls uint64
  WriteStallsDuration time.Duration
 }
@@ -180,11 +181,12 @@ func (m *Metrics) BenchmarkString(name string) string {
 // Runner runs a captured workload against a test database, collecting
 // metrics on performance.
 type Runner struct {
- RunDir string
- WorkloadFS vfs.FS
- WorkloadPath string
- Pacer Pacer
- Opts *pebble.Options
+ RunDir string
+ WorkloadFS vfs.FS
+ WorkloadPath string
+ Pacer Pacer
+ Opts *pebble.Options
+ MaxWriteBytes uint64
 
  // Internal state.
 
@@ -202,6 +204,7 @@ type Runner struct {
  stepsApplied chan workloadStep
 
  metrics struct {
+ writeBytes uint64
  writeStalls uint64
  writeStallsDurationNano uint64
  }
@@ -420,6 +423,7 @@ func (r *Runner) Wait() (Metrics, error) {
  m := Metrics{
  Final: pm,
  TotalWriteAmp: total.WriteAmp(),
+ WriteBytes: r.metrics.writeBytes,
  WriteStalls: r.metrics.writeStalls,
  WriteStallsDuration: time.Duration(r.metrics.writeStallsDurationNano),
  }
@@ -556,6 +560,7 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
 
  idx := r.workload.manifestIdx
 
+ var cumulativeWriteBytes uint64
  var flushBufs flushBuffers
  var v *manifest.Version
  var previousVersion *manifest.Version
@@ -576,6 +581,10 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
  }
 
  for ; idx < len(r.workload.manifests); idx++ {
+ if r.MaxWriteBytes != 0 && cumulativeWriteBytes > r.MaxWriteBytes {
+ break
+ }
+
  err := func() error {
  manifestName := r.workload.manifests[idx]
  f, err := r.WorkloadFS.Open(r.WorkloadFS.PathJoin(r.WorkloadPath, manifestName))
@@ -688,6 +697,7 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
  if err := loadFlushedSSTableKeys(s.flushBatch, r.WorkloadFS, r.WorkloadPath, newFiles, r.readerOpts, &flushBufs); err != nil {
  return errors.Wrapf(err, "flush in %q at offset %d", manifestName, rr.Offset())
  }
+ cumulativeWriteBytes += uint64(s.flushBatch.Len())
  case ingestStepKind:
  // Copy the ingested sstables into a staging area within the
  // run dir. This is necessary for two reasons:
@@ -705,6 +715,11 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
  if err := vfs.CopyAcrossFS(r.WorkloadFS, src, r.Opts.FS, dst); err != nil {
  return errors.Wrapf(err, "ingest in %q at offset %d", manifestName, rr.Offset())
  }
+ finfo, err := r.Opts.FS.Stat(dst)
+ if err != nil {
+ return errors.Wrapf(err, "stating %q", dst)
+ }
+ cumulativeWriteBytes += uint64(finfo.Size())
  s.tablesToIngest = append(s.tablesToIngest, dst)
  }
  case compactionStepKind:
@@ -716,13 +731,18 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
  return ctx.Err()
  case r.steps <- s:
  }
+
+ if r.MaxWriteBytes != 0 && cumulativeWriteBytes > r.MaxWriteBytes {
+ break
+ }
  }
  return nil
  }()
  if err != nil {
  return err
  }
  }
+ atomic.StoreUint64(&r.metrics.writeBytes, cumulativeWriteBytes)
  return nil
 }
 

replay/replay_test.go

@@ -21,6 +21,7 @@ import (
  "github.com/cockroachdb/pebble"
  "github.com/cockroachdb/pebble/internal/base"
  "github.com/cockroachdb/pebble/internal/datatest"
+ "github.com/cockroachdb/pebble/internal/humanize"
  "github.com/cockroachdb/pebble/internal/invariants"
  "github.com/cockroachdb/pebble/internal/testkeys"
  "github.com/cockroachdb/pebble/rangekey"
@@ -131,10 +132,11 @@ func runReplayTest(t *testing.T, path string) {
  ct.WaitForInflightCompactionsToEqual(target)
  return ""
  case "wait":
- if _, err := r.Wait(); err != nil {
+ m, err := r.Wait()
+ if err != nil {
  return err.Error()
  }
- return ""
+ return fmt.Sprintf("replayed %s in writes", humanize.Uint64(m.WriteBytes))
  case "close":
  if err := r.Close(); err != nil {
  return err.Error()

replay/testdata/replay

@@ -85,6 +85,7 @@ replay simple unpaced
 
 wait
 ----
+replayed 42 B in writes
 
 # NB: The file sizes are non-deterministic after replay (because compactions are
 # nondeterministic). We don't `tree` here as a result.

replay/testdata/replay_paced

@@ -39,6 +39,7 @@ wait-for-compactions
 
 wait
 ----
+replayed 42 B in writes
 
 scan-keys
 ----