liveness: lazily calculate the node connectivity inside NodeVitality

Before this commit, always calculate if we are connected to the node as
part of the NodeVitality creation. However, most uses of node vitality
don't need that calculation to be made. This commit addresses this
performance issue by lazily calculating the connected field only when
requested.

Fixes: #143144

Release note: None

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")
+}