cockroachdb / cockroachdb/cockroach
kvserver/allocator: pre-MMA test helpers pin legacy mode; migrate to MMA-aware setup
- Dominant language
- Go
- Stars
- 32.5k
- Forks
- 4.1k
- PR merge metrics
- PR metrics pending
Description
**Summary:**
When v26.3 made MMA the default load-based rebalancer (#169430), the long-standing allocator unit-test helpers (`CreateTestAllocator{,WithKnobs}` in [test_helpers.go](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/test_helpers.go) and `MakeAllocatorSync` in [allocator_test.go](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go)) began panicking. They construct an `mmaprototype.AllocatorState` but never call `SetStore`, so the moment any decision path reaches `mmaintegration.AllocatorSync.BuildMMARebalanceAdvisor` (e.g. `bestRebalanceTarget` → `BuildMMARebalanceAdvisor` → `mmaprototype/load.go:498`) the test nil-derefs.
As an expedient workaround, both helpers now call `kvserverbase.OverrideLoadBasedRebalancingMode(ctx, &st.SV, kvserverbase.LBRebalancingLeasesAndReplicas)`, pinning every test that uses them to the legacy single-metric path. This preserves the prior implicit behavior but means roughly 75 test call sites continue to exercise the dying allocator instead of MMA.
**Findings:**
- `CreateTestAllocator{,WithKnobs}`: 67 callers across 5 files
- `pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go` (54)
- `pkg/kv/kvserver/store_rebalancer_test.go` (9)
- `pkg/kv/kvserver/allocator_impl_test.go` (2)
- `pkg/kv/kvserver/allocator/allocatorimpl/allocator_scorer_test.go` (1)
- `pkg/kv/kvserver/store_test.go` (1)
- `MakeAllocatorSync` direct callers (bypass the helper): 8, all in `allocator_test.go`.
- Workaround pin sites added in #169430:
- [test_helpers.go:51](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/test_helpers.go#L51)
- [allocator_test.go:64](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go#L64)
**Code References:**
- Panic site: [pkg/kv/kvserver/allocator/mmaprototype/load.go:498](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/mmaprototype/load.go#L498)
- Decision path: [pkg/kv/kvserver/mmaintegration/thrashing.go:114](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/mmaintegration/thrashing.go#L114) → `mmaprototype.BuildMMARebalanceAdvisor`
**Next Steps:**
- [ ] Decide on the right end state: extend the test helpers to set up a usable `mmaprototype` (SetStore for each store the test creates), OR migrate tests off the helpers to a fixture that exercises the real Store/MMA lifecycle.
- [ ] Once helpers are MMA-safe, drop the `OverrideLoadBasedRebalancingMode(..., LBRebalancingLeasesAndReplicas)` calls in `test_helpers.go` and `allocator_test.go`.
- [ ] Re-run the affected packages with the pin removed; triage any new fallout file by file.
- [ ] Consider whether some allocator-level tests are inherently legacy-only and should be deleted rather than ported.
Epic CRDB-56265
Jira issue: CRDB-63453
Contributor guide
Assessment
This issue has not been assessed yet.