From cb8cc0a262a7595f8e03734e1ba8e20cfdc8b731 Mon Sep 17 00:00:00 2001 From: Mile Druzijanic Date: Tue, 28 Jul 2026 13:52:58 +0200 Subject: [PATCH] feat: add custom per command flags --- README.md | 39 +++++- bench/commands/spec.go | 122 +++++++++++++++++ bench/commands/spec_test.go | 137 ++++++++++++++----- benchmarks/HGET/dfly_bench/HGET_reproduce.md | 36 +++++ benchmarks/HSET/dfly_bench/HSET_reproduce.md | 38 +++++ cmd/run_commands.go | 77 ++++++++++- cmd/run_commands_dfly.go | 45 +++++- cmd/run_commands_test.go | 57 ++++++++ 8 files changed, 509 insertions(+), 42 deletions(-) create mode 100644 benchmarks/HGET/dfly_bench/HGET_reproduce.md create mode 100644 benchmarks/HSET/dfly_bench/HSET_reproduce.md create mode 100644 cmd/run_commands_test.go diff --git a/README.md b/README.md index 757e564..ce1e472 100644 --- a/README.md +++ b/README.md @@ -580,13 +580,14 @@ the redirected file is a single well-formed payload. ### `run-commands` Benchmark one Redis command against an existing setup. Requires `--id` -and `--command`; writes to `~/.dfbench/runs//commands//`. +and either `--command` (a built-in command) or `--command-template` (a +custom one); writes to `~/.dfbench/runs//commands//`. Run it once per command you want to profile. | Flag | Default | Notes | | ---- | ------- | ----- | | `--id` | (required) | Setup previously created by `dfbench setup`. | -| `--command` | (required) | Single command name, e.g. `SET`, `GET`, `LRANGE`. See the supported list above. | +| `--command` | (one of --command / --command-template) | Single built-in command name, e.g. `SET`, `GET`, `LRANGE`. See the supported list above. Mutually exclusive with `--command-template`. | | `--engines` | `dragonfly,redis,valkey` | Subset of the setup's installed engines. | | `--memtier-threads` | `32` | Total concurrent connections = threads * clients. | | `--memtier-clients` | `5` | Clients per thread. | @@ -615,6 +616,37 @@ dfbench run-commands --id percmd \ --key-maximum 1000000 --test-time 15 --trials 1 ``` +#### Custom commands + +To benchmark a command that isn't in the built-in registry, replace +`--command ` with `--command-template` (the two are mutually +exclusive). The template uses memtier's arbitrary-command form with the +same `__key__` / `__data__` placeholders as the built-in specs +(`__data__` is filled from `--memtier-data-size`). + +| Flag | Default | Notes | +| ---- | ------- | ----- | +| `--command-template` | `""` | Custom command, e.g. `"SETEX __key__ 3600 __data__"`. Mutually exclusive with `--command`. | +| `--command-name` | first word of the template | Label used in output files and `results.json`. | +| `--memtier-command-key-pattern` | `R` | memtier `--command-key-pattern` for the measured pass: `R` (random), `S` (sequential), `P` (parallel), `G` (gaussian). memtier-specific, so it carries the `--memtier-` prefix; the dfly variant uses `--dfly-bench-key-dist` instead. | +| `--command-mutates` | `false` | Mark the command as mutating: flush + re-preload before every trial. | +| `--command-drains` | `false` | Mark the command as draining preloaded data (like `DEL`, `LPOP`); implies `--command-mutates` and flags results approximate. | +| `--preload-template` | `""` | Optional preload command, e.g. `"SET __key__ __data__"`; must contain `__key__`. | +| `--preload-items` | `1` | Items written per key by the preload (list/set/zset cardinality). | +| `--command-key-maximum` | `0` (use `--key-maximum`) | Per-command key range override. Required when `--preload-items > 1` so a multi-item preload doesn't run against the full 100M range. | + +Example - a mutating custom command (`SETEX`) and a read command +(`GETRANGE`) preloaded with plain strings: + +```bash +dfbench run-commands --id percmd --engines dragonfly --dragonfly_num_shards -1 \ + --command-template "SETEX __key__ 3600 __data__" --command-name SETEX --command-mutates + +dfbench run-commands --id percmd --engines dragonfly --dragonfly_num_shards -1 \ + --command-template "GETRANGE __key__ 0 50" --command-name GETRANGE \ + --preload-template "SET __key__ __data__" +``` + Budget note: one run is roughly (engines) x (trials x test-time + warmup + preload + pauses); the default 3 engines x 3 x 300s lands around 50-60 minutes per command. Preloading a 100M keyspace for read @@ -667,7 +699,8 @@ server; they are unrelated. | Flag | Default | Notes | | ---- | ------- | ----- | | `--id` | (required) | Setup previously created by `dfbench setup`. | -| `--command` | (required) | Single command name, e.g. `SET`, `GET`, `LRANGE`. | +| `--command` | (one of --command / --command-template) | Single built-in command name, e.g. `SET`, `GET`, `LRANGE`. Mutually exclusive with `--command-template`. | +| `--command-template` (+ `--command-name`, `--command-mutates`, `--command-drains`, `--preload-template`, `--preload-items`, `--command-key-maximum`) | (see [Custom commands](#custom-commands)) | Define a command not in the registry. Same workload flags as `run-commands`. There is no key-pattern flag here (that is memtier-specific); the measured distribution comes from `--dfly-bench-key-dist`. | | `--engines` | `dragonfly` | Subset of the setup's installed engines (`dfly_bench` speaks RESP, so Redis/Valkey work too). | | `--dfly-bench-threads` | `32` | `dfly_bench --proactor_threads` (analogous to memtier threads). Total connections = threads * conns. | | `--dfly-bench-conns` | `5` | `dfly_bench -c`, connections per thread (analogous to memtier clients). | diff --git a/bench/commands/spec.go b/bench/commands/spec.go index b70a439..9572a35 100644 --- a/bench/commands/spec.go +++ b/bench/commands/spec.go @@ -153,3 +153,125 @@ func Lookup(raw string) (CommandSpec, error) { } return spec, nil } + +// validKeyPatterns is the set of memtier --command-key-pattern values a spec +// may use: R (uniform random), S (sequential), P (parallel), G (gaussian). +var validKeyPatterns = map[string]bool{"R": true, "S": true, "P": true, "G": true} + +// Validate checks the invariants every CommandSpec must satisfy, whether it +// comes from the built-in registry or from user-supplied CLI flags. It is +// the single source of truth for what makes a spec runnable. +func (s CommandSpec) Validate() error { + if s.Name == "" { + return fmt.Errorf("command name is empty") + } + if s.Name != strings.ToUpper(s.Name) { + return fmt.Errorf("command name %q must be upper-case", s.Name) + } + if s.Template == "" { + return fmt.Errorf("%s: empty template", s.Name) + } + if !validKeyPatterns[s.KeyPattern] { + return fmt.Errorf("%s: invalid key pattern %q (want one of R, S, P, G)", s.Name, s.KeyPattern) + } + if s.Drains && !s.Mutates { + return fmt.Errorf("%s: drains the keyspace but is not marked mutating", s.Name) + } + if s.Drains && s.Preload == nil { + return fmt.Errorf("%s: drains the keyspace but has no preload", s.Name) + } + if s.Preload != nil { + if s.Preload.Template == "" { + return fmt.Errorf("%s: preload with empty template", s.Name) + } + if !strings.Contains(s.Preload.Template, "__key__") { + return fmt.Errorf("%s: preload template %q has no __key__ placeholder", s.Name, s.Preload.Template) + } + if s.Preload.Items() < 1 { + return fmt.Errorf("%s: preload Items() = %d, want >= 1", s.Name, s.Preload.Items()) + } + // Multi-item preloads must bound the key range or the preload cost + // explodes (keyMax * items ops against the default 100M range). + if s.Preload.Items() > 1 && s.KeyMaximum == 0 { + return fmt.Errorf("%s: preload writes %d items/key but has no key-maximum override "+ + "(set --command-key-maximum)", s.Name, s.Preload.Items()) + } + } + return nil +} + +// CustomSpecInput carries the raw, CLI-provided fields for a user-defined +// command spec. BuildCustom fills in defaults and validates the result. +type CustomSpecInput struct { + Name string + Template string + KeyPattern string + Mutates bool + Drains bool + PreloadTemplate string + PreloadItems int + KeyMaximum int +} + +// BuildCustom assembles a CommandSpec from CLI-provided fields, applying +// defaults (name derived from the template's first word, key pattern R, +// Drains implying Mutates) and validating the result. It lets run-commands +// benchmark a command that isn't in the built-in registry. +func BuildCustom(in CustomSpecInput) (CommandSpec, error) { + template := strings.TrimSpace(in.Template) + if template == "" { + return CommandSpec{}, fmt.Errorf("custom command needs a non-empty --command-template") + } + + name := sanitizeName(in.Name) + if name == "" { + name = sanitizeName(firstToken(template)) + } + if name == "" { + return CommandSpec{}, fmt.Errorf("could not derive a command name from template %q; pass --command-name", template) + } + + keyPattern := strings.ToUpper(strings.TrimSpace(in.KeyPattern)) + if keyPattern == "" { + keyPattern = "R" + } + + spec := CommandSpec{ + Name: name, + Template: template, + KeyPattern: keyPattern, + Mutates: in.Mutates || in.Drains, + Drains: in.Drains, + KeyMaximum: in.KeyMaximum, + } + if pt := strings.TrimSpace(in.PreloadTemplate); pt != "" { + spec.Preload = &PreloadSpec{Template: pt, ItemsPerKey: in.PreloadItems} + } + + if err := spec.Validate(); err != nil { + return CommandSpec{}, err + } + return spec, nil +} + +// firstToken returns the first whitespace-delimited word of s. +func firstToken(s string) string { + fields := strings.Fields(s) + if len(fields) == 0 { + return "" + } + return fields[0] +} + +// sanitizeName upper-cases the command name and strips it to characters safe +// for result labels and file paths (labels look like "--t1"). +func sanitizeName(raw string) string { + var b strings.Builder + for _, r := range strings.ToUpper(strings.TrimSpace(raw)) { + switch { + case r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '_', r == '-': + b.WriteRune(r) + } + } + return b.String() +} diff --git a/bench/commands/spec_test.go b/bench/commands/spec_test.go index 074e7e2..18b9d6f 100644 --- a/bench/commands/spec_test.go +++ b/bench/commands/spec_test.go @@ -6,51 +6,122 @@ import ( ) func TestSupported_Valid(t *testing.T) { - validPatterns := map[string]bool{"R": true, "S": true, "P": true, "G": true} seen := make(map[string]bool, len(Supported)) for _, s := range Supported { - if s.Name == "" { - t.Fatal("spec with empty name") - } - if s.Name != strings.ToUpper(s.Name) { - t.Errorf("%s: name must be upper-case", s.Name) + // Validate is the shared source of truth for the core invariants. + if err := s.Validate(); err != nil { + t.Errorf("%s: %v", s.Name, err) } + // Registry-only expectations Validate does not enforce. if seen[s.Name] { t.Errorf("%s: duplicate spec", s.Name) } seen[s.Name] = true - - if s.Template == "" { - t.Errorf("%s: empty template", s.Name) - } if !strings.HasPrefix(s.Template, s.Name+" ") && s.Template != s.Name { t.Errorf("%s: template %q does not start with the command name", s.Name, s.Template) } - if !validPatterns[s.KeyPattern] { - t.Errorf("%s: invalid key pattern %q", s.Name, s.KeyPattern) - } - if s.Drains && !s.Mutates { - t.Errorf("%s: drains the keyspace but is not marked mutating", s.Name) - } - if s.Drains && s.Preload == nil { - t.Errorf("%s: drains the keyspace but has no preload", s.Name) + } +} + +func TestValidate_Rejections(t *testing.T) { + cases := []struct { + name string + spec CommandSpec + }{ + {"empty name", CommandSpec{Template: "SET __key__ __data__", KeyPattern: "R"}}, + {"lower-case name", CommandSpec{Name: "set", Template: "SET __key__", KeyPattern: "R"}}, + {"empty template", CommandSpec{Name: "SET", KeyPattern: "R"}}, + {"bad key pattern", CommandSpec{Name: "SET", Template: "SET __key__", KeyPattern: "X"}}, + {"drains without mutates", CommandSpec{Name: "DEL", Template: "DEL __key__", KeyPattern: "R", Drains: true, + Preload: &PreloadSpec{Template: "SET __key__ __data__"}}}, + {"drains without preload", CommandSpec{Name: "DEL", Template: "DEL __key__", KeyPattern: "R", Mutates: true, Drains: true}}, + {"preload without __key__", CommandSpec{Name: "GET", Template: "GET __key__", KeyPattern: "R", + Preload: &PreloadSpec{Template: "SET foo __data__"}}}, + {"multi-item preload without key-maximum", CommandSpec{Name: "LRANGE", Template: "LRANGE __key__ 0 99", KeyPattern: "R", + Preload: &PreloadSpec{Template: "RPUSH __key__ __data__", ItemsPerKey: 100}}}, + } + for _, c := range cases { + if err := c.spec.Validate(); err == nil { + t.Errorf("%s: expected validation error, got nil", c.name) } - if s.Preload != nil { - if s.Preload.Template == "" { - t.Errorf("%s: preload with empty template", s.Name) - } - if !strings.Contains(s.Preload.Template, "__key__") { - t.Errorf("%s: preload template %q has no __key__ placeholder", s.Name, s.Preload.Template) - } - if s.Preload.Items() < 1 { - t.Errorf("%s: preload Items() = %d, want >= 1", s.Name, s.Preload.Items()) - } - // Multi-item preloads must bound the key range or preload - // cost explodes (keyMax * items ops). - if s.Preload.Items() > 1 && s.KeyMaximum == 0 { - t.Errorf("%s: preload writes %d items/key but has no KeyMaximum override", s.Name, s.Preload.Items()) - } + } +} + +func TestBuildCustom_Defaults(t *testing.T) { + spec, err := BuildCustom(CustomSpecInput{Template: "SETEX __key__ 3600 __data__"}) + if err != nil { + t.Fatalf("BuildCustom: %v", err) + } + if spec.Name != "SETEX" { + t.Errorf("Name = %q, want SETEX (derived from template)", spec.Name) + } + if spec.KeyPattern != "R" { + t.Errorf("KeyPattern = %q, want default R", spec.KeyPattern) + } + if spec.Preload != nil { + t.Errorf("Preload = %+v, want nil when no preload template given", spec.Preload) + } +} + +func TestBuildCustom_NameSanitizedAndUppercased(t *testing.T) { + spec, err := BuildCustom(CustomSpecInput{Name: "my cmd!", Template: "GET __key__"}) + if err != nil { + t.Fatalf("BuildCustom: %v", err) + } + if spec.Name != "MYCMD" { + t.Errorf("Name = %q, want MYCMD (upper-cased, stripped of unsafe chars)", spec.Name) + } +} + +func TestBuildCustom_DrainsImpliesMutates(t *testing.T) { + spec, err := BuildCustom(CustomSpecInput{ + Template: "DEL __key__", + Drains: true, + PreloadTemplate: "SET __key__ __data__", + }) + if err != nil { + t.Fatalf("BuildCustom: %v", err) + } + if !spec.Mutates { + t.Error("Drains should imply Mutates") + } +} + +func TestBuildCustom_Preload(t *testing.T) { + spec, err := BuildCustom(CustomSpecInput{ + Template: "LRANGE __key__ 0 99", + PreloadTemplate: "RPUSH __key__ __data__", + PreloadItems: 100, + KeyMaximum: 1_000_000, + }) + if err != nil { + t.Fatalf("BuildCustom: %v", err) + } + if spec.Preload == nil || spec.Preload.Items() != 100 { + t.Errorf("Preload = %+v, want template with 100 items/key", spec.Preload) + } + if spec.KeyMaximum != 1_000_000 { + t.Errorf("KeyMaximum = %d, want 1000000", spec.KeyMaximum) + } +} + +func TestBuildCustom_Rejections(t *testing.T) { + cases := []struct { + name string + in CustomSpecInput + }{ + {"empty template", CustomSpecInput{}}, + {"blank template", CustomSpecInput{Template: " "}}, + {"bad key pattern", CustomSpecInput{Template: "GET __key__", KeyPattern: "Z"}}, + {"multi-item preload without key-maximum", CustomSpecInput{ + Template: "LRANGE __key__ 0 99", PreloadTemplate: "RPUSH __key__ __data__", PreloadItems: 100}}, + {"preload without __key__", CustomSpecInput{ + Template: "GET __key__", PreloadTemplate: "SET foo __data__"}}, + } + for _, c := range cases { + if _, err := BuildCustom(c.in); err == nil { + t.Errorf("%s: expected error, got nil", c.name) } } } diff --git a/benchmarks/HGET/dfly_bench/HGET_reproduce.md b/benchmarks/HGET/dfly_bench/HGET_reproduce.md new file mode 100644 index 0000000..07e7a2a --- /dev/null +++ b/benchmarks/HGET/dfly_bench/HGET_reproduce.md @@ -0,0 +1,36 @@ +### Stateful setup: + +``` +dfbench setup --id maxqps2 --tune-network --ubuntu-version 24.04 \ + --server-instance m7g.8xlarge --server-arch arm64 \ + --client-instance c6gn.8xlarge --client-arch arm64 \ + --engines dragonfly,redis,valkey --availability-zone us-east-1c +``` + +### Test run: + +``` +dfbench run-commands-dfly \ + --id maxqps2 --engines dragonfly,redis,valkey \ + \ + --command-template "HGET __key__ field" --command-name HGET \ + --preload-template "HSET __key__ field __data__" \ + --key-maximum 100000000 \ + --test-time 60 --warmup-time 10 --trials 1 \ + \ + --dfly-bench-threads 32 --dfly-bench-conns 5 --dfly-bench-pipeline 30 \ + --dfly-bench-data-size 128 --dfly-bench-key-dist U --dfly-bench-qps 0 \ + \ + --dragonfly_version s3://df-use1-pub/v2/dragonfly-aarch64.tar.gz \ + --dragonfly_num_shards 16 +``` + +### Expected results: + + HGET + Engine Throughput (median) p50 p99 p99.9 Avg Latency + ────── ─────────────────── ─── ─── ───── ─────────── + dragonfly 9.31M ops/s 0.468 ms 0.875 ms 1.163 ms 0.489 ms + redis 744.4K ops/s 6.522 ms 9.979 ms 13.329 ms 6.429 ms + valkey 791.8K ops/s 6.483 ms 6.990 ms 6.999 ms 6.044 ms + \ No newline at end of file diff --git a/benchmarks/HSET/dfly_bench/HSET_reproduce.md b/benchmarks/HSET/dfly_bench/HSET_reproduce.md new file mode 100644 index 0000000..f47d554 --- /dev/null +++ b/benchmarks/HSET/dfly_bench/HSET_reproduce.md @@ -0,0 +1,38 @@ +### Stateful setup: + +``` +dfbench setup --id maxqps2 --tune-network --ubuntu-version 24.04 \ + --server-instance m7g.8xlarge --server-arch arm64 \ + --client-instance c6gn.8xlarge --client-arch arm64 \ + --engines dragonfly,redis,valkey --availability-zone us-east-1c +``` + +### Test run: + +``` +dfbench run-commands-dfly \ + --id maxqps2 --engines dragonfly,redis,valkey \ + \ + --command-template "HSET __key__ field __data__" --command-name HSET --command-mutates \ + --key-maximum 100000000 \ + --test-time 300 --warmup-time 10 --trials 1 \ + \ + --dfly-bench-threads 32 --dfly-bench-conns 5 --dfly-bench-pipeline 30 \ + --dfly-bench-data-size 128 --dfly-bench-key-dist U --dfly-bench-qps 0 \ + \ + --dragonfly_version s3://df-use1-pub/v2/dragonfly-aarch64.tar.gz \ + --dragonfly_num_shards 16 --dragonfly_conn_use_incoming_cpu=true \ + \ + --redis_io_threads 8 --valkey_io_threads 8 +``` + +### Expected results: + + HSET + Engine Throughput (median) p50 p99 p99.9 Avg Latency + ────── ─────────────────── ─── ─── ───── ─────────── + dragonfly 7.99M ops/s 0.507 ms 1.985 ms 7.332 ms 0.583 ms + redis 638.5K ops/s 6.839 ms 11.613 ms 18.655 ms 7.503 ms + valkey 704.0K ops/s 6.573 ms 13.470 ms 17.760 ms 6.804 ms + + diff --git a/cmd/run_commands.go b/cmd/run_commands.go index a640d29..edeb6da 100644 --- a/cmd/run_commands.go +++ b/cmd/run_commands.go @@ -38,6 +38,16 @@ func newRunCommandsCommand() *cobra.Command { distinctSeed bool output string + // Custom command spec (alternative to --command; see spec.go). + commandTemplate string + commandName string + commandKeyPattern string + commandMutates bool + commandDrains bool + preloadTemplate string + preloadItems int + commandKeyMaximum int + // Dragonfly-specific tuning. dflyNumShards int dflyConnIncoming bool @@ -87,9 +97,18 @@ marked approximate in the output.`, if !ok { return fmt.Errorf("invalid --engines value %q: must be comma-separated list of dragonfly,redis,valkey", enginesFlag) } - spec, err := commands.Lookup(commandFlag) + spec, err := resolveCommandSpec(cmd.Flags(), commandFlag, commands.CustomSpecInput{ + Name: commandName, + Template: commandTemplate, + KeyPattern: commandKeyPattern, + Mutates: commandMutates, + Drains: commandDrains, + PreloadTemplate: preloadTemplate, + PreloadItems: preloadItems, + KeyMaximum: commandKeyMaximum, + }) if err != nil { - return fmt.Errorf("invalid --command: %w", err) + return err } if trials < 1 { return fmt.Errorf("--trials must be >= 1, got %d", trials) @@ -149,7 +168,29 @@ marked approximate in the output.`, f.StringVar(&enginesFlag, "engines", "dragonfly,redis,valkey", "Comma-separated list of engines to benchmark (must be a subset of the setup's installed engines)") f.StringVar(&commandFlag, "command", "", - "Command to benchmark (required), one of: "+strings.Join(commands.SupportedNames(), ",")) + "Command to benchmark, one of: "+strings.Join(commands.SupportedNames(), ",")+ + ". Not required when --command-template is given.") + + // Custom command spec: define a command not in the built-in registry. + // Mutually exclusive with --command. + f.StringVar(&commandTemplate, "command-template", "", + "Custom command in memtier arbitrary-command form, e.g. \"SETEX __key__ 3600 __data__\". "+ + "Fills __key__/__data__ like the built-in commands. Mutually exclusive with --command.") + f.StringVar(&commandName, "command-name", "", + "Label for the custom command (defaults to the first word of --command-template); used in output files and results.") + f.StringVar(&commandKeyPattern, "memtier-command-key-pattern", "R", + "memtier --command-key-pattern for a custom --command-template's measured pass: R (random), S (sequential), P (parallel), G (gaussian). Only used with --command-template.") + f.BoolVar(&commandMutates, "command-mutates", false, + "Mark the custom command as mutating the keyspace: flush + re-preload before every trial.") + f.BoolVar(&commandDrains, "command-drains", false, + "Mark the custom command as draining preloaded data (DEL, LPOP, ...); implies --command-mutates and marks results approximate.") + f.StringVar(&preloadTemplate, "preload-template", "", + "Optional preload command for the custom command, e.g. \"SET __key__ __data__\"; must contain __key__.") + f.IntVar(&preloadItems, "preload-items", 1, + "Items written per key by --preload-template (list/set/zset cardinality).") + f.IntVar(&commandKeyMaximum, "command-key-maximum", 0, + "Per-command key range override for the custom command (0 = use --key-maximum). Required when --preload-items > 1.") + f.IntVar(&threads, "memtier-threads", 32, "Number of memtier threads") f.IntVar(&clients, "memtier-clients", 5, "Number of clients per memtier thread") f.IntVar(&pipeline, "memtier-pipeline", 20, "Requests pipelined per response wait") @@ -198,6 +239,36 @@ marked approximate in the output.`, return cmd } +// resolveCommandSpec picks between a built-in registry command (--command) +// and a fully custom spec (--command-template). The two are mutually +// exclusive and exactly one must be provided. Shared by run-commands and +// run-commands-dfly. +func resolveCommandSpec(f *pflag.FlagSet, commandFlag string, custom commands.CustomSpecInput) (commands.CommandSpec, error) { + hasName := strings.TrimSpace(commandFlag) != "" + hasTemplate := f.Changed("command-template") || strings.TrimSpace(custom.Template) != "" + + switch { + case hasName && hasTemplate: + return commands.CommandSpec{}, fmt.Errorf("--command and --command-template are mutually exclusive; pass only one") + case hasTemplate: + spec, err := commands.BuildCustom(custom) + if err != nil { + return commands.CommandSpec{}, fmt.Errorf("invalid custom command: %w", err) + } + return spec, nil + case hasName: + spec, err := commands.Lookup(commandFlag) + if err != nil { + return commands.CommandSpec{}, fmt.Errorf("invalid --command: %w", err) + } + return spec, nil + default: + return commands.CommandSpec{}, fmt.Errorf( + "provide either --command (one of: %s) or --command-template", + strings.Join(commands.SupportedNames(), ", ")) + } +} + // engineTuning holds the per-engine tuning knobs resolved from the // engine-prefixed flags and their deprecated shared aliases. type engineTuning struct { diff --git a/cmd/run_commands_dfly.go b/cmd/run_commands_dfly.go index f55c40c..fdf4859 100644 --- a/cmd/run_commands_dfly.go +++ b/cmd/run_commands_dfly.go @@ -40,6 +40,15 @@ func newRunCommandsDflyCommand() *cobra.Command { dflyBenchRef string output string + // Custom command spec (alternative to --command; see spec.go). + commandTemplate string + commandName string + commandMutates bool + commandDrains bool + preloadTemplate string + preloadItems int + commandKeyMaximum int + // Dragonfly-specific server tuning (identical to run-commands). dflyNumShards int dflyConnIncoming bool @@ -78,9 +87,20 @@ compare'.`, if !ok { return fmt.Errorf("invalid --engines value %q: must be comma-separated list of dragonfly,redis,valkey", enginesFlag) } - spec, err := commands.Lookup(commandFlag) + // No KeyPattern here: dfly_bench drives the measured distribution + // via --dfly-bench-key-dist, not memtier's key pattern. BuildCustom + // defaults the (unused) spec KeyPattern to R. + spec, err := resolveCommandSpec(cmd.Flags(), commandFlag, commands.CustomSpecInput{ + Name: commandName, + Template: commandTemplate, + Mutates: commandMutates, + Drains: commandDrains, + PreloadTemplate: preloadTemplate, + PreloadItems: preloadItems, + KeyMaximum: commandKeyMaximum, + }) if err != nil { - return fmt.Errorf("invalid --command: %w", err) + return err } if trials < 1 { return fmt.Errorf("--trials must be >= 1, got %d", trials) @@ -136,7 +156,26 @@ compare'.`, f.StringVar(&enginesFlag, "engines", "dragonfly", "Comma-separated list of engines to benchmark (must be a subset of the setup's installed engines)") f.StringVar(&commandFlag, "command", "", - "Command to benchmark (required), one of: "+strings.Join(commands.SupportedNames(), ",")) + "Command to benchmark, one of: "+strings.Join(commands.SupportedNames(), ",")+ + ". Not required when --command-template is given.") + + // Custom command spec: define a command not in the built-in registry. + // Mutually exclusive with --command. + f.StringVar(&commandTemplate, "command-template", "", + "Custom command in dfly_bench arbitrary-command form, e.g. \"SETEX __key__ 3600 __data__\". "+ + "Fills __key__/__data__ like the built-in commands. Mutually exclusive with --command.") + f.StringVar(&commandName, "command-name", "", + "Label for the custom command (defaults to the first word of --command-template); used in output files and results.") + f.BoolVar(&commandMutates, "command-mutates", false, + "Mark the custom command as mutating the keyspace: flush + re-preload before every trial.") + f.BoolVar(&commandDrains, "command-drains", false, + "Mark the custom command as draining preloaded data (DEL, LPOP, ...); implies --command-mutates and marks results approximate.") + f.StringVar(&preloadTemplate, "preload-template", "", + "Optional preload command for the custom command, e.g. \"SET __key__ __data__\"; must contain __key__.") + f.IntVar(&preloadItems, "preload-items", 1, + "Items written per key by --preload-template (list/set/zset cardinality).") + f.IntVar(&commandKeyMaximum, "command-key-maximum", 0, + "Per-command key range override for the custom command (0 = use --key-maximum). Required when --preload-items > 1.") // dfly_bench (client load generator) flags. Named --dfly-bench-* to make // it explicit these tune the dfly_bench tool, not the Dragonfly server // (--dragonfly_*), and to parallel run-commands' --memtier-* flags. diff --git a/cmd/run_commands_test.go b/cmd/run_commands_test.go new file mode 100644 index 0000000..700fef7 --- /dev/null +++ b/cmd/run_commands_test.go @@ -0,0 +1,57 @@ +package cmd + +import ( + "testing" + + "github.com/spf13/pflag" + + "github.com/dragonflydb/benchmarking/bench/commands" +) + +// newCommandFlags returns a flag set with just the --command-template flag +// registered, matching what resolveCommandSpec inspects via Changed. +func newCommandFlags(templateChanged bool, template string) *pflag.FlagSet { + f := pflag.NewFlagSet("test", pflag.ContinueOnError) + var tmpl string + f.StringVar(&tmpl, "command-template", "", "") + if templateChanged { + _ = f.Set("command-template", template) + } + return f +} + +func TestResolveCommandSpec_Registry(t *testing.T) { + f := newCommandFlags(false, "") + spec, err := resolveCommandSpec(f, "SET", commands.CustomSpecInput{}) + if err != nil { + t.Fatalf("resolveCommandSpec: %v", err) + } + if spec.Name != "SET" { + t.Errorf("Name = %q, want SET", spec.Name) + } +} + +func TestResolveCommandSpec_Custom(t *testing.T) { + f := newCommandFlags(true, "GET __key__") + spec, err := resolveCommandSpec(f, "", commands.CustomSpecInput{Template: "GET __key__"}) + if err != nil { + t.Fatalf("resolveCommandSpec: %v", err) + } + if spec.Name != "GET" || spec.Template != "GET __key__" { + t.Errorf("spec = %+v, want custom GET", spec) + } +} + +func TestResolveCommandSpec_MutuallyExclusive(t *testing.T) { + f := newCommandFlags(true, "GET __key__") + if _, err := resolveCommandSpec(f, "SET", commands.CustomSpecInput{Template: "GET __key__"}); err == nil { + t.Fatal("expected error when both --command and --command-template are set") + } +} + +func TestResolveCommandSpec_Neither(t *testing.T) { + f := newCommandFlags(false, "") + if _, err := resolveCommandSpec(f, "", commands.CustomSpecInput{}); err == nil { + t.Fatal("expected error when neither --command nor --command-template is set") + } +}