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

Author: wenyihu6

pkg/base/testing_knobs.go

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

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/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