mmaprototype: add Status to cluster state test

This allows TestClusterState to update the store status.
It continues to be true that no semantics are attached to the store status at
this point.

Author: tbg

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

@@ -63,6 +63,53 @@ func parseSecondaryLoadVector(t *testing.T, in string) SecondaryLoadVector {
 	return vec
 }
 
+func parseStatusFromArgs(t *testing.T, d *datadriven.TestData) Status {
+	var status Status
+	if d.HasArg("health") {
+		healthStr := dd.ScanArg[string](t, d, "health")
+		found := false
+		for i := Health(0); i < healthCount; i++ {
+			if i.String() == healthStr {
+				status.Health = i
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Fatalf("unknown health: %s", healthStr)
+		}
+	}
+	if d.HasArg("leases") {
+		leaseStr := dd.ScanArg[string](t, d, "leases")
+		found := false
+		for i := LeaseDisposition(0); i < leaseDispositionCount; i++ {
+			if i.String() == leaseStr {
+				status.Disposition.Lease = i
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Fatalf("unknown lease disposition: %s", leaseStr)
+		}
+	}
+	if d.HasArg("replicas") {
+		replicaStr := dd.ScanArg[string](t, d, "replicas")
+		found := false
+		for i := ReplicaDisposition(0); i < replicaDispositionCount; i++ {
+			if i.String() == replicaStr {
+				status.Disposition.Replica = i
+				found = true
+				break
+			}
+		}
+		if !found {
+			t.Fatalf("unknown replica disposition: %s", replicaStr)
+		}
+	}
+	return status
+}
+
 func parseStoreLoadMsg(t *testing.T, in string) StoreLoadMsg {
 	var msg StoreLoadMsg
 	for _, v := range strings.Fields(in) {
@@ -333,8 +380,10 @@ func TestClusterState(t *testing.T) {
 						ss := cs.stores[storeID]
 						ns := cs.nodes[ss.NodeID]
 						fmt.Fprintf(&buf,
-							"store-id=%v node-id=%v reported=%v adjusted=%v node-reported-cpu=%v node-adjusted-cpu=%v seq=%d\n",
-							ss.StoreID, ss.NodeID, ss.reportedLoad, ss.adjusted.load, ns.ReportedCPU, ns.adjustedCPU, ss.loadSeqNum,
+							"store-id=%v node-id=%v status=%s reported=%v adjusted=%v node-reported-cpu=%v node-adjusted-cpu=%v seq"+
+								"=%d\n",
+							ss.StoreID, ss.NodeID, ss.status, ss.reportedLoad, ss.adjusted.load, ns.ReportedCPU, ns.adjustedCPU,
+							ss.loadSeqNum,
 						)
 						for ls, topk := range ss.adjusted.topKRanges {
 							n := topk.len()
@@ -354,9 +403,22 @@ func TestClusterState(t *testing.T) {
 					for _, next := range strings.Split(d.Input, "\n") {
 						sal := parseStoreAttributedAndLocality(t, next)
 						cs.setStore(sal)
+						// For convenience, in these tests, stores start out
+						// healthy.
+						cs.stores[sal.StoreID].status = Status{Health: HealthOK}
 					}
 					return printNodeListMeta()
 
+				case "set-store-status":
+					storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
+					ss, ok := cs.stores[storeID]
+					if !ok {
+						t.Fatalf("store %d not found", storeID)
+					}
+					status := parseStatusFromArgs(t, d)
+					ss.status = MakeStatus(status.Health, status.Disposition.Lease, status.Disposition.Replica)
+					return ss.status.String()
+
 				case "store-load-msg":
 					msg := parseStoreLoadMsg(t, d.Input)
 					cs.processStoreLoadMsg(context.Background(), &msg)

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

@@ -36,9 +36,9 @@ store-id=1
 # The top-k uses CPURate.
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:100, write-bandwidth:40, byte-size:50] adjusted=[cpu:100, write-bandwidth:40, byte-size:50] node-reported-cpu=100 node-adjusted-cpu=100 seq=1
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:100, write-bandwidth:40, byte-size:50] adjusted=[cpu:100, write-bandwidth:40, byte-size:50] node-reported-cpu=100 node-adjusted-cpu=100 seq=1
   top-k-ranges (local-store-id=1) dim=CPURate: r4 r3
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:45, write-bandwidth:40, byte-size:50] adjusted=[cpu:45, write-bandwidth:40, byte-size:50] node-reported-cpu=45 node-adjusted-cpu=45 seq=1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:45, write-bandwidth:40, byte-size:50] adjusted=[cpu:45, write-bandwidth:40, byte-size:50] node-reported-cpu=45 node-adjusted-cpu=45 seq=1
   top-k-ranges (local-store-id=1) dim=CPURate: r3 r2
 
 
@@ -67,9 +67,9 @@ store-id=1
 # The top-2 ranges for store 2 are based on ByteSize.
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:100, write-bandwidth:40, byte-size:50] adjusted=[cpu:100, write-bandwidth:40, byte-size:50] node-reported-cpu=100 node-adjusted-cpu=100 seq=1
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:100, write-bandwidth:40, byte-size:50] adjusted=[cpu:100, write-bandwidth:40, byte-size:50] node-reported-cpu=100 node-adjusted-cpu=100 seq=1
   top-k-ranges (local-store-id=1) dim=CPURate: r4 r3
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:45, write-bandwidth:40, byte-size:50] adjusted=[cpu:45, write-bandwidth:40, byte-size:50] node-reported-cpu=45 node-adjusted-cpu=45 seq=2
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:45, write-bandwidth:40, byte-size:50] adjusted=[cpu:45, write-bandwidth:40, byte-size:50] node-reported-cpu=45 node-adjusted-cpu=45 seq=2
   top-k-ranges (local-store-id=1) dim=ByteSize: r1 r2
 
 # StoreLeaseholderMsg not containing r1 and r4 since no longer the leaseholder.
@@ -86,7 +86,7 @@ store-id=1
 # r1 and r4 no longer mentioned in the top-k.
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:100, write-bandwidth:40, byte-size:50] adjusted=[cpu:100, write-bandwidth:40, byte-size:50] node-reported-cpu=100 node-adjusted-cpu=100 seq=1
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:100, write-bandwidth:40, byte-size:50] adjusted=[cpu:100, write-bandwidth:40, byte-size:50] node-reported-cpu=100 node-adjusted-cpu=100 seq=1
   top-k-ranges (local-store-id=1) dim=CPURate: r3 r2
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:45, write-bandwidth:40, byte-size:50] adjusted=[cpu:45, write-bandwidth:40, byte-size:50] node-reported-cpu=45 node-adjusted-cpu=45 seq=2
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:45, write-bandwidth:40, byte-size:50] adjusted=[cpu:45, write-bandwidth:40, byte-size:50] node-reported-cpu=45 node-adjusted-cpu=45 seq=2
   top-k-ranges (local-store-id=1) dim=ByteSize: r2 r3

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

@@ -13,10 +13,10 @@ store-load-msg
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
 
-store-leaseholder-msg 
+store-leaseholder-msg
 store-id=1
   range-id=1 load=[80,80,80] raft-cpu=20 config=(num_replicas=3 constraints={'+region=us-west-1:1'} voter_constraints={'+region=us-west-1:1'})
     store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
@@ -30,9 +30,9 @@ range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=1
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] 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=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
 
 make-pending-changes range-id=1
   rebalance-replica: remove-store-id=1 add-store-id=2
@@ -52,9 +52,9 @@ range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=0 seq=2
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=0 seq=2
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:88, write-bandwidth:88, byte-size:88] node-reported-cpu=0 node-adjusted-cpu=88 seq=1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:88, write-bandwidth:88, byte-size:88] node-reported-cpu=0 node-adjusted-cpu=88 seq=1
 
 # TODO: this test is confusing since it is sending StoreLeaseholderMsgs from
 # s1 and s2 to the same clusterState, but s1 and s2 are not on the same node.
@@ -66,9 +66,9 @@ store-id=2
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=0 seq=2
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=0 seq=2
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:88, write-bandwidth:88, byte-size:88] node-reported-cpu=0 node-adjusted-cpu=88 seq=1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:88, write-bandwidth:88, byte-size:88] node-reported-cpu=0 node-adjusted-cpu=88 seq=1
   top-k-ranges (local-store-id=2) dim=ByteSize: r1
 
 get-pending-changes
@@ -94,9 +94,9 @@ change-id=2 store-id=1 node-id=1 range-id=1 load-delta=[cpu:-80, write-bandwidth
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=0 seq=2
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=80 node-adjusted-cpu=0 seq=2
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=2
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=2
   top-k-ranges (local-store-id=2) dim=ByteSize: r1
 
 
@@ -110,7 +110,7 @@ pending(0)
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=3
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=3
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=2
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:80, write-bandwidth:80, byte-size:80] adjusted=[cpu:80, write-bandwidth:80, byte-size:80] node-reported-cpu=80 node-adjusted-cpu=80 seq=2
   top-k-ranges (local-store-id=2) dim=ByteSize: r1

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

@@ -13,8 +13,8 @@ store-load-msg
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:20, write-bandwidth:80, byte-size:80] node-reported-cpu=20 node-adjusted-cpu=20 seq=1
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:20, write-bandwidth:80, byte-size:80] node-reported-cpu=20 node-adjusted-cpu=20 seq=1
 
 store-leaseholder-msg 
 store-id=1
@@ -25,9 +25,9 @@ store-id=1
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:20, write-bandwidth:80, byte-size:80] node-reported-cpu=20 node-adjusted-cpu=20 seq=1
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:20, write-bandwidth:80, byte-size:80] node-reported-cpu=20 node-adjusted-cpu=20 seq=1
   top-k-ranges (local-store-id=1) dim=CPURate: r1
 
 ranges
@@ -55,9 +55,9 @@ range-id=1 load=[cpu:80, write-bandwidth:80, byte-size:80] raft-cpu=20
 # applied.
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=20 node-adjusted-cpu=0 seq=2
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=20 node-adjusted-cpu=0 seq=2
   top-k-ranges (local-store-id=1) dim=CPURate: r1
 
 store-leaseholder-msg
@@ -68,9 +68,9 @@ store-id=1
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=20 node-adjusted-cpu=0 seq=2
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:20, write-bandwidth:80, byte-size:80] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=20 node-adjusted-cpu=0 seq=2
 
 get-pending-changes
 ----
@@ -85,9 +85,9 @@ store-load-msg
 
 get-load-info
 ----
-store-id=1 node-id=1 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+store-id=1 node-id=1 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
   top-k-ranges (local-store-id=1) dim=CPURate: r1
-store-id=2 node-id=2 status=unknown, refusing=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=3
+store-id=2 node-id=2 status=ok accepting all reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=3
 
 get-pending-changes
 ----

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

@@ -0,0 +1,21 @@
+set-store
+  store-id=1 node-id=1 attrs=purple locality-tiers=region=foo
+----
+node-id=1 locality-tiers=region=foo,node=1
+  store-id=1 attrs=purple locality-code=1:2:
+
+set-store-status store-id=1 health=unhealthy replicas=refusing leases=shedding
+----
+unhealthy refusing=replicas shedding=leases
+
+set-store-status store-id=1 health=unhealthy replicas=shedding leases=shedding
+----
+unhealthy shedding=leases,replicas
+
+get-load-info
+----
+store-id=1 node-id=1 status=unhealthy shedding=leases,replicas reported=[cpu:0, write-bandwidth:0, byte-size:0] adjusted=[cpu:0, write-bandwidth:0, byte-size:0] node-reported-cpu=0 node-adjusted-cpu=0 seq=0
+
+set-store-status store-id=1 health=ok
+----
+ok accepting all