mma: [dnm] add cluster state message handling

TODO(kvoli)... Informs: #103320 Release note: None
cockroachdb · Feb 26, 2025 · 8d394bf · 8d394bf
1 parent 09d4505
commit 8d394bf
Show file tree

Hide file tree

Showing 11 changed files with 1,619 additions and 195 deletions.
diff --git a/pkg/kv/kvserver/allocator/mma/BUILD.bazel b/pkg/kv/kvserver/allocator/mma/BUILD.bazel
@@ -16,13 +16,16 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/roachpb",
+        "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_cockroachdb_redact//:redact",
     ],
 )
 
 go_test(
     name = "mma_test",
     srcs = [
+        "cluster_state_test.go",
         "constraint_matcher_test.go",
         "constraint_test.go",
         "load_test.go",
@@ -32,6 +35,10 @@ go_test(
     embed = [":mma"],
     deps = [
         "//pkg/roachpb",
+        "//pkg/spanconfig/spanconfigtestutils",
+        "//pkg/testutils/datapathutils",
+        "//pkg/util/log",
+        "//pkg/util/timeutil",
         "@com_github_cockroachdb_datadriven//:datadriven",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//require",

diff --git a/pkg/kv/kvserver/allocator/mma/allocator_state.go b/pkg/kv/kvserver/allocator/mma/allocator_state.go
@@ -9,6 +9,7 @@ import (
 	"sync"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 )
 
 type allocatorState struct {
@@ -27,9 +28,9 @@ type allocatorState struct {
 	changeRangeLimiter *storeChangeRateLimiter
 }
 
-func newAllocatorState() *allocatorState {
+func newAllocatorState(ts timeutil.TimeSource) *allocatorState {
 	interner := newStringInterner()
-	cs := newClusterState(interner)
+	cs := newClusterState(ts, interner)
 	return &allocatorState{
 		cs:                     cs,
 		rangesNeedingAttention: map[roachpb.RangeID]struct{}{},