rpc: return DRPC status from ConnHealth and GetBreaker when enabled

Currently, ConnHealth and GetBreakerForAddr methods in nodedialer always return
gRPC status even when DRPC is enabled. When DRPC is enabled, they should return
the connection health and breaker status of DRPC instead. This change fixes
that behavior.

Fixes: 153860
Epic: CRDB-51459
Release note: None

Author: cthumuluru-crdb

pkg/kv/kvserver/client_raft_test.go

@@ -2850,6 +2850,10 @@ func TestReportUnreachableHeartbeats(t *testing.T) {
 	ctx := context.Background()
 	tc := testcluster.StartTestCluster(t, 3,
 		base.TestClusterArgs{
+			ServerArgs: base.TestServerArgs{
+				// TODO(server): enabled DRPC once serverutils adds support for DRPC.
+				DefaultDRPCOption: base.TestDRPCDisabled,
+			},
 			ReplicationMode: base.ReplicationManual,
 		})
 	defer tc.Stopper().Stop(ctx)
@@ -2911,6 +2915,10 @@ func TestReportUnreachableRemoveRace(t *testing.T) {
 	ctx := context.Background()
 	tc := testcluster.StartTestCluster(t, 3,
 		base.TestClusterArgs{
+			ServerArgs: base.TestServerArgs{
+				// TODO(server): enabled DRPC once serverutils adds support for DRPC.
+				DefaultDRPCOption: base.TestDRPCDisabled,
+			},
 			ReplicationMode: base.ReplicationManual,
 		})
 	defer tc.Stopper().Stop(ctx)

pkg/rpc/context.go

@@ -1392,12 +1392,31 @@ func (rpcCtx *Context) ConnHealth(
 	if rpcCtx.GetLocalInternalClientForAddr(nodeID) != nil {
 		return nil
 	}
+
+	if !rpcbase.DRPCEnabled(context.Background(), rpcCtx.Settings) {
+		return rpcCtx.grpcConnHealth(target, nodeID, class)
+	}
+	return rpcCtx.drpcConnHealth(target, nodeID, class)
+}
+
+func (rpcCtx *Context) grpcConnHealth(
+	target string, nodeID roachpb.NodeID, class rpcbase.ConnectionClass,
+) error {
 	if p, ok := rpcCtx.peers.get(peerKey{target, nodeID, class}); ok {
 		return p.c.Health()
 	}
 	return ErrNotHeartbeated
 }
 
+func (rpcCtx *Context) drpcConnHealth(
+	target string, nodeID roachpb.NodeID, class rpcbase.ConnectionClass,
+) error {
+	if p, ok := rpcCtx.drpcPeers.get(peerKey{target, nodeID, class}); ok {
+		return p.c.Health()
+	}
+	return ErrNotHeartbeated
+}
+
 type transportType bool
 
 const (
@@ -1499,14 +1518,32 @@ func (rpcCtx *Context) dialOptsLocal() ([]grpc.DialOption, error) {
 func (rpcCtx *Context) GetBreakerForAddr(
 	nodeID roachpb.NodeID, class rpcbase.ConnectionClass, addr net.Addr,
 ) (*circuitbreaker.Breaker, bool) {
-	sAddr := addr.String()
-	rpcCtx.peers.mu.RLock()
-	defer rpcCtx.peers.mu.RUnlock()
-	p, ok := rpcCtx.peers.mu.m[peerKey{
-		TargetAddr: sAddr,
+	k := peerKey{
+		TargetAddr: addr.String(),
 		NodeID:     nodeID,
 		Class:      class,
-	}]
+	}
+
+	if !rpcbase.DRPCEnabled(context.Background(), rpcCtx.Settings) {
+		return rpcCtx.grpcGetBreakerForAddr(k)
+	}
+	return rpcCtx.drpcGetBreakerForAddr(k)
+}
+
+func (rpcCtx *Context) grpcGetBreakerForAddr(k peerKey) (*circuitbreaker.Breaker, bool) {
+	rpcCtx.peers.mu.RLock()
+	defer rpcCtx.peers.mu.RUnlock()
+	p, ok := rpcCtx.peers.mu.m[k]
+	if !ok {
+		return nil, false
+	}
+	return p.b, true
+}
+
+func (rpcCtx *Context) drpcGetBreakerForAddr(k peerKey) (*circuitbreaker.Breaker, bool) {
+	rpcCtx.drpcPeers.mu.RLock()
+	defer rpcCtx.drpcPeers.mu.RUnlock()
+	p, ok := rpcCtx.drpcPeers.mu.m[k]
 	if !ok {
 		return nil, false
 	}

pkg/rpc/nodedialer/BUILD.bazel

@@ -52,6 +52,9 @@ go_test(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_stretchr_testify//assert",
         "@com_github_stretchr_testify//require",
+        "@io_storj_drpc//drpcmigrate",
+        "@io_storj_drpc//drpcmux",
+        "@io_storj_drpc//drpcserver",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
     ],

pkg/rpc/nodedialer/nodedialer.go

@@ -282,8 +282,11 @@ func (n *Dialer) ConnHealthTryDial(nodeID roachpb.NodeID, class rpcbase.Connecti
 		return err
 	}
 	// NB: This will always return `ErrNotHeartbeated` since the heartbeat will
-	// not be done by the time `Health` is called since GRPCDialNode is async.
-	return n.rpcContext.GRPCDialNode(addr.String(), nodeID, locality, class).Health()
+	// not be done by the time `Health` is called since DialNode is async.
+	if !rpcbase.DRPCEnabled(context.Background(), n.rpcContext.Settings) {
+		return n.rpcContext.GRPCDialNode(addr.String(), nodeID, locality, class).Health()
+	}
+	return n.rpcContext.DRPCDialNode(addr.String(), nodeID, locality, class).Health()
 }
 
 // ConnHealthTryDialInstance returns nil if we have an open connection of the
@@ -298,7 +301,11 @@ func (n *Dialer) ConnHealthTryDialInstance(id base.SQLInstanceID, addr string) e
 		addr, roachpb.NodeID(id), rpcbase.DefaultClass); err == nil {
 		return nil
 	}
-	return n.rpcContext.GRPCDialPod(addr, id, roachpb.Locality{}, rpcbase.DefaultClass).Health()
+
+	if !rpcbase.DRPCEnabled(context.Background(), n.rpcContext.Settings) {
+		return n.rpcContext.GRPCDialPod(addr, id, roachpb.Locality{}, rpcbase.DefaultClass).Health()
+	}
+	return n.rpcContext.DRPCDialPod(addr, id, roachpb.Locality{}, rpcbase.DefaultClass).Health()
 }
 
 // GetCircuitBreaker retrieves the circuit breaker for connections to the