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/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", 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 */)