Merge #144106

144106: kvserver: add mixed version race condition test for policy refresher r=arulajmani a=wenyihu6


**kvserver: add PolicyRefresher.TestingKnobs**

This commit adds PolicyRefresher.TestingKnobs, which enables testing race
scenarios where the sender hasn't observed a cluster version upgrade when
creating the policy map, but the PolicyRefresher later does.

Epic: none
Release note: none

---

**kvserver: add mixed version race condition test for policy refresher**

This commit adds TestReplicaClosedTSPolicyWithPolicyRefresherInMixedVersionCluster.

TestReplicaClosedTSPolicyWithPolicyRefresherInMixedVersionCluster verifies
that the closed timestamp policy refresher behaves correctly in a
mixed-version cluster.

Particularly, a race condition like below might occur:
1. Side transport prepares a policy map before cluster upgrade is complete.
2. Cluster upgrade completes.
3. Policy refresher sees the upgrade and quickly updates replica policies to
use latency-based policies.
4. Replica tries to use a latency-based policy but the policy map from step 1
doesn't include it yet.

The logic in replica.getTargetByPolicy handles this race condition by falling
back to no-latency based policies if no-latency based policies were included
from the map provided by the side transport sender.

This test simulates a race condition by using a testing knob to allow the
policy refresher to use latency based policies on replicas while the rest of
the system is still on an older version. It verifies that even if the policy
refresher sees an upgrade to V25_2 and replicas starts holding latency based
policies, replicas will correctly fall back to non-latency-based policies if
the sender hasn’t yet sent the updated latency-based policies.

Part of: #143888
Release note: none



Co-authored-by: wenyihu6 <[email protected]>

Author: craig[bot]

pkg/base/testing_knobs.go

@@ -57,4 +57,5 @@ type TestingKnobs struct {
 	TableMetadata                  ModuleTestingKnobs
 	LicenseTestingKnobs            ModuleTestingKnobs
 	VecIndexTestingKnobs           ModuleTestingKnobs
+	PolicyRefresherTestingKnobs    ModuleTestingKnobs
 }

pkg/kv/kvserver/BUILD.bazel

@@ -429,6 +429,7 @@ go_test(
         "//pkg/kv/kvserver/benignerror",
         "//pkg/kv/kvserver/closedts",
         "//pkg/kv/kvserver/closedts/ctpb",
+        "//pkg/kv/kvserver/closedts/policyrefresher",
         "//pkg/kv/kvserver/closedts/tracker",
         "//pkg/kv/kvserver/concurrency",
         "//pkg/kv/kvserver/concurrency/isolation",

pkg/kv/kvserver/closedts/policyrefresher/BUILD.bazel

@@ -2,7 +2,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "policyrefresher",
-    srcs = ["policy_refresher.go"],
+    srcs = [
+        "config.go",
+        "policy_refresher.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/policyrefresher",
     visibility = ["//visibility:public"],
     deps = [

pkg/kv/kvserver/closedts/policyrefresher/config.go

@@ -0,0 +1,24 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package policyrefresher
+
+import (
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
+)
+
+// TestingKnobs contains testing knobs for the policy refresher.
+type TestingKnobs struct {
+	// InjectedLatencies returns a callback that returns a map of node IDs to
+	// latencies that will be used by the policy refresher instead of observed
+	// latencies from rpc context without checking for cluster settings or
+	// version compatibility.
+	InjectedLatencies func() map[roachpb.NodeID]time.Duration
+}
+
+// TestingKnobs implements the ModuleTestingKnobs interface.
+func (*TestingKnobs) ModuleTestingKnobs() {}

pkg/kv/kvserver/closedts/policyrefresher/policy_refresher.go

@@ -29,6 +29,7 @@ import (
 type PolicyRefresher struct {
 	stopper  *stop.Stopper
 	settings *cluster.Settings
+	knobs    *TestingKnobs
 
 	// getLeaseholderReplicas returns the set of replicas that are currently
 	// leaseholders of the node.
@@ -64,6 +65,7 @@ func NewPolicyRefresher(
 	settings *cluster.Settings,
 	getLeaseholderReplicas func() []Replica,
 	getNodeLatencies func() map[roachpb.NodeID]time.Duration,
+	knobs *TestingKnobs,
 ) *PolicyRefresher {
 	if getLeaseholderReplicas == nil || getNodeLatencies == nil {
 		log.Fatalf(context.Background(), "getLeaseholderReplicas and getNodeLatencies must be non-nil")
@@ -72,6 +74,7 @@ func NewPolicyRefresher(
 	refresher := &PolicyRefresher{
 		stopper:                stopper,
 		settings:               settings,
+		knobs:                  knobs,
 		getLeaseholderReplicas: getLeaseholderReplicas,
 		getNodeLatencies:       getNodeLatencies,
 		refreshNotificationCh:  make(chan struct{}, 1),
@@ -127,6 +130,10 @@ func (pr *PolicyRefresher) updateLatencyCache() {
 // getCurrentLatencies returns the current latency information if auto-tuning is
 // enabled and the cluster has been fully upgraded to v25.2, or nil otherwise.
 func (pr *PolicyRefresher) getCurrentLatencies() map[roachpb.NodeID]time.Duration {
+	// Testing knobs only.
+	if pr.knobs != nil && pr.knobs.InjectedLatencies != nil {
+		return pr.knobs.InjectedLatencies()
+	}
 	if !closedts.LeadForGlobalReadsAutoTuneEnabled.Get(&pr.settings.SV) || !pr.settings.Version.IsActive(context.TODO(), clusterversion.V25_2) {
 		return nil
 	}

pkg/kv/kvserver/closedts/policyrefresher/policy_refresher_test.go

@@ -29,7 +29,9 @@ func newNoopPolicyRefresher(stopper *stop.Stopper, settings *cluster.Settings) *
 		func() []Replica { return nil },
 		func() map[roachpb.NodeID]time.Duration {
 			return nil
-		})
+		},
+		nil,
+	)
 }
 
 type mockSpanConfig struct {
@@ -134,7 +136,7 @@ func TestPolicyRefreshOnRefreshIntervalUpdate(t *testing.T) {
 	getLeaseholders := func() []Replica { return []Replica{r} }
 	getLatencies := func() map[roachpb.NodeID]time.Duration { return nil }
 
-	pr := NewPolicyRefresher(stopper, st, getLeaseholders, getLatencies)
+	pr := NewPolicyRefresher(stopper, st, getLeaseholders, getLatencies, nil)
 	require.NotNil(t, pr)
 
 	// Start the refresher.
@@ -222,7 +224,7 @@ func TestPolicyRefresherOnLatencyIntervalUpdate(t *testing.T) {
 	closedts.RangeClosedTimestampPolicyLatencyRefreshInterval.Override(
 		ctx, &st.SV, 1*time.Hour)
 
-	pr := NewPolicyRefresher(stopper, st, getLeaseholders, getLatencies)
+	pr := NewPolicyRefresher(stopper, st, getLeaseholders, getLatencies, nil)
 	require.NotNil(t, pr)
 	pr.Run(ctx)
 

pkg/kv/kvserver/closedts/sidetransport/sender_test.go

@@ -415,10 +415,8 @@ func TestSenderWithLatencyTracker(t *testing.T) {
 		}
 	}
 
-	// Add a leaseholder with replicas in different regions.
-	r := newMockReplica(15, ctpb.LEAD_FOR_GLOBAL_READS_WITH_NO_LATENCY_INFO, 1, 2, 3)
 	s, stopper := newMockSender(connFactory)
-	policyRefresher := policyrefresher.NewPolicyRefresher(stopper, st, s.GetLeaseholders, getLatencyFn)
+	policyRefresher := policyrefresher.NewPolicyRefresher(stopper, st, s.GetLeaseholders, getLatencyFn, nil)
 	defer stopper.Stop(ctx)
 	policyRefresher.Run(ctx)
 
@@ -439,6 +437,9 @@ func TestSenderWithLatencyTracker(t *testing.T) {
 	require.Nil(t, up.Removed)
 	require.Nil(t, up.AddedOrUpdated)
 
+	// Add a leaseholder with replicas in different regions.
+	r := newMockReplica(15, ctpb.LEAD_FOR_GLOBAL_READS_WITH_NO_LATENCY_INFO, 1, 2, 3)
+
 	// Verify policy updates when adding a leaseholder with far-away replicas.
 	s.RegisterLeaseholder(ctx, r, 1)
 	testutils.SucceedsSoon(t, func() error {

pkg/kv/kvserver/replica_closedts_test.go

@@ -14,12 +14,14 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/ctpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/policyrefresher"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -33,6 +35,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/errors"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"
@@ -1201,3 +1204,137 @@ func TestRefreshPolicyWithVariousLatencies(t *testing.T) {
 		})
 	}
 }
+
+// TestReplicaClosedTSPolicyWithPolicyRefresherInMixedVersionCluster verifies
+// that the closed timestamp policy refresher behaves correctly in a
+// mixed-version cluster.
+//
+// Particularly, a race condition like below might occur:
+// 1. Side transport prepares a policy map before cluster upgrade is complete.
+// 2. Cluster upgrade completes.
+// 3. Policy refresher sees the upgrade and quickly updates replica policies to
+// use latency-based policies.
+// 4. Replica tries to use a latency-based policy but the policy map from step 1
+// doesn't include it yet.
+//
+// The logic in replica.getTargetByPolicy handles this race condition by falling
+// back to no-latency based policies if no-latency based policies were included
+// from the map provided by the side transport sender.
+//
+// This test simulates a race condition by using a testing knob to allow the
+// policy refresher to use latency based policies on replicas while the rest of
+// the system is still on an older version. It verifies that even if the policy
+// refresher sees an upgrade to V25_2 and replicas starts holding latency based
+// policies, replicas will correctly fall back to non-latency-based policies if
+// the sender hasn’t yet sent the updated latency-based policies.
+func TestReplicaClosedTSPolicyWithPolicyRefresherInMixedVersionCluster(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	prevVer := clusterversion.V25_1.Version()
+	st := cluster.MakeTestingClusterSettingsWithVersions(prevVer, prevVer, true)
+	// Helper function to check if a policy is a newly introduced latency-based policy.
+	isLatencyBasedPolicy := func(policy ctpb.RangeClosedTimestampPolicy) bool {
+		return policy >= ctpb.LEAD_FOR_GLOBAL_READS_LATENCY_LESS_THAN_20MS &&
+			policy <= ctpb.LEAD_FOR_GLOBAL_READS_LATENCY_EQUAL_OR_GREATER_THAN_300MS
+	}
+
+	// Set small intervals for faster testing.
+	closedts.RangeClosedTimestampPolicyRefreshInterval.Override(ctx, &st.SV, 5*time.Millisecond)
+	closedts.RangeClosedTimestampPolicyLatencyRefreshInterval.Override(ctx, &st.SV, 5*time.Millisecond)
+	closedts.SideTransportCloseInterval.Override(ctx, &st.SV, 5*time.Millisecond)
+
+	type latencyMap struct {
+		mu syncutil.Mutex
+		m  map[roachpb.NodeID]time.Duration
+	}
+
+	latencies := latencyMap{m: make(map[roachpb.NodeID]time.Duration)}
+	upgradeForPolicyRefresher := func() {
+		latencies.mu.Lock()
+		defer latencies.mu.Unlock()
+		latencies.m = map[roachpb.NodeID]time.Duration{
+			1: 10 * time.Millisecond,
+			2: 20 * time.Millisecond,
+			3: 50 * time.Millisecond,
+		}
+	}
+
+	tc := testcluster.StartTestCluster(t, 3,
+		base.TestClusterArgs{
+			ReplicationMode: base.ReplicationManual,
+			ServerArgs: base.TestServerArgs{
+				Knobs: base.TestingKnobs{
+					PolicyRefresherTestingKnobs: &policyrefresher.TestingKnobs{
+						InjectedLatencies: func() map[roachpb.NodeID]time.Duration {
+							latencies.mu.Lock()
+							defer latencies.mu.Unlock()
+							return latencies.m
+						},
+					},
+				},
+				Settings: st,
+			},
+		})
+	defer tc.Stopper().Stop(ctx)
+
+	key := tc.ScratchRange(t)
+	// Split the range at the table prefix and replicate it across all nodes.
+	tc.SplitRangeOrFatal(t, key)
+	tc.AddVotersOrFatal(t, key, tc.Target(1), tc.Target(2))
+
+	// Get the store and replica for testing.
+	store := tc.GetFirstStoreFromServer(t, 0)
+	replica := store.LookupReplica(roachpb.RKey(key))
+	require.NotNil(t, replica)
+	spanConfig, err := replica.LoadSpanConfig(ctx)
+	spanConfig.GlobalReads = true
+	require.NoError(t, err)
+	require.NotNil(t, spanConfig)
+	replica.SetSpanConfig(*spanConfig, roachpb.Span{Key: key})
+
+	hasLatencyBasedPolicies := func(snapshot *ctpb.Update) bool {
+		// Verify no latency-based policies are being sent.
+		latencyBasedPolicyClosedTimestamps := len(snapshot.ClosedTimestamps) == int(roachpb.MAX_CLOSED_TIMESTAMP_POLICY)
+		// Verify no latency-based policies is chosen by ranges.
+		hasLatencyBasedPolicyForAllRanges := func() bool {
+			for _, policy := range snapshot.AddedOrUpdated {
+				if isLatencyBasedPolicy(policy.Policy) {
+					return true
+				}
+			}
+			return false
+		}
+		return !latencyBasedPolicyClosedTimestamps && !hasLatencyBasedPolicyForAllRanges()
+	}
+
+	// Verify that no latency-based policies should be sent initially.
+	require.Never(t, func() bool {
+		snapshot := store.GetStoreConfig().ClosedTimestampSender.GetSnapshot()
+		expected := !hasLatencyBasedPolicies(snapshot) || len(snapshot.AddedOrUpdated) == 0
+		return !expected
+	}, 5*time.Second, 50*time.Millisecond)
+
+	// Upgrade the cluster version for policy refresher.
+	upgradeForPolicyRefresher()
+
+	// Replicas should now start holding latency based policies.
+	testutils.SucceedsSoon(t, func() error {
+		leaseholders := store.GetStoreConfig().ClosedTimestampSender.GetLeaseholders()
+		for _, lh := range leaseholders {
+			if policy := lh.(*kvserver.Replica).GetCachedClosedTimestampPolicyForTesting(); isLatencyBasedPolicy(policy) {
+				return nil
+			}
+		}
+		return errors.New("expected some leaseholder to have a latency-based policy but none had one")
+	})
+
+	// Sender does not see the cluster version upgrade yet. Replicas should fall
+	// back to no-latency based policies when side transport senders consult with
+	// leaseholders on policies to be sent.
+	require.Never(t, func() bool {
+		snapshot := store.GetStoreConfig().ClosedTimestampSender.GetSnapshot()
+		expected := !hasLatencyBasedPolicies(snapshot) || len(snapshot.AddedOrUpdated) == 0
+		return !expected
+	}, 10*time.Second, 50*time.Millisecond)
+}

pkg/server/server.go

@@ -678,8 +678,15 @@ func NewServer(cfg Config, stopper *stop.Stopper) (serverctl.ServerStartupInterf
 
 	ctSender := sidetransport.NewSender(stopper, st, clock, kvNodeDialer)
 	ctReceiver := sidetransport.NewReceiver(nodeIDContainer, stopper, stores, nil /* testingKnobs */)
-	policyRefresher := policyrefresher.NewPolicyRefresher(stopper, st, ctSender.GetLeaseholders,
-		rpcContext.RemoteClocks.AllLatencies)
+	var policyRefresher *policyrefresher.PolicyRefresher
+	{
+		var knobs *policyrefresher.TestingKnobs
+		if policyRefresherKnobs := cfg.TestingKnobs.PolicyRefresherTestingKnobs; policyRefresherKnobs != nil {
+			knobs = policyRefresherKnobs.(*policyrefresher.TestingKnobs)
+		}
+		policyRefresher = policyrefresher.NewPolicyRefresher(stopper, st, ctSender.GetLeaseholders,
+			rpcContext.RemoteClocks.AllLatencies, knobs)
+	}
 
 	// The Executor will be further initialized later, as we create more
 	// of the server's components. There's a circular dependency - many things