kvserver/mmaprototype: add BenchmarkRebalanceStores by tbg · Pull Request #165284 · cockroachdb/cockroach

tbg · 2026-03-10T09:31:55Z

Summary

Reduce per-iteration allocations in the MMA rebalance loop from ~2000 to ~160
allocs/op, primarily by guarding expensive logging and pooling hot slices/maps.

name                old time/op    new time/op    delta
RebalanceStores-10     246µs ± 3%     104µs ± 2%  -57.87%  (p=0.000 n=10+10)

name                old alloc/op   new alloc/op   delta
RebalanceStores-10     101kB ± 0%      19kB ± 0%  -81.17%  (p=0.000 n=9+8)

name                old allocs/op  new allocs/op  delta
RebalanceStores-10     1.99k ± 0%     0.16k ± 0%  -91.75%  (p=0.000 n=10+10)

Commit arc

Add BenchmarkRebalanceStores — microbenchmark for the rebalance loop
(4 stores, 1000 ranges, 100 hot). Uses undoPendingChange to restore state
between iterations. Baseline: 2013 allocs/op.
Annotate top allocation sites — TODO comments at each site found via
memory profiling, guiding the subsequent commits.
Guard logging in sortTargetCandidateSetAndPick — extract
formatCandidatesLog to a package-level function with an internal
ExpensiveLogEnabled check, removing redundant guards at all 6 call sites.
2013 → 1689 allocs/op.
Guard top-K debug logging — wrap the redact.StringBuilder +
LoadVector formatting block with ExpensiveLogEnabled.
1689 → 778 allocs/op.
Guard lease-transfer candidate logging — prevent boxing candsPL into
interface{} when verbose logging is off. 778 → 678 allocs/op.
Pool slice and map allocations in rebalance loop — introduce generic
slicePool[T] and mapPool[K, V] wrappers around sync.Pool, replacing
manual pool boilerplate and the rebalanceEnv.scratch struct (which had
aliasing hazards). Five pools cover slices (storeAndLeasePreference,
StoreID, candidateInfo) and maps (NodeID→*NodeLoad,
StoreID→struct{}). candidatesToMoveLease accepts a reusable buffer.
678 → 160 allocs/op.

Epic: CRDB-55052

trunk-io · 2026-03-10T09:32:00Z

😎 Merged successfully - details.

cockroach-teamcity · 2026-03-10T09:32:10Z

This change is

pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_rebalance_stores.go

angeladietz

some nits but lgtm overall, nothing blocking!

I'm also curious how you attributed the number of allocs to each line (commit 2). was it just with pprof (& your agent)?

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_rebalance_stores_bench_test.go

pkg/kv/kvserver/allocator/mmaprototype/allocator_state.go

pkg/kv/kvserver/allocator/mmaprototype/load.go

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_rebalance_stores.go

pkg/kv/kvserver/allocator/mmaprototype/constraint.go

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_rebalance_stores_bench_test.go

tbg

Yes I told the agent to run the benchmark with memprofile and to top50 the output by allocation count, then iterated.

@tbg partially reviewed 10 files and all commit messages, made 3 comments, resolved 5 discussions, and dismissed @angeladietz from a discussion.
Reviewable status: complete! 0 of 0 LGTMs obtained (waiting on angeladietz).

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_rebalance_stores_bench_test.go

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_rebalance_stores.go

Add a microbenchmark for the MMA rebalance loop (rebalanceStores). The benchmark uses the existing undoPendingChange production code to reverse mutations after each iteration, plus targeted resets of leaked fields (lastFailedChange, overload times, meansMemo cache). Setup: 4 stores, 1000 ranges (100 hot, 900 cold), store 1 overloaded. BenchmarkRebalanceStores 100 323k ns/op 102kB/op 2013 allocs/op Epic: none Release note: None Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

Add TODO(tbg) comments at each allocation site found via BenchmarkRebalanceStores memory profiling (4 stores, 1000 ranges, 100 hot). The allocs/op counts are from a 1000-iteration run. Epic: none Release note: None Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

The `formatCandidatesLog` closure and `redact.StringBuilder` were allocated on every call to `sortTargetCandidateSetAndPick`, even when verbose logging at level 2 was disabled. Additionally, the standalone VEventf for discarded candidates boxed multiple arguments into interfaces unconditionally. Extract `formatCandidatesLog` to a package-level function and guard all call sites (plus the standalone VEventf) with `log.ExpensiveLogEnabled`. BenchmarkRebalanceStores 1999 -> 1689 allocs/op (-310) Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

…bled The top-K debug logging block in `processSheddingStore` builds a `redact.StringBuilder` and formats `LoadVector` values for every range in the top-K set on every call, even when verbose logging at level 2 is disabled. Wrap the entire block with `log.ExpensiveLogEnabled`. BenchmarkRebalanceStores 1689 -> 778 allocs/op (-911) Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

…Enabled The VEventf call logging lease-transfer candidates boxes `candsPL` (a `storeSet`) to `interface{}`, triggering formatting even when verbose logging at level 2 is disabled. Guard with `log.ExpensiveLogEnabled`. BenchmarkRebalanceStores 778 -> 678 allocs/op (-100) Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

Reduce per-iteration allocations by pooling hot slices and maps via generic `slicePool[T]` and `mapPool[K, V]` wrappers around `sync.Pool`. These handle the `New` func, type assertion on `get`, and `clear`+reset on `put`. Five pools are introduced: - `storeAndLeasePreferencePool`, `storeIDPool`, `candidateInfoPool` for slices - `nodeLoadMapPool`, `storeIDStructMapPool` for maps The `rebalanceEnv.scratch` struct (which had aliasing hazards) is removed in favor of these pools. `candidatesToMoveLease` now accepts a reusable buffer. 678 → 160 allocs/op. Co-Authored-By: roachdev-claude <roachdev-claude-bot@cockroachlabs.com>

tbg · 2026-03-16T10:35:54Z

/trunk merge

tbg force-pushed the tbg/mma-rebalance-bench branch 4 times, most recently from f7acd25 to 6e2a60c Compare March 10, 2026 09:42