cockroachdb / cockroachdb/cockroach
kvserver/allocator: restore MMA coverage on allocator unit tests pinned to legacy mode
- Dominant language
- Go
- Stars
- 32.5k
- Forks
- 4.1k
- PR merge metrics
- PR metrics pending
Description
**Summary:**
[#169430](https://github.com/cockroachdb/cockroach/pull/169430) made MMA the default load-based rebalancer in v26.3 by flipping the `kv.allocator.load_based_rebalancing` cluster setting default to `auto`. As an expedient unblock, the allocator unit-test helpers `CreateTestAllocator{,WithKnobs}` ([test_helpers.go:51](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/test_helpers.go#L51)) and `MakeAllocatorSync` ([allocator_test.go:64](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go#L64)) override the mode to `LBRebalancingLeasesAndReplicas`. This pins ~75 test functions to the legacy code path. Many of those tests are not legacy-implementation-specific and would (if able to run) provide coverage against MMA. As wenyihu6 [noted](https://github.com/cockroachdb/cockroach/pull/169430#pullrequestreview-4232669781) on the PR, we lose MMA coverage on the allocator-agnostic ones.
**Findings — classification of affected tests:**
**A. Legacy-allocator-specific (pin OK; delete with legacy):**
- `store_rebalancer_test.go`: `TestChooseLeaseToTransfer`, `TestChooseRangeToRebalance{Random,AcrossHeterogeneousZones,IgnoresRangeOnBestStores,OffHotNodes}`, `TestNoLeaseTransferToBehindReplicas`, `TestStoreRebalancerIOOverloadCheck`, `TestStoreRebalancerHotRangesLogging`, `TestReplicateQueueMaxSize`.
- `allocator_test.go`: `TestAllocatorRebalanceByQPS`, `TestAllocatorRemoveBasedOnQPS`, `TestAllocatorRebalanceByCount`, `TestAllocatorRebalanceBasedOnRangeCount`, `TestAllocatorRebalanceThrashing`, `TestAllocatorTransferLeaseTargetLoadBased`, `TestAllocatorShouldTransferLease`, `TestAllocatorShouldTransferLeaseDraining`, `TestAllocatorShouldTransferSuspected`, `TestAllocatorShouldTransferLeaseIOOverload`.
- `allocator_impl_test.go`: `TestAllocatorRebalanceTarget`.
**B. MMA-specific (pin SILENTLY NEUTRALIZES — needs real MMA setup):**
- `TestAllocatorTransferLeaseTargetMMAConflict`
- `TestAllocatorRebalanceMMAConflict`
- `TestBestRebalanceTarget`
**C. Allocator-agnostic (should also run under MMA):**
- *No MMA path traversed:* `TestAllocatorSimpleRetrieval`, `TestAllocator{NoAvailableDisks,FullDiskError,AllocateVoterIOOverloadCheck,TwoDatacenters,ExistingReplica,ReplaceDecommissioningReplica,ReplaceFailsOnConstrainedDecommissioningReplica,MultipleStoresPerNode,MultipleStoresPerNodeLopsided,RebalanceDeadNodes,RebalanceDifferentLocalitySizes,LeasePreferences,LeasePreferencesMultipleStoresPerLocality,RemoveBasedOnDiversity,ConstraintsAndVoterConstraints,AllocateTargetLocality,RebalanceTargetLocality}`, `TestAllocate{CandidatesExcludeNonReadyNodes,CandidatesNumReplicasConstraints}`, `TestAllocatorNonVoterAllocationExcludesVoterNodes`, `TestRemoveCandidatesNumReplicasConstraints`, `TestVotersCanRebalanceToNonVoterStores`, `TestNonVotersCannotRebalanceToVoterStores`, `TestRebalanceCandidatesNumReplicasConstraints`, `TestAllocatorRemoveTargetBasedOnCapacity`, `TestAllocatorComputeAction*` (~10 variants), `TestAllocatorRemoveLearner`, `TestNonVoterPrioritizationInVoterAdditions`, `TestAllocatorRebalanceTargetVoterConstraintUnsatisfied`, `TestAllocatorComputeActionNoStorePool`, `TestAllocatorRebalanceDeterminism`, `TestAllocatorRebalanceWithScatter`, `TestAllocatorRebalanceAway`.
- *DO traverse MMA path (`TransferLeaseTarget` / `RebalanceVoter`):* `TestAllocatorTransferLeaseTarget`, `TestAllocatorTransferLeaseTargetIOOverloadCheck`, `TestAllocatorTransferLeaseToReplicasNeeding{Snapshot,Catchup}`, `TestAllocatorTransferLeaseTargetConstraints`, `TestAllocatorTransferLeaseTargetDraining`, `TestAllocatorRebalanceNonVoters`, `TestAllocatorRebalanceIOOverloadCheck`.
- *Other:* `TestAllocatorThrottled` (`allocator_impl_test.go`), `TestAllocatorCheckRange` (`store_test.go`).
**Code References:**
- [pkg/kv/kvserver/allocator/allocatorimpl/test_helpers.go:51](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/test_helpers.go#L51) — pin in `CreateTestAllocatorWithKnobs`
- [pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go:64](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go#L64) — pin in `MakeAllocatorSync`
- [pkg/kv/kvserver/mmaintegration/thrashing.go:108](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/mmaintegration/thrashing.go#L108) — `BuildMMARebalanceAdvisor` returns Noop unless mode is `multi-metric and count`
- [pkg/kv/kvserver/mmaintegration/allocator_sync.go](https://github.com/cockroachdb/cockroach/blob/master/pkg/kv/kvserver/mmaintegration/allocator_sync.go) — existing `OverrideIsInConflictWithMMA` knob
**Next Steps:**
- [ ] Cheap unblock: extend `mmaintegration.TestingKnobs` with a knob that also short-circuits `BuildMMARebalanceAdvisor` (return `NoopMMARebalanceAdvisor`). Replace the pin in `CreateTestAllocator{,WithKnobs}` and `MakeAllocatorSync` with that knob. Category C tests then run with the MMA-default cluster setting; MMA dispatch path is exercised but does not panic on uninitialized `mmaprototype` state.
- [ ] Re-classify category B tests. They were intended to assert MMA conflict behavior but are silently neutralized today. Migrate them to construct a usable `mmaprototype.AllocatorState` (call `SetStore` + feed `ProcessStoreLoadMsg`) so MMA contributes to the assertion.
- [ ] For a sample of category C tests that go through the MMA path (e.g. `TestAllocatorTransferLeaseTargetIOOverloadCheck`), do the same MMA setup so we get real cross-implementation coverage.
- [ ] Once #169444 retires the pre-MMA helper paradigm entirely, drop the category-A tests.
Related: #169430, #169444.
Epic CRDB-56265
Jira issue: CRDB-63675
Contributor guide
Assessment
This issue has not been assessed yet.