Merge cockroachdb#157818 cockroachdb#157827

157818: mmaprototype: extract and unit test maxFractionPendingIncDec r=wenyihu6 a=tbg

Touches cockroachdb#157757.

Epic: CRDB-55052

157827: kvserver: deflake TestReplicateRestartAfterTruncationWithRemoveAndReAdd r=wenyihu6 a=tbg

This test could spend >10s with non-live nodes (at least under stress), so it
would trip the timeout in a check in `executeChangeReplicasTxn` that tries to
ensure that we don't replicate into an unavailable configuration.

Instead of skipping these tests under stress, we enable the testing knob that
disables the check.

This passed 2k iters of

```
./dev test pkg/kv/kvserver -f TestReplicateRestartAfterTruncationWithRemoveAndReAdd --stress
```

which before (the one time I tried before fixing) failed after maybe 25% of that.

Closes cockroachdb#157203.
Closes cockroachdb#156689.

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

Author: craig[bot]

pkg/kv/kvserver/allocator/mmaprototype/BUILD.bazel

@@ -39,6 +39,7 @@ go_test(
     name = "mmaprototype_test",
     srcs = [
         "allocator_state_test.go",
+        "cluster_state_rebalance_stores_test.go",
         "cluster_state_test.go",
         "constraint_matcher_test.go",
         "constraint_test.go",

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

@@ -935,32 +935,49 @@ func (ss *storeState) computePendingChangesReflectedInLatestLoad(
 }
 
 func (ss *storeState) computeMaxFractionPending() {
-	fracIncrease := 0.0
-	fracDecrease := 0.0
-	for i := range ss.reportedLoad {
-		if ss.reportedLoad[i] == ss.adjusted.load[i] && ss.reportedLoad[i] == 0 {
-			// Avoid setting ss.maxFractionPendingIncrease and
-			// ss.maxFractionPendingDecrease to 1000 when the reported load and
-			// adjusted load are both 0 since some dimension is expected to have zero
-			// (e.g. write bandwidth during read-only workloads).
-			continue
-		}
-		if ss.reportedLoad[i] == 0 {
-			fracIncrease = 1000
-			fracDecrease = 1000
-			break
-		}
-		f := math.Abs(float64(ss.adjusted.load[i]-ss.reportedLoad[i])) / float64(ss.reportedLoad[i])
-		if ss.adjusted.load[i] > ss.reportedLoad[i] {
-			if f > fracIncrease {
-				fracIncrease = f
+	ss.maxFractionPendingIncrease, ss.maxFractionPendingDecrease = computeMaxFractionPendingIncDec(ss.reportedLoad, ss.adjusted.load)
+}
+
+func computeMaxFractionPendingIncDec(rep, adj LoadVector) (maxFracInc, maxFracDec float64) {
+	for i := range rep {
+		inc, dec := func(rep, adj LoadValue) (inc, dec float64) {
+			// The fraction pending expresses the absolute difference of the adjusted
+			// and reported load as a multiple of the reported load. Note that this
+			// is the case even if the adjusted load is negative: if, say, the adjusted
+			// load is -50 and the reported load is 100, it is still correct to say that
+			// a "magnitude 1.5x" change is pending (from 100 to -50).
+			diff := adj - rep
+
+			switch {
+			case diff == 0:
+				// Reported and adjusted are equal, so nothing is pending.
+				// This also handles the case in which both are zero.
+				// We don't need to update maxFracInc or maxFracDec because
+				// they started at zero and only go up from there.
+				return 0, 0
+			case rep == 0:
+				// The adjusted load is nonzero, but the reported one is zero. We can't
+				// express the load change as a multiple of zero. Arbitrarily assign large
+				// value to both increase and decrease, indicating that no more changes
+				// should be made until either the pending change clears (and we get a
+				// zero diff above) or we register positive reported load.
+				return 1000, 1000
+			case diff > 0:
+				// Vanilla case of adjusted > reported, i.e. we have load incoming.
+				// We don't need to update maxFracDec.
+				return math.Abs(float64(diff) / float64(rep)), 0
+			case diff < 0:
+				// Vanilla case of adjusted < reported, i.e. we have load incoming.
+				// We don't need to update maxFracInc.
+				return 0, math.Abs(float64(diff) / float64(rep))
+			default:
+				panic("impossible")
 			}
-		} else if f > fracDecrease {
-			fracDecrease = f
-		}
+		}(rep[i], adj[i])
+		maxFracInc = max(maxFracInc, inc)
+		maxFracDec = max(maxFracDec, dec)
 	}
-	ss.maxFractionPendingIncrease = fracIncrease
-	ss.maxFractionPendingDecrease = fracDecrease
+	return maxFracInc, maxFracDec
 }
 
 func newStoreState() *storeState {

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

@@ -0,0 +1,115 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package mmaprototype
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestComputeMaxFractionPendingIncDec(t *testing.T) {
+	tests := []struct {
+		name           string
+		reported       LoadVector
+		adjusted       LoadVector
+		expectedMaxInc float64
+		expectedMaxDec float64
+	}{
+		{
+			name:           "both_zero",
+			reported:       LoadVector{},
+			adjusted:       LoadVector{},
+			expectedMaxInc: 0,
+			expectedMaxDec: 0,
+		},
+		{
+			name:           "reported_zero_adjusted_nonzero",
+			reported:       LoadVector{},
+			adjusted:       LoadVector{CPURate: 100, WriteBandwidth: 200, ByteSize: 300},
+			expectedMaxInc: 1000,
+			expectedMaxDec: 1000,
+		},
+		{
+			name:           "reported_nonzero_adjusted_zero",
+			reported:       LoadVector{CPURate: 100, WriteBandwidth: 200, ByteSize: 300},
+			adjusted:       LoadVector{CPURate: 0, WriteBandwidth: 0, ByteSize: 0},
+			expectedMaxInc: 0,
+			expectedMaxDec: 1.0, // abs((0-100)/100) = 1.0, same for all dimensions
+		},
+		{
+			name:           "simple_increase",
+			reported:       LoadVector{CPURate: 100},
+			adjusted:       LoadVector{CPURate: 150},
+			expectedMaxInc: 0.5,
+			expectedMaxDec: 0,
+		},
+		{
+			name:           "simple_decrease",
+			reported:       LoadVector{CPURate: 100},
+			adjusted:       LoadVector{CPURate: 50},
+			expectedMaxInc: 0,
+			expectedMaxDec: 0.5,
+		},
+		{
+			name:           "multiple_dimensions_increase",
+			reported:       LoadVector{CPURate: 100, WriteBandwidth: 200, ByteSize: 300},
+			adjusted:       LoadVector{CPURate: 150, WriteBandwidth: 250, ByteSize: 400},
+			expectedMaxInc: 0.5, // max(0.5, 0.25, 0.333...)
+			expectedMaxDec: 0,
+		},
+		{
+			name:           "multiple_dimensions_decrease",
+			reported:       LoadVector{CPURate: 100, WriteBandwidth: 200, ByteSize: 300},
+			adjusted:       LoadVector{CPURate: 50, WriteBandwidth: 100, ByteSize: 150},
+			expectedMaxInc: 0,
+			expectedMaxDec: 0.5, // max(0.5, 0.5, 0.5)
+		},
+		{
+			name:           "mixed_increase_and_decrease",
+			reported:       LoadVector{CPURate: 100, WriteBandwidth: 200, ByteSize: 300},
+			adjusted:       LoadVector{CPURate: 150, WriteBandwidth: 100, ByteSize: 300},
+			expectedMaxInc: 0.5, // from CPURate
+			expectedMaxDec: 0.5, // from WriteBandwidth
+		},
+		{
+			name:           "negative_adjusted_load",
+			reported:       LoadVector{CPURate: 100},
+			adjusted:       LoadVector{CPURate: -50},
+			expectedMaxInc: 0,
+			expectedMaxDec: 1.5, // abs(-50 - 100) / 100 = 1.5
+		},
+		{
+			name:           "fractional_change",
+			reported:       LoadVector{CPURate: 1000},
+			adjusted:       LoadVector{CPURate: 1100},
+			expectedMaxInc: 0.1,
+			expectedMaxDec: 0,
+		},
+		{
+			name:           "max_across_dimensions",
+			reported:       LoadVector{CPURate: 100, WriteBandwidth: 200, ByteSize: 300},
+			adjusted:       LoadVector{CPURate: 120, WriteBandwidth: 300, ByteSize: 450},
+			expectedMaxInc: 0.5, // max(0.2, 0.5, 0.5)
+			expectedMaxDec: 0,
+		},
+		{
+			name:           "one_dimension_zero_others_change",
+			reported:       LoadVector{CPURate: 0, WriteBandwidth: 100, ByteSize: 200},
+			adjusted:       LoadVector{CPURate: 50, WriteBandwidth: 150, ByteSize: 100},
+			expectedMaxInc: 1000, // from CPURate (reported=0)
+			expectedMaxDec: 1000, // from CPURate (reported=0)
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			maxInc, maxDec := computeMaxFractionPendingIncDec(tt.reported, tt.adjusted)
+			require.InDelta(t, tt.expectedMaxInc, maxInc, 1e-9, "max fraction increase")
+			require.InDelta(t, tt.expectedMaxDec, maxDec, 1e-9, "max fraction decrease")
+		})
+	}
+}

pkg/kv/kvserver/client_raft_test.go

@@ -36,7 +36,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/logstore"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/rditer"
-	raft "github.com/cockroachdb/cockroach/pkg/raft"
+	"github.com/cockroachdb/cockroach/pkg/raft"
 	"github.com/cockroachdb/cockroach/pkg/raft/raftpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
@@ -2294,6 +2294,12 @@ func runReplicateRestartAfterTruncation(t *testing.T, removeBeforeTruncateAndReA
 					StickyVFSRegistry: fs.NewStickyRegistry(),
 					WallClock:         manualClock,
 				},
+				// In this test, under duress, we can occasionally fail enough
+				// heartbeats to exceed the 10s timeout in execChangeReplicasTxn to
+				// verify liveness of a post-change quorum.
+				//
+				// See: https://github.com/cockroachdb/cockroach/issues/156689
+				Store: &kvserver.StoreTestingKnobs{AllowDangerousReplicationChanges: true},
 			},
 		}
 	}