Merge pull request #152455 from iskettaneh/backport25.2-144272

release-25.2: liveness: lazily calculate the node connectivity inside NodeVitality

Author: iskettaneh

pkg/kv/kvserver/asim/state/liveness.go

@@ -89,11 +89,13 @@ func convertNodeStatusToNodeVitality(
 		l.Expiration = liveTs
 		l.Draining = true
 	}
+	ncs := livenesspb.NewNodeConnectionStatus(nid, nil)
+	ncs.SetIsConnected(true)
 	entry := l.CreateNodeVitality(
 		now,               /* now */
 		hlc.Timestamp{},   /* descUpdateTime */
 		hlc.Timestamp{},   /* descUnavailableTime */
-		true,              /* connected */
+		ncs,               /* nodeConnectionStatus */
 		timeUntilNodeDead, /* timeUntilNodeDead */
 		0,                 /* timeAfterNodeSuspect */
 	)

pkg/kv/kvserver/liveness/BUILD.bazel

@@ -18,7 +18,6 @@ go_library(
         "//pkg/kv/kvpb",
         "//pkg/kv/kvserver/liveness/livenesspb",
         "//pkg/roachpb",
-        "//pkg/rpc",
         "//pkg/rpc/nodedialer",
         "//pkg/server/telemetry",
         "//pkg/settings",

pkg/kv/kvserver/liveness/cache.go

@@ -13,7 +13,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/gossip"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness/livenesspb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
-	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -244,14 +243,13 @@ func (c *Cache) convertToNodeVitality(l livenesspb.Liveness) livenesspb.NodeVita
 	// even before the first gossip arrives for a store.
 
 	// NB: nodeDialer is nil in some tests.
-	connected := c.nodeDialer == nil || c.nodeDialer.ConnHealth(l.NodeID, rpc.SystemClass) == nil
 	lastDescUpdate := c.lastDescriptorUpdate(l.NodeID)
 
 	return l.CreateNodeVitality(
 		c.clock.Now(),
 		lastDescUpdate.lastUpdateTime,
 		lastDescUpdate.lastUnavailableTime,
-		connected,
+		livenesspb.NewNodeConnectionStatus(l.NodeID, c.nodeDialer),
 		TimeUntilNodeDead.Get(&c.st.SV),
 		TimeAfterNodeSuspect.Get(&c.st.SV),
 	)

pkg/kv/kvserver/liveness/client_test.go

@@ -417,3 +417,41 @@ func TestLivenessRangeGetsPeriodicallyCompacted(t *testing.T) {
 		return nil
 	})
 }
+
+// TestNodeConnectionStatusIsConnected verifies that NodeConnectionStatus caches
+// the connection status after the first IsConnected() call.
+func TestNodeConnectionStatusIsConnected(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	testCases := []struct {
+		name             string
+		initialConnected bool
+	}{
+		{
+			name:             "initially connected",
+			initialConnected: true,
+		},
+		{
+			name:             "initially disconnected",
+			initialConnected: false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Create a mock node dialer with the initial connection state.
+			connHealth := livenesspb.NewMockNodeConnectionHealth(tc.initialConnected)
+			ncs := livenesspb.NewNodeConnectionStatus(roachpb.NodeID(1), connHealth)
+
+			// First call should calculate and cache the connection status.
+			require.Equal(t, tc.initialConnected, ncs.IsConnected())
+
+			// Modify the mock to simulate a connection state change.
+			connHealth.SetConnected(!tc.initialConnected)
+
+			// Subsequent calls should use the cached value (still the initial state).
+			require.Equal(t, tc.initialConnected, ncs.IsConnected())
+		})
+	}
+}

pkg/kv/kvserver/liveness/livenesspb/BUILD.bazel

@@ -13,6 +13,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/roachpb",
+        "//pkg/rpc",
         "//pkg/util/hlc",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",

pkg/kv/kvserver/liveness/livenesspb/liveness.go

@@ -8,9 +8,11 @@ package livenesspb
 import (
 	"context"
 	"fmt"
+	"sync/atomic"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/redact"
@@ -159,6 +161,61 @@ type IsLiveMapEntry struct {
 // IsLiveMap is a type alias for a map from NodeID to IsLiveMapEntry.
 type IsLiveMap map[roachpb.NodeID]IsLiveMapEntry
 
+// NodeConnectionHealth represents the minimal interface needed for checking if
+// a node RPC connection is alive.
+type NodeConnectionHealth interface {
+	// ConnHealth returns nil if we have an open connection of the request
+	// class that works successfully. Otherwise, it returns an error.
+	ConnHealth(nodeID roachpb.NodeID, class rpc.ConnectionClass) error
+}
+
+// NodeConnectionStatus is a lightweight wrapper around the
+// NodeConnectionHealth, where calculating the connection status is done lazily
+// upon the first call to IsConnected(). This is useful as a member in structs
+// where we sometimes want to calculate the connection status, and sometimes we
+// don't.
+type NodeConnectionStatus struct {
+	nodeConnectionHealth NodeConnectionHealth
+	nodeID               roachpb.NodeID
+	// calculatedConnected specifies whether we did calculate the connected field
+	// or not. If we haven't, we can't rely on the value of connected.
+	calculatedConnected atomic.Bool
+	// connected is an atomic boolean that tracks whether we are connected to the node.
+	connected atomic.Bool
+}
+
+func NewNodeConnectionStatus(
+	nodeID roachpb.NodeID, nodeConnectionHealth NodeConnectionHealth,
+) *NodeConnectionStatus {
+	ncs := &NodeConnectionStatus{
+		nodeConnectionHealth: nodeConnectionHealth,
+		nodeID:               nodeID,
+	}
+	return ncs
+}
+
+// SetIsConnected changes the connection status of the node.
+func (ncs *NodeConnectionStatus) SetIsConnected(connected bool) {
+	ncs.connected.Store(connected)
+	ncs.calculatedConnected.Store(true)
+}
+
+// IsConnected checks if we are connected to the supplied nodeID. It only
+// performs the calculation the first time. Future calls will use the cached
+// version.
+func (ncs *NodeConnectionStatus) IsConnected() bool {
+	if !ncs.calculatedConnected.Load() {
+		// Calculate the connection status if we haven't done that before.
+		// Some tests will set the nodeDialer to nil, so we need to check for that.
+		connected := ncs.nodeConnectionHealth == nil ||
+			ncs.nodeConnectionHealth.ConnHealth(ncs.nodeID, rpc.SystemClass) == nil
+		ncs.SetIsConnected(connected)
+		ncs.calculatedConnected.Store(true)
+	}
+
+	return ncs.connected.Load()
+}
+
 // NodeVitality should be used any place other than epoch leases where it is
 // necessary to determine if a node is currently alive and what its health is.
 // Aliveness and deadness are concepts that refer to our best guess of the
@@ -173,8 +230,9 @@ type NodeVitality struct {
 	draining bool
 	// membership is whether the node is active or in a state of decommissioning.
 	membership MembershipStatus
-	// connected is whether we are currently directly connect to this node.
-	connected bool
+	// nodeConnectionStatus calculates whether we are currently directly connect
+	// to this node.
+	nodeConnectionStatus *NodeConnectionStatus
 
 	// When the record is created. Records are not held for long, but they should
 	// always give consistent results when asked.
@@ -285,7 +343,7 @@ func (nv NodeVitality) IsLive(usage VitalityUsage) bool {
 			return nv.isAliveAndConnected()
 		}
 	case NetworkMap:
-		return nv.connected
+		return nv.nodeConnectionStatus.IsConnected()
 	case LossOfQuorum:
 		return nv.isAlive()
 	case ReplicaGCQueue:
@@ -315,7 +373,7 @@ func (nv NodeVitality) isAvailableNotDraining() bool {
 }
 
 func (nv NodeVitality) isAliveAndConnected() bool {
-	return nv.isAvailableNotDraining() && nv.connected
+	return nv.isAvailableNotDraining() && nv.nodeConnectionStatus.IsConnected()
 }
 
 // isAliveEpoch is used for epoch leases. It is similar to isAlive, but doesn't
@@ -535,7 +593,7 @@ func (l Liveness) CreateNodeVitality(
 	now hlc.Timestamp,
 	descUpdateTime hlc.Timestamp,
 	descUnavailableTime hlc.Timestamp,
-	connected bool,
+	connectionStatus *NodeConnectionStatus,
 	timeUntilNodeDead time.Duration,
 	timeAfterNodeSuspect time.Duration,
 ) NodeVitality {
@@ -546,7 +604,7 @@ func (l Liveness) CreateNodeVitality(
 		nodeID:               l.NodeID,
 		draining:             l.Draining,
 		membership:           l.Membership,
-		connected:            connected,
+		nodeConnectionStatus: connectionStatus,
 		now:                  now,
 		descUpdateTime:       descUpdateTime,
 		descUnavailableTime:  descUnavailableTime,

pkg/kv/kvserver/liveness/livenesspb/liveness_test_helper.go

@@ -10,8 +10,10 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
+	"github.com/cockroachdb/errors"
 )
 
 // testNodeVitalityEntry is here to minimize the impact on tests of changing to
@@ -49,10 +51,13 @@ func TestCreateNodeVitality(ids ...roachpb.NodeID) TestNodeVitality {
 
 func (e testNodeVitalityEntry) convert() NodeVitality {
 	now := e.clock.Now()
+	ncs := NewNodeConnectionStatus(0, nil)
 	if e.Alive {
-		return e.Liveness.CreateNodeVitality(now, now, hlc.Timestamp{}, true, time.Second, time.Second)
+		ncs.SetIsConnected(true)
+		return e.Liveness.CreateNodeVitality(now, now, hlc.Timestamp{}, ncs, time.Second, time.Second)
 	} else {
-		return e.Liveness.CreateNodeVitality(now, now.AddDuration(-time.Hour), hlc.Timestamp{}, false, time.Second, time.Second)
+		ncs.SetIsConnected(false)
+		return e.Liveness.CreateNodeVitality(now, now.AddDuration(-time.Hour), hlc.Timestamp{}, ncs, time.Second, time.Second)
 	}
 }
 
@@ -159,21 +164,49 @@ func (tnv TestNodeVitality) RestartNode(id roachpb.NodeID) {
 // FakeNodeVitality creates a node vitality record that is either dead or alive
 // by all accounts.
 func FakeNodeVitality(alive bool) NodeVitality {
+	ncs := NewNodeConnectionStatus(1, nil)
 	if alive {
+		ncs.SetIsConnected(true)
 		return NodeVitality{
 			nodeID:               1,
-			connected:            true,
+			nodeConnectionStatus: ncs,
 			now:                  hlc.Timestamp{}.AddDuration(time.Nanosecond),
 			timeUntilNodeDead:    time.Second,
 			timeAfterNodeSuspect: time.Second,
 			livenessExpiration:   hlc.Timestamp{}.AddDuration(2 * time.Nanosecond),
 			livenessEpoch:        1,
 		}
 	} else {
+		ncs.SetIsConnected(false)
 		return NodeVitality{
-			nodeID:        1,
-			connected:     false,
-			livenessEpoch: 1,
+			nodeID:               1,
+			nodeConnectionStatus: ncs,
+			livenessEpoch:        1,
 		}
 	}
 }
+
+// MockNodeConnectionHealth implements the NodeConnectionHealth interface for
+// testing. It just has one field to simulate the connection state.
+type MockNodeConnectionHealth struct {
+	connected bool
+}
+
+// NewMockNodeConnectionHealth creates a new MockNodeConnectionHealth with the
+// given connection state.
+func NewMockNodeConnectionHealth(connected bool) *MockNodeConnectionHealth {
+	return &MockNodeConnectionHealth{connected: connected}
+}
+
+// SetConnected sets the connection state of the mock node connection health.
+func (m *MockNodeConnectionHealth) SetConnected(connected bool) {
+	m.connected = connected
+}
+
+// ConnHealth implements the NodeDialer interface.
+func (m *MockNodeConnectionHealth) ConnHealth(_ roachpb.NodeID, _ rpc.ConnectionClass) error {
+	if m.connected {
+		return nil
+	}
+	return errors.Errorf("not connected")
+}