mmaprototype: add datadriven test for rebalanceStores

Author: tbg

pkg/kv/kvserver/allocator/mmaprototype/cluster_state_test.go

@@ -8,6 +8,7 @@ package mmaprototype
 import (
 	"context"
 	"fmt"
+	"math/rand"
 	"slices"
 	"sort"
 	"strconv"
@@ -513,6 +514,13 @@ func TestClusterState(t *testing.T) {
 				case "get-pending-changes":
 					return printPendingChangesTest(testingGetPendingChanges(t, cs))
 
+				case "rebalance-stores":
+					storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
+					rng := rand.New(rand.NewSource(0))
+					dsm := newDiversityScoringMemo()
+					cs.rebalanceStores(context.Background(), storeID, rng, dsm)
+					return printPendingChangesTest(testingGetPendingChanges(t, cs))
+
 				case "tick":
 					seconds := dd.ScanArg[int](t, d, "seconds")
 					ts.Advance(time.Second * time.Duration(seconds))

pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/rebalance_stores_basic

@@ -0,0 +1,68 @@
+# Basic test that rebalanceStores issues a lease transfer when seeing a single replica
+# contributing to overload on the first out of three single-store nodes.
+
+set-store
+  store-id=1 node-id=1 attrs=purple locality-tiers=region=us-west-1,zone=us-west-1a
+  store-id=2 node-id=2 attrs=yellow locality-tiers=region=us-east-1,zone=us-east-1a
+  store-id=3 node-id=3 attrs=green locality-tiers=region=us-central-1,zone=us-central-1a
+----
+node-id=1 locality-tiers=region=us-west-1,zone=us-west-1a,node=1
+  store-id=1 attrs=purple locality-code=1:2:3:
+node-id=2 locality-tiers=region=us-east-1,zone=us-east-1a,node=2
+  store-id=2 attrs=yellow locality-code=4:5:6:
+node-id=3 locality-tiers=region=us-central-1,zone=us-central-1a,node=3
+  store-id=3 attrs=green locality-code=7:8:9:
+
+store-load-msg
+  store-id=1 node-id=1 load=[80,0,0] capacity=[100,100,100] secondary-load=0 load-time=0s
+----
+
+store-load-msg
+  store-id=2 node-id=2 load=[10,0,0] capacity=[100,100,100] secondary-load=0 load-time=0s
+----
+
+store-load-msg
+  store-id=3 node-id=3 load=[10,0,0] capacity=[100,100,100] secondary-load=0 load-time=0s
+----
+
+get-load-info
+----
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:0, byte-size:0] adjusted=[cpu:80, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=80 seq=1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:10, write-bandwidth:0, byte-size:0] adjusted=[cpu:10, write-bandwidth:0, byte-size:0] node-reported-cpu=10 node-adjusted-cpu=10 seq=1
+store-id=3 node-id=3 status=ok accepting all reported=[cpu:10, write-bandwidth:0, byte-size:0] adjusted=[cpu:10, write-bandwidth:0, byte-size:0] node-reported-cpu=10 node-adjusted-cpu=10 seq=1
+
+store-leaseholder-msg
+store-id=1
+  range-id=1 load=[60,0,0] raft-cpu=20 config=(num_replicas=3 constraints={} voter_constraints={})
+    store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
+    store-id=2 replica-id=2 type=VOTER_FULL
+    store-id=3 replica-id=3 type=VOTER_FULL
+----
+
+ranges
+----
+range-id=1 local-store=1 load=[cpu:60, write-bandwidth:0, byte-size:0] raft-cpu=20
+  store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
+  store-id=2 replica-id=2 type=VOTER_FULL
+  store-id=3 replica-id=3 type=VOTER_FULL
+
+get-load-info
+----
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:0, byte-size:0] adjusted=[cpu:80, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=80 seq=1
+  top-k-ranges (local-store-id=1) dim=CPURate: r1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:10, write-bandwidth:0, byte-size:0] adjusted=[cpu:10, write-bandwidth:0, byte-size:0] node-reported-cpu=10 node-adjusted-cpu=10 seq=1
+  top-k-ranges (local-store-id=1) dim=WriteBandwidth: r1
+store-id=3 node-id=3 status=ok accepting all reported=[cpu:10, write-bandwidth:0, byte-size:0] adjusted=[cpu:10, write-bandwidth:0, byte-size:0] node-reported-cpu=10 node-adjusted-cpu=10 seq=1
+  top-k-ranges (local-store-id=1) dim=WriteBandwidth: r1
+
+# TODO(tbg): maybe it'll be good to at least optionally be able to have the
+# trace printed in the output here.
+rebalance-stores store-id=1
+----
+pending(2)
+change-id=1 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-40, write-bandwidth:0, byte-size:0] start=0s
+  prev=(replica-id=1 type=VOTER_FULL leaseholder=true)
+  next=(replica-id=1 type=VOTER_FULL)
+change-id=2 store-id=2 node-id=2 range-id=1 load-delta=[cpu:44, write-bandwidth:0, byte-size:0] start=0s
+  prev=(replica-id=2 type=VOTER_FULL)
+  next=(replica-id=2 type=VOTER_FULL leaseholder=true)