kvserver/closedts: change lastClosed to a map

This commit changes streamState shared between sender and
receiver from a fixed size slice to a map, similar to
trackedRange. Currently, it does not change any existing
behavior. But future commits use this to make sure only
lastClosed will be populated for enums with a corresponding
roachpb.RangeClosedTimestampPolicy until the cluster has
been fully upgraded to v25.2.

Part of: #59680
Release note: none

Author: wenyihu6

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

@@ -246,8 +246,11 @@ func (r *incomingStream) processUpdate(ctx context.Context, msg *ctpb.Update) {
 			if !ok {
 				log.Fatalf(ctx, "attempting to unregister a missing range: r%d", rangeID)
 			}
-			r.stores.ForwardSideTransportClosedTimestampForRange(
-				ctx, rangeID, r.mu.lastClosed[info.policy], info.lai)
+			ts, ok := r.mu.lastClosed[info.policy]
+			if !ok {
+				log.Fatalf(ctx, "missing closed timestamp policy %v for range r%d", info.policy, rangeID)
+			}
+			r.stores.ForwardSideTransportClosedTimestampForRange(ctx, rangeID, ts, info.lai)
 		}
 		r.mu.RUnlock()
 	}
@@ -258,9 +261,7 @@ func (r *incomingStream) processUpdate(ctx context.Context, msg *ctpb.Update) {
 
 	// Reset all the state on snapshots.
 	if msg.Snapshot {
-		for i := range r.mu.lastClosed {
-			r.mu.lastClosed[i] = hlc.Timestamp{}
-		}
+		r.mu.lastClosed = make(map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp, len(r.mu.lastClosed))
 		r.mu.tracked = make(map[roachpb.RangeID]trackedRange, len(r.mu.tracked))
 	} else if msg.SeqNum != r.mu.lastSeqNum+1 {
 		log.Fatalf(ctx, "expected closed timestamp side-transport message with sequence number "+

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

@@ -97,9 +97,12 @@ type Sender struct {
 type streamState struct {
 	// lastSeqNum is the sequence number of the last message published.
 	lastSeqNum ctpb.SeqNum
-	// lastClosed is the closed timestamp published for each policy in the
-	// last message.
-	lastClosed [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp
+	// lastClosed is the closed timestamp published for each policy in the last
+	// message. During mixed version clusters with nodes running both v25.1 and
+	// v25.2, lastClosed will only be populated for enums with a corresponding
+	// roachpb.RangeClosedTimestampPolicy (LAG_BY_CLUSTER_SETTING and
+	// LEAD_FOR_GLOBAL_READS) until the cluster is fully upgraded to v25.2.
+	lastClosed map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp
 	// tracked maintains the information that was communicated to connections in
 	// the last sent message (implicitly or explicitly). A range enters this
 	// structure as soon as it's included in a message, and exits it when it's
@@ -151,7 +154,7 @@ type Replica interface {
 	BumpSideTransportClosed(
 		ctx context.Context,
 		now hlc.ClockTimestamp,
-		targetByPolicy [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp,
+		targetByPolicy map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp,
 	) BumpSideTransportClosedResult
 }
 
@@ -210,6 +213,7 @@ func newSenderWithConnFactory(
 		buf:         newUpdatesBuf(),
 	}
 	s.trackedMu.tracked = make(map[roachpb.RangeID]trackedRange)
+	s.trackedMu.lastClosed = make(map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp)
 	s.leaseholdersMu.leaseholders = make(map[roachpb.RangeID]leaseholder)
 	s.connsMu.conns = make(map[roachpb.NodeID]conn)
 	return s
@@ -305,10 +309,14 @@ func (s *Sender) publish(ctx context.Context) hlc.ClockTimestamp {
 	defer s.trackedMu.Unlock()
 	log.VEventf(ctx, 4, "side-transport generating a new message")
 	s.trackedMu.closingFailures = [MaxReason]int{}
-
+	// TODO(wenyihu6): a cluster version check is needed here once we add new
+	// policies to ctpb.RangeClosedTimestampPolicies that do not have a
+	// corresponding roachpb.RangeClosedTimestampPolicy.
+	numPolicies := int(roachpb.MAX_CLOSED_TIMESTAMP_POLICY)
+	s.trackedMu.lastClosed = make(map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp, numPolicies)
 	msg := &ctpb.Update{
 		NodeID:           s.nodeID,
-		ClosedTimestamps: make([]ctpb.Update_GroupUpdate, len(s.trackedMu.lastClosed)),
+		ClosedTimestamps: make([]ctpb.Update_GroupUpdate, numPolicies),
 	}
 
 	// Determine the message's sequence number.
@@ -326,8 +334,7 @@ func (s *Sender) publish(ctx context.Context) hlc.ClockTimestamp {
 	lagTargetDuration := closedts.TargetDuration.Get(&s.st.SV)
 	leadTargetOverride := closedts.LeadForGlobalReadsOverride.Get(&s.st.SV)
 	sideTransportCloseInterval := closedts.SideTransportCloseInterval.Get(&s.st.SV)
-	for i := range s.trackedMu.lastClosed {
-		pol := ctpb.RangeClosedTimestampPolicy(i)
+	for pol := ctpb.RangeClosedTimestampPolicy(0); pol < ctpb.RangeClosedTimestampPolicy(numPolicies); pol++ {
 		target := closedts.TargetForPolicy(
 			now,
 			maxClockOffset,
@@ -482,14 +489,14 @@ func (s *Sender) GetSnapshot() *ctpb.Update {
 		// of incremental messages.
 		SeqNum:           s.trackedMu.lastSeqNum,
 		Snapshot:         true,
-		ClosedTimestamps: make([]ctpb.Update_GroupUpdate, len(s.trackedMu.lastClosed)),
+		ClosedTimestamps: make([]ctpb.Update_GroupUpdate, 0, len(s.trackedMu.lastClosed)),
 		AddedOrUpdated:   make([]ctpb.Update_RangeUpdate, 0, len(s.trackedMu.tracked)),
 	}
 	for pol, ts := range s.trackedMu.lastClosed {
-		msg.ClosedTimestamps[pol] = ctpb.Update_GroupUpdate{
-			Policy:          ctpb.RangeClosedTimestampPolicy(pol),
+		msg.ClosedTimestamps = append(msg.ClosedTimestamps, ctpb.Update_GroupUpdate{
+			Policy:          pol,
 			ClosedTimestamp: ts,
-		}
+		})
 	}
 	for rid, r := range s.trackedMu.tracked {
 		msg.AddedOrUpdated = append(msg.AddedOrUpdated, ctpb.Update_RangeUpdate{
@@ -886,7 +893,7 @@ func (s streamState) String() string {
 		} else {
 			agoMsg = fmt.Sprintf("%s in the future", -ago)
 		}
-		fmt.Fprintf(sb, "%s:%s (%s)", roachpb.RangeClosedTimestampPolicy(policy), closedTS, agoMsg)
+		fmt.Fprintf(sb, "%s:%s (%s)", policy, closedTS, agoMsg)
 	}
 
 	// List the tracked ranges.

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

@@ -48,10 +48,11 @@ 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) 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) BumpSideTransportClosed(
-	_ context.Context, _ hlc.ClockTimestamp, _ [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp,
+	_ context.Context, _ hlc.ClockTimestamp, _ map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp,
 ) BumpSideTransportClosedResult {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -120,7 +121,7 @@ func newMockReplica(id roachpb.RangeID, nodes ...roachpb.NodeID) *mockReplica {
 		rangeID: id,
 		canBump: true,
 		lai:     5,
-		policy:   ctpb.LAG_BY_CLUSTER_SETTING,
+		policy:  ctpb.LAG_BY_CLUSTER_SETTING,
 	}
 	r.mu.desc = desc
 	return r
@@ -137,7 +138,7 @@ func newMockReplicaEx(id roachpb.RangeID, replicas ...roachpb.ReplicationTarget)
 		rangeID: id,
 		canBump: true,
 		lai:     5,
-		policy:   ctpb.LAG_BY_CLUSTER_SETTING,
+		policy:  ctpb.LAG_BY_CLUSTER_SETTING,
 	}
 	r.mu.desc = desc
 	return r

pkg/kv/kvserver/replica.go

@@ -7,7 +7,7 @@ package kvserver
 
 import (
 	"context"
-	fmt "fmt"
+	"fmt"
 	"slices"
 	"sync"
 	"sync/atomic"

pkg/kv/kvserver/replica_closedts.go

@@ -10,6 +10,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/sidetransport"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -32,7 +33,7 @@ import (
 func (r *Replica) BumpSideTransportClosed(
 	ctx context.Context,
 	now hlc.ClockTimestamp,
-	targetByPolicy [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp,
+	targetByPolicy map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp,
 ) sidetransport.BumpSideTransportClosedResult {
 	var res sidetransport.BumpSideTransportClosedResult
 	r.mu.Lock()

pkg/kv/kvserver/replica_closedts_test.go

@@ -18,6 +18,7 @@ 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/ctpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/concurrency/lock"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvserverbase"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
@@ -262,8 +263,8 @@ func TestBumpSideTransportClosed(t *testing.T) {
 			setup: func(a setupArgs) (chan struct{}, chan error, error) {
 				// Manually bump the assigned closed timestamp to a time below
 				// where the test will attempt to bump it to.
-				var targets [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp
-				targets[roachpb.LAG_BY_CLUSTER_SETTING] = a.target.Add(-1, 0)
+				targets := map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp{}
+				targets[ctpb.LAG_BY_CLUSTER_SETTING] = a.target.Add(-1, 0)
 				return nil, nil, testutils.SucceedsSoonError(func() error {
 					res := a.repl.BumpSideTransportClosed(ctx, a.now, targets)
 					if !res.OK {
@@ -279,8 +280,8 @@ func TestBumpSideTransportClosed(t *testing.T) {
 			setup: func(a setupArgs) (chan struct{}, chan error, error) {
 				// Manually bump the assigned closed timestamp to a time equal
 				// to where the test will attempt to bump it to.
-				var targets [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp
-				targets[roachpb.LAG_BY_CLUSTER_SETTING] = a.target
+				targets := map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp{}
+				targets[ctpb.LAG_BY_CLUSTER_SETTING] = a.target
 				return nil, nil, testutils.SucceedsSoonError(func() error {
 					res := a.repl.BumpSideTransportClosed(ctx, a.now, targets)
 					if !res.OK {
@@ -296,8 +297,8 @@ func TestBumpSideTransportClosed(t *testing.T) {
 			setup: func(a setupArgs) (chan struct{}, chan error, error) {
 				// Manually bump the assigned closed timestamp to a time above
 				// where the test will attempt to bump it to.
-				var targets [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp
-				targets[roachpb.LAG_BY_CLUSTER_SETTING] = a.target.Add(1, 0)
+				targets := map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp{}
+				targets[ctpb.LAG_BY_CLUSTER_SETTING] = a.target.Add(1, 0)
 				return nil, nil, testutils.SucceedsSoonError(func() error {
 					res := a.repl.BumpSideTransportClosed(ctx, a.now, targets)
 					if !res.OK {
@@ -373,8 +374,8 @@ func TestBumpSideTransportClosed(t *testing.T) {
 			} else {
 				target, exp = test.computeTarget(repl)
 			}
-			var targets [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp
-			targets[roachpb.LAG_BY_CLUSTER_SETTING] = target
+			targets := map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp{}
+			targets[ctpb.LAG_BY_CLUSTER_SETTING] = target
 
 			// Run the setup function to get the replica in the desired state.
 			var unblockFilterC chan struct{}
@@ -716,13 +717,13 @@ func BenchmarkBumpSideTransportClosed(b *testing.B) {
 
 	manual.Pause()
 	now := s.Clock().NowAsClockTimestamp()
-	var targets [roachpb.MAX_CLOSED_TIMESTAMP_POLICY]hlc.Timestamp
+	targets := map[ctpb.RangeClosedTimestampPolicy]hlc.Timestamp{}
 
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
 		// Advance time and the closed timestamp target.
 		now = now.ToTimestamp().Add(1, 0).UnsafeToClockTimestamp()
-		targets[roachpb.LAG_BY_CLUSTER_SETTING] = now.ToTimestamp()
+		targets[ctpb.LAG_BY_CLUSTER_SETTING] = now.ToTimestamp()
 
 		// Perform the call.
 		res := r.BumpSideTransportClosed(ctx, now, targets)