Merge #158373

158373: mmaprototype: filter ineligible lease targets before computing means r=wenyihu6 a=tbg

Add health and store-level disposition checks to lease target filtering,
and move all eligibility filtering (including the existing per-replica
disposition check) to happen before computing load means.

Computing means only over eligible targets is semantically correct: when
asking "is store X above or below average?", we mean among stores that
could actually receive the lease. Including ineligible stores (dead,
unhealthy, refusing) distorts the average.

Example: with eligible targets at 30/40/50 QPS and a dead store at 0 QPS:
- Including dead store: mean=30, so 50 QPS looks "significantly above average"
- Excluding dead store: mean=40, so 50 QPS is "slightly above average"

The latter correctly reflects reality among viable candidates.

Add a datadriven directive and testdata file to exercise the new
retainReadyLeaseTargetStoresOnly function, covering health, store-level
disposition, and per-replica disposition filtering. Also extend the test
DSL to support lease-disposition on replica lines and allow set-store-status
to set individual status fields without triggering validation assertions.

- [ ] TODO: dig up the pre/post-means-ness of the SMA and highlight differences to reviewers.

Part of #156776.

Epic: CRDB-55052

Co-authored-by: Tobias Grieger <[email protected]>

Author: craig[bot]

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

@@ -112,6 +112,9 @@ func newRebalanceEnv(
 
 type sheddingStore struct {
 	roachpb.StoreID
+	// storeLoadSummary is relative to the entire cluster (not a set of valid
+	// replacement stores for some particular replica), see the comment where
+	// sheddingStores are constructed.
 	storeLoadSummary
 }
 
@@ -242,6 +245,49 @@ func (re *rebalanceEnv) rebalanceStores(
 	return re.changes
 }
 
+// rebalanceStore attempts to shed load from a single overloaded store via lease
+// transfers and/or replica rebalances.
+//
+// # Candidate Filtering Strategy
+//
+// When selecting rebalance targets, we filter candidates in two phases. The key
+// consideration is which stores should be included when computing load means,
+// since the means determine whether a store looks "underloaded" (good target)
+// or "overloaded" (bad target).
+//
+// **Pre-means filtering** excludes stores with a non-OK disposition from the
+// mean. The disposition serves as the source of truth for "can this store
+// accept work?" — whether due to drain, maintenance, disk capacity, or any
+// other reason. Note that unhealthy stores always have non-OK disposition, so
+// disposition indirectly reflects health.
+//
+// Ill-disposed stores cannot participate as targets and it is correct to
+// exclude them from the mean, as the mean helps us answer the question "among
+// viable targets, who is underloaded?". If ill-disposed storeswere included in
+// the mean, they could skew it down, making viable targets look too overloaded
+// to be considered, and preventing resolution of a load imbalance. They could
+// also skew the mean up and make actually-overloaded targets look underloaded,
+// but this is less of an issue: since both source and target are evaluated
+// relative to the same mean before accepting a transfer, their relative
+// positions are preserved, and we wouldn't accept a transfer that leaves the
+// target worse off than the source.
+//
+// **Post-means filtering** prevents some of the remaining candidates from being
+// chosen as targets for tactical reasons:
+//   - Stores that would be worse off than the source (relative to the mean
+//     across the candidate set), to reduce thrashing. For example, if a store
+//     were overloaded due to a single very hot range, moving that range to
+//     another store would just shift the problem elsewhere and cause thrashing.
+//   - Stores with too much pending inflight work (reduced confidence in load arithmetic).
+//
+// NB: TestClusterState tests various scenarios and edge cases that demonstrate
+// filtering behavior and outcomes.
+//
+// TODO(tbg): The above describes the intended design. Lease transfers follow this
+// pattern (see retainReadyLeaseTargetStoresOnly). Replica transfers are still
+// being cleaned up: they currently compute means over all constraint-matching
+// stores and filter only afterward, which is not intentional.
+// Tracking issue: https://github.com/cockroachdb/cockroach/pull/158373
 func (re *rebalanceEnv) rebalanceStore(
 	ctx context.Context, store sheddingStore, localStoreID roachpb.StoreID,
 ) {
@@ -595,6 +641,9 @@ func (re *rebalanceEnv) rebalanceLeasesFromLocalStoreID(
 		if len(candsPL) <= 1 {
 			continue // leaseholder is the only candidate
 		}
+
+		candsPL = retainReadyLeaseTargetStoresOnly(ctx, candsPL, re.stores, rangeID)
+
 		clear(re.scratch.nodes)
 		means := computeMeansForStoreSet(re, candsPL, re.scratch.nodes, re.scratch.stores)
 		sls := re.computeLoadSummary(ctx, store.StoreID, &means.storeLoad, &means.nodeLoad)
@@ -607,12 +656,6 @@ func (re *rebalanceEnv) rebalanceLeasesFromLocalStoreID(
 		}
 		var candsSet candidateSet
 		for _, cand := range cands {
-			if disp := re.stores[cand.storeID].adjusted.replicas[rangeID].LeaseDisposition; disp != LeaseDispositionOK {
-				// Don't transfer lease to a store that is lagging.
-				log.KvDistribution.Infof(ctx, "skipping store s%d for lease transfer: lease disposition %v",
-					cand.storeID, disp)
-				continue
-			}
 			candSls := re.computeLoadSummary(ctx, cand.storeID, &means.storeLoad, &means.nodeLoad)
 			candsSet.candidates = append(candsSet.candidates, candidateInfo{
 				StoreID:              cand.storeID,
@@ -685,3 +728,26 @@ func (re *rebalanceEnv) rebalanceLeasesFromLocalStoreID(
 	// We iterated through all top-K ranges without running into any limits.
 	return leaseTransferCount
 }
+
+// retainReadyLeaseTargetStoresOnly filters the input set to only those stores that
+// are ready to accept a lease for the given range. A store is not ready if it
+// is not healthy, or does not accept leases at either the store or replica level.
+//
+// The input storeSet is mutated (and used to for the returned result).
+func retainReadyLeaseTargetStoresOnly(
+	ctx context.Context, in storeSet, stores map[roachpb.StoreID]*storeState, rangeID roachpb.RangeID,
+) storeSet {
+	out := in[:0]
+	for _, storeID := range in {
+		s := stores[storeID].status
+		switch {
+		case s.Disposition.Lease != LeaseDispositionOK:
+			log.KvDistribution.VEventf(ctx, 2, "skipping s%d for lease transfer: lease disposition %v (health %v)", storeID, s.Disposition.Lease, s.Health)
+		case stores[storeID].adjusted.replicas[rangeID].LeaseDisposition != LeaseDispositionOK:
+			log.KvDistribution.VEventf(ctx, 2, "skipping s%d for lease transfer: replica lease disposition %v (health %v)", storeID, stores[storeID].adjusted.replicas[rangeID].LeaseDisposition, s.Health)
+		default:
+			out = append(out, storeID)
+		}
+	}
+	return out
+}

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

@@ -74,8 +74,7 @@ func parseSecondaryLoadVector(t *testing.T, in string) SecondaryLoadVector {
 	return vec
 }
 
-func parseStatusFromArgs(t *testing.T, d *datadriven.TestData) Status {
-	var status Status
+func parseStatusFromArgs(t *testing.T, d *datadriven.TestData, status *Status) {
 	if d.HasArg("health") {
 		healthStr := dd.ScanArg[string](t, d, "health")
 		found := false
@@ -118,7 +117,6 @@ func parseStatusFromArgs(t *testing.T, d *datadriven.TestData) Status {
 			t.Fatalf("unknown replica disposition: %s", replicaStr)
 		}
 	}
-	return status
 }
 
 func parseStoreLoadMsg(t *testing.T, in string) StoreLoadMsg {
@@ -215,6 +213,18 @@ func parseStoreLeaseholderMsg(t *testing.T, in string) StoreLeaseholderMsg {
 					replType, err := parseReplicaType(parts[1])
 					require.NoError(t, err)
 					repl.ReplicaType.ReplicaType = replType
+				case "lease-disposition":
+					found := false
+					for i := LeaseDisposition(0); i < leaseDispositionCount; i++ {
+						if i.String() == parts[1] {
+							repl.LeaseDisposition = i
+							found = true
+							break
+						}
+					}
+					if !found {
+						t.Fatalf("unknown lease disposition: %s", parts[1])
+					}
 				default:
 					t.Fatalf("unknown argument: %s", parts[0])
 				}
@@ -345,6 +355,9 @@ func TestClusterState(t *testing.T) {
 		func(t *testing.T, path string) {
 			ts := timeutil.NewManualTime(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC))
 			cs := newClusterState(ts, newStringInterner())
+			tr := tracing.NewTracer()
+			tr.SetRedactable(true)
+			defer tr.Close()
 
 			printNodeListMeta := func(t *testing.T) string {
 				nodeList := []int{}
@@ -401,6 +414,11 @@ func TestClusterState(t *testing.T) {
 			// Recursively invoked in `include` directive.
 			var invokeFn func(t *testing.T, d *datadriven.TestData) string
 			invokeFn = func(t *testing.T, d *datadriven.TestData) string {
+				// Start a recording span for each command. Commands that want to
+				// include the trace in their output can call finishAndGet().
+				ctx, finishAndGet := tracing.ContextWithRecordingSpan(
+					context.Background(), tr, d.Cmd,
+				)
 				switch d.Cmd {
 				case "include":
 					loc := dd.ScanArg[string](t, d, "path")
@@ -477,8 +495,11 @@ func TestClusterState(t *testing.T) {
 					if !ok {
 						t.Fatalf("store %d not found", storeID)
 					}
-					status := parseStatusFromArgs(t, d)
-					ss.status = MakeStatus(status.Health, status.Disposition.Lease, status.Disposition.Replica)
+					// NB: we intentionall bypass the assertion in MakeStatus
+					// here so that we can test all combinations of health,
+					// lease, and replica dispositions, even those that we never
+					// want to see in production.
+					parseStatusFromArgs(t, d, &ss.status)
 					return ss.status.String()
 
 				case "store-load-msg":
@@ -576,10 +597,6 @@ func TestClusterState(t *testing.T) {
 					storeID := dd.ScanArg[roachpb.StoreID](t, d, "store-id")
 					rng := rand.New(rand.NewSource(0))
 					dsm := newDiversityScoringMemo()
-					tr := tracing.NewTracer()
-					tr.SetRedactable(true)
-					defer tr.Close()
-					ctx, finishAndGet := tracing.ContextWithRecordingSpan(context.Background(), tr, "rebalance-stores")
 					re := newRebalanceEnv(cs, rng, dsm, cs.ts.Now())
 
 					if n, ok := dd.ScanArgOpt[int](t, d, "max-lease-transfer-count"); ok {
@@ -603,6 +620,15 @@ func TestClusterState(t *testing.T) {
 					ts.Advance(time.Second * time.Duration(seconds))
 					return fmt.Sprintf("t=%v", ts.Now().Sub(testingBaseTime))
 
+				case "retain-ready-lease-target-stores-only":
+					in := dd.ScanArg[[]roachpb.StoreID](t, d, "in")
+					rangeID := dd.ScanArg[roachpb.RangeID](t, d, "range-id")
+					out := retainReadyLeaseTargetStoresOnly(ctx, storeSet(in), cs.stores, rangeID)
+					rec := finishAndGet()
+					var sb redact.StringBuilder
+					rec.SafeFormatMinimal(&sb)
+					return fmt.Sprintf("%s%v\n", sb.String(), out)
+
 				default:
 					panic(fmt.Sprintf("unknown command: %v", d.Cmd))
 				}

pkg/kv/kvserver/allocator/mmaprototype/testdata/cluster_state/lease_disposition.txt

@@ -0,0 +1,139 @@
+# Test retainReadyLeaseTargetStoresOnly which filters stores based on:
+# 1. Store-level lease disposition (must be LeaseDispositionOK)
+# 2. Per-replica lease disposition (must be LeaseDispositionOK)
+#
+# Note: Health is handled indirectly via disposition. Unhealthy stores always
+# have non-OK disposition, so we don't need a separate health check.
+
+set-store
+  store-id=1 node-id=1 locality-tiers=region=us
+  store-id=2 node-id=2 locality-tiers=region=us
+  store-id=3 node-id=3 locality-tiers=region=us
+----
+node-id=1 locality-tiers=region=us,node=1
+  store-id=1 attrs=
+node-id=2 locality-tiers=region=us,node=2
+  store-id=2 attrs=
+node-id=3 locality-tiers=region=us,node=3
+  store-id=3 attrs=
+
+# Set up some stores. The specifics don't matter since we're only testing
+# the lease health checks.
+store-load-msg
+  store-id=1 node-id=1 load=[100,0,0] capacity=[200,100,100] load-time=0s
+  store-id=2 node-id=2 load=[50,0,0] capacity=[200,100,100] load-time=0s
+  store-id=3 node-id=3 load=[50,0,0] capacity=[200,100,100] load-time=0s
+----
+
+store-leaseholder-msg
+store-id=1
+  range-id=1 load=[10,0,0]
+  config=(num_replicas=3)
+    store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
+    store-id=2 replica-id=2 type=VOTER_FULL leaseholder=false
+    store-id=3 replica-id=3 type=VOTER_FULL leaseholder=false
+----
+
+# All stores healthy and accepting leases - all retained.
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+[1 2 3]
+
+# Only input stores matter.
+retain-ready-lease-target-stores-only in=(1,3) range-id=1
+----
+[1 3]
+
+# Mark s2 as shedding replicas, which should have no effect
+# since we're looking at leases.
+set-store-status store-id=2 replicas=shedding
+----
+ok shedding=replicas
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+[1 2 3]
+
+# Make store 2 unhealthy - it should be filtered out via disposition.
+# Unhealthy stores must have non-OK disposition (invariant).
+set-store-status store-id=2 health=unhealthy leases=shedding replicas=shedding
+----
+unhealthy shedding=leases,replicas
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+skipping s2 for lease transfer: lease disposition shedding (health unhealthy)
+[1 3]
+
+# Different kind of unhealthy. Same result.
+set-store-status store-id=2 health=unknown leases=shedding replicas=shedding
+----
+unknown shedding=leases,replicas
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+skipping s2 for lease transfer: lease disposition shedding (health unknown)
+[1 3]
+
+# Restore store 2, make store 3 refuse leases at store level.
+set-store-status store-id=2 health=ok leases=ok replicas=ok
+----
+ok accepting all
+
+set-store-status store-id=3 leases=refusing
+----
+ok refusing=leases
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+skipping s3 for lease transfer: lease disposition refusing (health ok)
+[1 2]
+
+# Shedding and refusing are treated the same.
+set-store-status store-id=3 health=ok leases=shedding
+----
+ok shedding=leases
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+skipping s3 for lease transfer: lease disposition shedding (health ok)
+[1 2]
+
+# Restore store 3.
+set-store-status store-id=3 leases=ok
+----
+ok accepting all
+
+# All stores ready again.
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+[1 2 3]
+
+# Mark r1 as not accepting leases targeting s2 via per-replica disposition.
+store-leaseholder-msg
+store-id=1
+  range-id=1 load=[10,0,0]
+  config=(num_replicas=3)
+    store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
+    store-id=2 replica-id=2 type=VOTER_FULL leaseholder=false lease-disposition=refusing
+    store-id=3 replica-id=3 type=VOTER_FULL leaseholder=false
+----
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+skipping s2 for lease transfer: replica lease disposition refusing (health ok)
+[1 3]
+
+# Restore s2's replica disposition.
+store-leaseholder-msg
+store-id=1
+  range-id=1 load=[10,0,0]
+  config=(num_replicas=3)
+    store-id=1 replica-id=1 type=VOTER_FULL leaseholder=true
+    store-id=2 replica-id=2 type=VOTER_FULL leaseholder=false lease-disposition=ok
+    store-id=3 replica-id=3 type=VOTER_FULL leaseholder=false
+----
+
+retain-ready-lease-target-stores-only in=(1,2,3) range-id=1
+----
+[1 2 3]