From c09e0f8d66ceb82a5dec8f39953c7865a8352b10 Mon Sep 17 00:00:00 2001 From: Tobias Grieger Date: Wed, 20 May 2026 15:50:55 +0200 Subject: [PATCH 1/2] kvserver: co-construct AllocatorSync in NewStore Production code (server.go) pairs MMAllocator with an AllocatorSync on the StoreConfig. A number of test entry points populate MMAllocator (via TestStoreConfig) and StorePool but leave AllocatorSync unset. The MMA store rebalancer worked around this with a buildutil-gated fallback that lazily constructed an AllocatorSync in test builds and panicked otherwise. That gate was tripping benchmarks (which are compiled without the crdb_test tag), e.g. BenchmarkStoreGetReplica and BenchmarkMVCCGCWithForegroundTraffic. Co-construct AllocatorSync inside NewStore when MMAllocator is set but AllocatorSync is not, mirroring what server.go already does explicitly. The fallback in newMMAStoreRebalancer is replaced with a plain non-nil assertion, and the now-redundant explicit construction in localtestcluster is dropped. Fixes #170455. Fixes #170456. Release note: None Epic: none Co-Authored-By: roachdev-claude --- pkg/kv/kvserver/mma_store_rebalancer.go | 11 +---------- pkg/kv/kvserver/store.go | 9 +++++++++ pkg/testutils/localtestcluster/local_test_cluster.go | 5 ++--- 3 files changed, 12 insertions(+), 13 deletions(-) diff --git a/pkg/kv/kvserver/mma_store_rebalancer.go b/pkg/kv/kvserver/mma_store_rebalancer.go index 805904e21954..eebe50ee0b5e 100644 --- a/pkg/kv/kvserver/mma_store_rebalancer.go +++ b/pkg/kv/kvserver/mma_store_rebalancer.go @@ -18,7 +18,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration" "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/settings/cluster" - "github.com/cockroachdb/cockroach/pkg/util/buildutil" "github.com/cockroachdb/cockroach/pkg/util/log" "github.com/cockroachdb/cockroach/pkg/util/stop" "github.com/cockroachdb/cockroach/pkg/util/timeutil" @@ -62,15 +61,7 @@ func newMMAStoreRebalancer( mma.SetDiskUtilThresholds(opts.RebalanceToThreshold, opts.ShedAndBlockAllThreshold) as := s.cfg.AllocatorSync if as == nil { - // Production code (server.go) is expected to populate AllocatorSync on - // the StoreConfig. A handful of test entry points (TestStoreConfig and - // the assorted store_rebalancer_test setups) populate StorePool but not - // AllocatorSync; fall back to constructing one here so those tests don't - // trip a nil-receiver panic from the background rebalancer goroutine. - if !buildutil.CrdbTestBuild { - panic(errors.AssertionFailedf("AllocatorSync must be set on StoreConfig outside of test builds")) - } - as = mmaintegration.NewAllocatorSync(s.cfg.StorePool, mma, st, nil /* knobs */) + panic(errors.AssertionFailedf("AllocatorSync must be set on StoreConfig")) } return &mmaStoreRebalancer{ store: (*mmaStore)(s), diff --git a/pkg/kv/kvserver/store.go b/pkg/kv/kvserver/store.go index d3b2bdf4e426..2b6f687ada94 100644 --- a/pkg/kv/kvserver/store.go +++ b/pkg/kv/kvserver/store.go @@ -1491,6 +1491,15 @@ func NewStore( if !cfg.Valid() { log.KvExec.Fatalf(ctx, "invalid store configuration: %+v", &cfg) } + // Production code (server.go) pairs MMAllocator with an AllocatorSync on + // the StoreConfig. A number of test entry points populate MMAllocator (via + // TestStoreConfig) and StorePool but leave AllocatorSync unset; co-construct + // it here so the rebalancer and queues always observe a non-nil receiver. + if cfg.AllocatorSync == nil && cfg.MMAllocator != nil { + cfg.AllocatorSync = mmaintegration.NewAllocatorSync( + cfg.StorePool, cfg.MMAllocator, cfg.Settings, nil, /* knobs */ + ) + } iot := ioThresholds{} iot.Replace(nil, 1.0) // init as empty s := &Store{ diff --git a/pkg/testutils/localtestcluster/local_test_cluster.go b/pkg/testutils/localtestcluster/local_test_cluster.go index 2cea8532e2b3..0d226fb6748c 100644 --- a/pkg/testutils/localtestcluster/local_test_cluster.go +++ b/pkg/testutils/localtestcluster/local_test_cluster.go @@ -27,7 +27,6 @@ import ( "github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvstorage" "github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness" "github.com/cockroachdb/cockroach/pkg/kv/kvserver/load" - "github.com/cockroachdb/cockroach/pkg/kv/kvserver/mmaintegration" "github.com/cockroachdb/cockroach/pkg/kv/kvserver/storeliveness" "github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities/tenantcapabilitiesauthorizer" "github.com/cockroachdb/cockroach/pkg/roachpb" @@ -250,8 +249,8 @@ func (ltc *LocalTestCluster) Start(t testing.TB, initFactory InitFactoryFn) { ) cfg.MMAllocator = mmaprototype.NewAllocatorState(timeutil.DefaultTimeSource{}, rand.New(rand.NewSource(timeutil.Now().UnixNano()))) - cfg.AllocatorSync = mmaintegration.NewAllocatorSync( - cfg.StorePool, cfg.MMAllocator, cfg.Settings, nil /* knobs */) + // AllocatorSync is co-constructed by kvserver.NewStore from MMAllocator and + // StorePool. cfg.Transport = kvserver.NewDummyRaftTransport(cfg.AmbientCtx, cfg.Settings, ltc.Clock) cfg.ClosedTimestampReceiver = sidetransport.NewReceiver(nc, ltc.stopper, ltc.Stores, nil /* testingKnobs */) From ca3a5de768fe0caa6453c1f47c66a921e15ec40e Mon Sep 17 00:00:00 2001 From: Tobias Grieger Date: Wed, 20 May 2026 17:41:10 +0200 Subject: [PATCH 2/2] kvserver: regenerate BUILD.bazel for localtestcluster Follow-up to dropping the mmaintegration import from pkg/testutils/localtestcluster/local_test_cluster.go. Release note: None Epic: none Co-Authored-By: roachdev-claude --- pkg/testutils/localtestcluster/BUILD.bazel | 1 - 1 file changed, 1 deletion(-) diff --git a/pkg/testutils/localtestcluster/BUILD.bazel b/pkg/testutils/localtestcluster/BUILD.bazel index fc27392c60f9..b69e974a30d6 100644 --- a/pkg/testutils/localtestcluster/BUILD.bazel +++ b/pkg/testutils/localtestcluster/BUILD.bazel @@ -22,7 +22,6 @@ go_library( "//pkg/kv/kvserver/kvstorage", "//pkg/kv/kvserver/liveness", "//pkg/kv/kvserver/load", - "//pkg/kv/kvserver/mmaintegration", "//pkg/kv/kvserver/storeliveness", "//pkg/multitenant/tenantcapabilities/tenantcapabilitiesauthorizer", "//pkg/roachpb",