Merge #143795

143795: kvserver: use policy refresher for closed timestamp policy r=arulajmani a=wenyihu6

**kvserver: use policy refresher for closed timestamp policy**

Previously, closed timestamp policies were computed on-demand during
`r.closedTimestampPolicyRLocked` under RLock.

This commit changes it so that is refreshed async during specific events
instead of on-demand during `r.closedTimestampPolicyRLocked`. This helps
reduce the time spent under the `r.mu.RLock()` in `closedTimestampPolicyRLocked`
and also helps pave the foundation for future changes to use latency based
closed timestamp policies.

Policies are now refreshed on demand during specific events:
- when there is a leaseholder change leasePostApplyLocked
- when there is a config change r.SetSpanConfig
- or periodically at the interval of kv.closed_timestamp.policy_refresh_interval.

Part of: #59680
Release note: none

---

**kvserver: rename registerLeaseholder to AndRefreshPolicy**

This commit renames registerLeaseholder to registerLeaseholderAndRefreshPolicy
to reflect its changes. In addition to registering the replica as a leaseholder,
it now also triggersa closed timestamp policy refresh on the replica.

Epic: none
Release note: none

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

Author: craig[bot]

pkg/kv/kvserver/BUILD.bazel

@@ -141,6 +141,7 @@ go_library(
         "//pkg/kv/kvserver/benignerror",
         "//pkg/kv/kvserver/closedts",
         "//pkg/kv/kvserver/closedts/ctpb",
+        "//pkg/kv/kvserver/closedts/policyrefresher",
         "//pkg/kv/kvserver/closedts/sidetransport",
         "//pkg/kv/kvserver/closedts/tracker",
         "//pkg/kv/kvserver/concurrency",

pkg/kv/kvserver/client_split_test.go

@@ -4140,7 +4140,12 @@ func TestStoreRangeSplitAndMergeWithGlobalReads(t *testing.T) {
 
 	// Verify that the closed timestamp policy is set up.
 	repl := store.LookupReplica(roachpb.RKey(descKey))
-	require.Equal(t, repl.ClosedTimestampPolicy(), roachpb.LEAD_FOR_GLOBAL_READS)
+	testutils.SucceedsSoon(t, func() error {
+		if actual := repl.ClosedTimestampPolicy(); actual != roachpb.LEAD_FOR_GLOBAL_READS {
+			return errors.Newf("expected LEAD_FOR_GLOBAL_READS, got %s", actual)
+		}
+		return nil
+	})
 
 	// Write to the range, which has the effect of bumping the closed timestamp.
 	pArgs := putArgs(descKey, []byte("foo"))

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

@@ -16,6 +16,7 @@ go_library(
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/closedts",
         "//pkg/kv/kvserver/closedts/ctpb",
+        "//pkg/kv/kvserver/closedts/policyrefresher",
         "//pkg/roachpb",
         "//pkg/rpc",
         "//pkg/rpc/nodedialer",

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

@@ -21,6 +21,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"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/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
@@ -132,6 +133,7 @@ type leaseholder struct {
 // Replica represents a *Replica object, but with only the capabilities needed
 // by the closed timestamp side transport to accomplish its job.
 type Replica interface {
+	policyrefresher.Replica
 	// Accessors.
 	StoreID() roachpb.StoreID
 	GetRangeID() roachpb.RangeID
@@ -474,6 +476,18 @@ func (s *Sender) publish(ctx context.Context) hlc.ClockTimestamp {
 	return now
 }
 
+// GetLeaseholders returns a slice of all replicas that are currently
+// leaseholders on this node.
+func (s *Sender) GetLeaseholders() []policyrefresher.Replica {
+	s.leaseholdersMu.Lock()
+	defer s.leaseholdersMu.Unlock()
+	leaseholders := make([]policyrefresher.Replica, 0, len(s.leaseholdersMu.leaseholders))
+	for _, lh := range s.leaseholdersMu.leaseholders {
+		leaseholders = append(leaseholders, lh.Replica)
+	}
+	return leaseholders
+}
+
 // GetSnapshot generates an update that contains all the sender's state (as
 // opposed to being an incremental delta since a previous message). The returned
 // msg will have the `snapshot` field set, and a sequence number indicating

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

@@ -48,9 +48,9 @@ type mockReplica struct {
 
 var _ Replica = &mockReplica{}
 
-func (m *mockReplica) StoreID() roachpb.StoreID                                    { return m.storeID }
-func (m *mockReplica) GetRangeID() roachpb.RangeID                                 { return m.rangeID }
-func (m *mockReplica) RefreshLatency(latencyInfo map[roachpb.NodeID]time.Duration) {}
+func (m *mockReplica) StoreID() roachpb.StoreID                                   { return m.storeID }
+func (m *mockReplica) GetRangeID() roachpb.RangeID                                { return m.rangeID }
+func (m *mockReplica) RefreshPolicy(latencyInfo map[roachpb.NodeID]time.Duration) {}
 func (m *mockReplica) BumpSideTransportClosed(
 	_ context.Context, _ hlc.ClockTimestamp, _ map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp,
 ) BumpSideTransportClosedResult {

pkg/kv/kvserver/replica.go

@@ -1083,6 +1083,12 @@ type Replica struct {
 		// GC threshold for the range.
 		pendingGCThreshold hlc.Timestamp
 	}
+
+	// cachedClosedTimestampPolicy is the cached closed timestamp policy of the
+	// range. It is updated asynchronously by listening on span configuration
+	// changes, leaseholder changes, and periodically at the interval of
+	// kv.closed_timestamp.policy_refresh_interval by PolicyRefresher.
+	cachedClosedTimestampPolicy atomic.Int32
 }
 
 // String returns the string representation of the replica using an
@@ -1174,6 +1180,7 @@ func (r *Replica) SetSpanConfig(conf roachpb.SpanConfig, sp roachpb.Span) bool {
 	r.mu.conf = conf
 	r.mu.spanConfigExplicitlySet = true
 	r.mu.confSpan = sp
+	r.store.policyRefresher.EnqueueReplicaForRefresh(r)
 	return oldConf.HasConfigurationChange(conf)
 }
 
@@ -1314,23 +1321,43 @@ func toClientClosedTsPolicy(
 }
 
 // closedTimestampPolicyRLocked returns the closed timestamp policy of the
-// range, which is updated asynchronously by listening in on span configuration
-// changes.
+// range, which is updated asynchronously by listening on span configuration
+// changes, leaseholder changes, and periodically at the interval of
+// kv.closed_timestamp.policy_refresh_interval.
 //
 // NOTE: an exported version of this method which does not require the replica
 // lock exists in helpers_test.go. Move here if needed.
 func (r *Replica) closedTimestampPolicyRLocked() ctpb.RangeClosedTimestampPolicy {
-	if r.mu.conf.GlobalReads {
-		if !r.shMu.state.Desc.ContainsKey(roachpb.RKey(keys.NodeLivenessPrefix)) {
-			return ctpb.LEAD_FOR_GLOBAL_READS_WITH_NO_LATENCY_INFO
-		}
+	// TODO(wenyi): try to remove the need for this key comparison under RLock.
+	// See more in #143648.
+	if r.shMu.state.Desc.ContainsKey(roachpb.RKey(keys.NodeLivenessPrefix)) {
+		return ctpb.LAG_BY_CLUSTER_SETTING
+	}
+	return ctpb.RangeClosedTimestampPolicy(r.cachedClosedTimestampPolicy.Load())
+}
+
+// RefreshPolicy updates the replica's cached closed timestamp policy based on
+// its span configuration. Note that the given map can be nil.
+//
+// TODO(wenyihu6): update this function to consult the supplied map and
+// clarify what happens when the map is nil
+func (r *Replica) RefreshPolicy(_ map[roachpb.NodeID]time.Duration) {
+	policy := func() ctpb.RangeClosedTimestampPolicy {
+		desc, conf := r.DescAndSpanConfig()
 		// The node liveness range ignores zone configs and always uses a
 		// LAG_BY_CLUSTER_SETTING closed timestamp policy. If it was to begin
 		// closing timestamps in the future, it would break liveness updates,
 		// which perform a 1PC transaction with a commit trigger and can not
 		// tolerate being pushed into the future.
+		if desc.ContainsKey(roachpb.RKey(keys.NodeLivenessPrefix)) {
+			return ctpb.LAG_BY_CLUSTER_SETTING
+		}
+		if !conf.GlobalReads {
+			return ctpb.LAG_BY_CLUSTER_SETTING
+		}
+		return ctpb.LEAD_FOR_GLOBAL_READS_WITH_NO_LATENCY_INFO
 	}
-	return ctpb.LAG_BY_CLUSTER_SETTING
+	r.cachedClosedTimestampPolicy.Store(int32(policy()))
 }
 
 // NodeID returns the ID of the node this replica belongs to.

pkg/kv/kvserver/replica_closedts_internal_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/batcheval"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/ctpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -30,6 +31,12 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
+// TestingSetCachedClosedTimestampPolicy sets the closed timestamp policy on r
+// to be the given policy. It is a test-only helper method.
+func (r *Replica) TestingSetCachedClosedTimestampPolicy(policy ctpb.RangeClosedTimestampPolicy) {
+	r.cachedClosedTimestampPolicy.Store(int32(policy))
+}
+
 func TestSideTransportClosed(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)

pkg/kv/kvserver/replica_closedts_test.go

@@ -18,11 +18,13 @@ import (
 	"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/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
@@ -937,3 +939,171 @@ func testNonBlockingReadsWithReaderFn(
 	atomic.StoreInt32(&done, 1)
 	require.NoError(t, g.Wait())
 }
+
+// TestClosedTimestampPolicyRefreshOnSetSpanConfig tests that SetSpanConfig
+// correctly triggers the closed timestamp policy refresh.
+func TestClosedTimestampPolicyRefreshOnSetSpanConfig(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{})
+	defer tc.Stopper().Stop(ctx)
+
+	scratchKey := tc.ScratchRange(t)
+
+	repl := tc.GetFirstStoreFromServer(t, 0).LookupReplica(roachpb.RKey(scratchKey))
+	require.Equal(t, roachpb.LAG_BY_CLUSTER_SETTING, repl.GetRangeInfo(ctx).ClosedTimestampPolicy)
+
+	spanConfig, err := repl.LoadSpanConfig(ctx)
+	spanConfig.GlobalReads = true
+	require.NoError(t, err)
+	require.NotNil(t, spanConfig)
+	repl.SetSpanConfig(*spanConfig, roachpb.Span{Key: scratchKey})
+
+	// Trigger policy refresh.
+	testutils.SucceedsSoon(t, func() error {
+		if repl.GetRangeInfo(ctx).ClosedTimestampPolicy != roachpb.LEAD_FOR_GLOBAL_READS {
+			return errors.New("expected LEAD_FOR_GLOBAL_READS")
+		}
+		return nil
+	})
+}
+
+// TestClosedTimestampPolicyRefreshIntervalOnLivenessRanges tests that the
+// closed timestamp policy is correctly applied to the node liveness range. That
+// is, even if we try to set the node liveness range to have global reads, the
+// closed timestamp policy should still be LAG_BY_CLUSTER_SETTING. Read more in
+// replica.closedTimestampPolicyRLocked.
+func TestClosedTimestampPolicyRefreshIntervalOnLivenessRanges(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{})
+	defer tc.Stopper().Stop(ctx)
+
+	// Get the node liveness range descriptor.
+	livenessRangeDesc, err := tc.LookupRange(keys.NodeLivenessPrefix)
+	require.NoError(t, err)
+
+	// Check liveness range policy.
+	livenessRepl := tc.GetFirstStoreFromServer(t, 0).LookupReplica(livenessRangeDesc.StartKey)
+	require.Equal(t, roachpb.LAG_BY_CLUSTER_SETTING, livenessRepl.GetRangeInfo(ctx).ClosedTimestampPolicy)
+
+	spanConfig, err := livenessRepl.LoadSpanConfig(ctx)
+	spanConfig.GlobalReads = true
+	require.NoError(t, err)
+	require.NotNil(t, spanConfig)
+	livenessRepl.SetSpanConfig(*spanConfig, roachpb.Span{Key: keys.NodeLivenessPrefix})
+
+	require.Never(t, func() bool {
+		expectedState := livenessRepl.GetRangeInfo(ctx).ClosedTimestampPolicy == roachpb.LAG_BY_CLUSTER_SETTING
+		return !expectedState
+	}, 3*time.Second, 500*time.Millisecond)
+}
+
+// TestSideTransportLeaseholder verifies that a range's leaseholder is properly
+// tracked by the closed timestamp side transport, even when the range is
+// receiving writes and the side transport interval is disabled.
+func TestSideTransportLeaseholder(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ctx := context.Background()
+	st := cluster.MakeTestingClusterSettings()
+	// Disable side transport interval to verify tracking works even without
+	// active transport.
+	closedts.SideTransportCloseInterval.Override(ctx, &st.SV, 0)
+	tc := serverutils.StartCluster(t, 3, base.TestClusterArgs{
+		ReplicationMode: base.ReplicationManual,
+		ServerArgs: base.TestServerArgs{
+			Settings: st,
+		},
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	// Get store and create test range.
+	store, err := tc.Server(0).GetStores().(*kvserver.Stores).GetStore(tc.Server(0).GetFirstStoreID())
+	require.NoError(t, err)
+	scratchKey := tc.ScratchRange(t)
+	tc.AddVotersOrFatal(t, scratchKey, tc.Target(1))
+	tc.AddNonVotersOrFatal(t, scratchKey, tc.Target(2))
+	repl := store.LookupReplica(roachpb.RKey(scratchKey))
+	require.NotNil(t, repl)
+
+	// Start goroutine that continuously writes to the range to create write load.
+	go func() {
+		for {
+			select {
+			case <-time.After(10 * time.Millisecond):
+				pArgs := putArgs(scratchKey, []byte("value"))
+				if _, pErr := kv.SendWrapped(ctx, store.DB().NonTransactionalSender(), pArgs); pErr != nil {
+					log.Errorf(ctx, "failed to put value: %s", pErr)
+				}
+			case <-tc.Stopper().ShouldQuiesce():
+				return
+			}
+		}
+	}()
+
+	// Verify that the range appears in the closed timestamp sender's leaseholders
+	// list despite write load and disabled side transport.
+	testutils.SucceedsSoon(t, func() error {
+		closedTsSender := store.GetStoreConfig().ClosedTimestampSender
+		leaseholders := closedTsSender.GetLeaseholders()
+		for _, lh := range leaseholders {
+			if lh.(*kvserver.Replica).RangeID == repl.RangeID {
+				return nil
+			}
+		}
+		return errors.Errorf("range %d not found in leaseholders slice", repl.RangeID)
+	})
+}
+
+// TestClosedTimestampPolicyRefreshIntervalOnLeaseTransfers tests that the
+// closed timestamp policy is correctly refreshed on a range after a lease
+// transfer.
+func TestClosedTimestampPolicyRefreshIntervalOnLeaseTransfers(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	tc := testcluster.StartTestCluster(t, 3, base.TestClusterArgs{
+		ReplicationMode: base.ReplicationManual,
+	})
+	defer tc.Stopper().Stop(ctx)
+
+	scratchKey := tc.ScratchRange(t)
+	desc := tc.AddVotersOrFatal(t, scratchKey, tc.Target(1), tc.Target(2))
+
+	repl1 := tc.GetFirstStoreFromServer(t, 0).LookupReplica(roachpb.RKey(scratchKey))
+	require.Equal(t, roachpb.LAG_BY_CLUSTER_SETTING, repl1.GetRangeInfo(ctx).ClosedTimestampPolicy)
+
+	repl2 := tc.GetFirstStoreFromServer(t, 1).LookupReplica(roachpb.RKey(scratchKey))
+	require.Equal(t, roachpb.LAG_BY_CLUSTER_SETTING, repl2.GetRangeInfo(ctx).ClosedTimestampPolicy)
+
+	spanConfig, err := repl2.LoadSpanConfig(ctx)
+	spanConfig.GlobalReads = true
+	require.NoError(t, err)
+	require.NotNil(t, spanConfig)
+	repl2.SetSpanConfig(*spanConfig, roachpb.Span{Key: scratchKey})
+	testutils.SucceedsSoon(t, func() error {
+		if repl2.GetRangeInfo(ctx).ClosedTimestampPolicy != roachpb.LEAD_FOR_GLOBAL_READS {
+			return errors.New("expected LEAD_FOR_GLOBAL_READS")
+		}
+		return nil
+	})
+
+	// Force repl2 policy to be LAG_BY_CLUSTER_SETTING.
+	repl2.TestingSetCachedClosedTimestampPolicy(ctpb.LAG_BY_CLUSTER_SETTING)
+	require.Equal(t, roachpb.LAG_BY_CLUSTER_SETTING, repl2.GetRangeInfo(ctx).ClosedTimestampPolicy)
+
+	// Ensure that transferring the lease to repl2 does trigger a lease refresh.
+	require.NoError(t, tc.TransferRangeLease(desc, tc.Target(1)))
+	testutils.SucceedsSoon(t, func() error {
+		if actual := repl2.GetRangeInfo(ctx).ClosedTimestampPolicy; actual != roachpb.LEAD_FOR_GLOBAL_READS {
+			return errors.Newf("expected LEAD_FOR_GLOBAL_READS but got %v", actual)
+		}
+		return nil
+	})
+}

pkg/kv/kvserver/replica_init.go

@@ -277,6 +277,7 @@ func newUninitializedReplicaWithoutRaftGroup(
 		EnabledWhenLeaderLevel: r.raftMu.flowControlLevel,
 		Knobs:                  r.store.TestingKnobs().FlowControlTestingKnobs,
 	})
+	r.RefreshPolicy(nil)
 	return r
 }
 

pkg/kv/kvserver/replica_proposal.go

@@ -581,7 +581,7 @@ func (r *Replica) leasePostApplyLocked(
 
 	// Inform the store of this lease.
 	if iAmTheLeaseHolder {
-		r.store.registerLeaseholder(ctx, r, newLease.Sequence)
+		r.store.registerLeaseholderAndRefreshPolicy(ctx, r, newLease.Sequence)
 	} else {
 		r.store.unregisterLeaseholder(ctx, r)
 	}