cmd/pebble: allow setting custom Pebble Options through a CLI flag

cmd/pebble/replay.go

@@ -5,6 +5,7 @@
 package main
 
 import (
+ "bytes"
  "context"
  "flag"
  "fmt"
@@ -14,6 +15,7 @@ import (
  "sort"
  "strconv"
  "strings"
+ "unicode"
 
  "github.com/cockroachdb/errors"
  "github.com/cockroachdb/pebble"
@@ -43,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().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(
  &c.runDir, "run-dir", c.runDir, "the directory to use for the replay data directory; defaults to a random dir in pwd")
  cmd.Flags().BoolVar(
@@ -59,6 +63,7 @@ type replayConfig struct {
  count int
  streamLogs bool
  ignoreCheckpoint bool
+ optionsString string
 
  cleanUpFuncs []func() error
 }
@@ -193,6 +198,9 @@ func (c *replayConfig) initOptions(r *replay.Runner) error {
  return err
  }
  }
+ if err := parseCustomOptions(c.optionsString, r.Opts); err != nil {
+ return err
+ }
  // TODO(jackson): If r.Opts.Comparer == nil, peek at the workload's
  // manifests and pull the comparer out of them.
  //
@@ -207,6 +215,49 @@ func (c *replayConfig) initOptions(r *replay.Runner) error {
  return nil
 }
 
+// parseCustomOptions parses Pebble Options passed through a CLI flag.
+// Ordinarily Pebble Options are specified through an INI file with newlines
+// delimiting fields. That doesn't translate well to a CLI interface, so this
+// function accepts fields are that delimited by any whitespace. This is the
+// same format that CockroachDB accepts Pebble Options through the --store flag,
+// and this code is copied from there.
+func parseCustomOptions(optsStr string, opts *pebble.Options) error {
+ if optsStr == "" {
+ return nil
+ }
+ // Pebble options are supplied in the Pebble OPTIONS ini-like
+ // format, but allowing any whitespace to delimit lines. Convert
+ // the options to a newline-delimited format. This isn't a trivial
+ // character replacement because whitespace may appear within a
+ // stanza, eg ["Level 0"].
+ value := strings.TrimSpace(optsStr)
+ var buf bytes.Buffer
+ for len(value) > 0 {
+ i := strings.IndexFunc(value, func(r rune) bool {
+ return r == '[' || unicode.IsSpace(r)
+ })
+ switch {
+ case i == -1:
+ buf.WriteString(value)
+ value = value[len(value):]
+ case value[i] == '[':
+ // If there's whitespace within [ ], we write it verbatim.
+ j := i + strings.IndexRune(value[i:], ']')
+ buf.WriteString(value[:j+1])
+ value = value[j+1:]
+ case unicode.IsSpace(rune(value[i])):
+ // NB: This doesn't handle multibyte whitespace.
+ buf.WriteString(value[:i])
+ buf.WriteRune('\n')
+ value = strings.TrimSpace(value[i+1:])
+ }
+ }
+ return opts.Parse(buf.String(), &pebble.ParseHooks{
+ NewComparer: makeComparer,
+ NewMerger: makeMerger,
+ })
+}
+
 func (c *replayConfig) cleanUp() error {
  for _, f := range c.cleanUpFuncs {
  if err := f(); err != nil {

cmd/pebble/replay_test.go

@@ -0,0 +1,58 @@
+// Copyright 2023 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 main
+
+import (
+ "testing"
+
+ "github.com/cockroachdb/pebble"
+ "github.com/stretchr/testify/require"
+)
+
+func TestParseOptionsStr(t *testing.T) {
+ type testCase struct {
+ optionsStr string
+ options *pebble.Options
+ }
+
+ testCases := []testCase{
+ {
+ optionsStr: `[Options] max_concurrent_compactions=9`,
+ options: &pebble.Options{MaxConcurrentCompactions: func() int { return 9 }},
+ },
+ {
+ optionsStr: `[Options] bytes_per_sync=90000`,
+ options: &pebble.Options{BytesPerSync: 90000},
+ },
+ {
+ optionsStr: `[Options] [Level "0"] target_file_size=222`,
+ options: &pebble.Options{Levels: []pebble.LevelOptions{
+ {TargetFileSize: 222},
+ }},
+ },
+ {
+ optionsStr: `[Options] lbase_max_bytes=10 max_open_files=20 [Level "0"] target_file_size=30 [Level "1"] index_block_size=40`,
+ options: &pebble.Options{
+ LBaseMaxBytes: 10,
+ MaxOpenFiles: 20,
+ Levels: []pebble.LevelOptions{
+ {TargetFileSize: 30},
+ {IndexBlockSize: 40},
+ },
+ },
+ },
+ }
+
+ for _, tc := range testCases {
+ o := new(pebble.Options)
+ require.NoError(t, parseCustomOptions(tc.optionsStr, o))
+ o.EnsureDefaults()
+ got := o.String()
+
+ tc.options.EnsureDefaults()
+ want := tc.options.String()
+ require.Equal(t, want, got)
+ }
+}