rpc: add DRPC dial methods

This commit adds DRPC dial methods to `rpc.nodedialer` and `rpc.Context`.
In the future commits, helper functions for creation RPC clients are
update to create gRPC or DRPC clients base on `rpc.experimental_drpc.enabled`
cluster setting.

Epic: CRDB-48923
Fixes: #148383
Release note: none

Author: cthumuluru-crdb

pkg/rpc/context.go

@@ -51,6 +51,7 @@ import (
 	"google.golang.org/grpc/encoding"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
+	"storj.io/drpc"
 )
 
 // NewServer sets up an RPC server. Depending on the ServerOptions, the Server
@@ -226,7 +227,10 @@ type Context struct {
 
 	localInternalClient RestrictedInternalClient
 
+	// peers is a map of gRPC connections to other nodes.
 	peers peerMap[*grpc.ClientConn]
+	// drpcPeers is a map of DRPC connections to other nodes.
+	drpcPeers peerMap[drpc.Conn]
 
 	// dialbackMap is a map of currently executing dialback connections. This map
 	// is typically empty or close to empty. It only holds entries that are being
@@ -1998,6 +2002,16 @@ func (rpcCtx *Context) GRPCUnvalidatedDial(
 	return rpcCtx.grpcDialNodeInternal(target, 0, locality, rpcbase.SystemClass)
 }
 
+// DRPCUnvalidatedDial uses DRPCDialNode and disables validation of the
+// node ID between client and server. This function should only be
+// used with the gossip client and CLI commands which can talk to any
+// node. This method implies a SystemClass.
+func (rpcCtx *Context) DRPCUnvalidatedDial(
+	target string, locality roachpb.Locality,
+) *DRPCConnection {
+	return rpcCtx.drpcDialNodeInternal(target, 0, locality, rpcbase.SystemClass)
+}
+
 // GRPCDialNode calls grpc.Dial with options appropriate for the
 // context and class (see the comment on ConnectionClass).
 //
@@ -2019,6 +2033,27 @@ func (rpcCtx *Context) GRPCDialNode(
 	return rpcCtx.grpcDialNodeInternal(target, remoteNodeID, remoteLocality, class)
 }
 
+// DRPCDialNode dials a DRPC connection with options appropriate for the
+// context and class (see the comment on ConnectionClass).
+//
+// The remoteNodeID becomes a constraint on the expected node ID of
+// the remote node; this is checked during heartbeats. The caller is
+// responsible for ensuring the remote node ID is known prior to using
+// this function.
+func (rpcCtx *Context) DRPCDialNode(
+	target string,
+	remoteNodeID roachpb.NodeID,
+	remoteLocality roachpb.Locality,
+	class rpcbase.ConnectionClass,
+) *DRPCConnection {
+	if remoteNodeID == 0 && !rpcCtx.TestingAllowNamedRPCToAnonymousServer {
+		log.Fatalf(
+			rpcCtx.makeDialCtx(target, remoteNodeID, class),
+			"%v", errors.AssertionFailedf("invalid node ID 0 in DRPCDialNode()"))
+	}
+	return rpcCtx.drpcDialNodeInternal(target, remoteNodeID, remoteLocality, class)
+}
+
 // GRPCDialPod wraps GRPCDialNode and treats the `remoteInstanceID`
 // argument as a `NodeID` which it converts. This works because the
 // tenant gRPC server is initialized using the `InstanceID` so it
@@ -2035,6 +2070,22 @@ func (rpcCtx *Context) GRPCDialPod(
 	return rpcCtx.GRPCDialNode(target, roachpb.NodeID(remoteInstanceID), remoteLocality, class)
 }
 
+// DRPCDialPod wraps DRPCDialNode and treats the `remoteInstanceID`
+// argument as a `NodeID` which it converts. This works because the
+// tenant DRPC server is initialized using the `InstanceID` so it
+// accepts our connection as matching the ID we're dialing.
+//
+// Since DRPCDialNode accepts a separate `target` and `NodeID` it
+// requires no further modification to work between pods.
+func (rpcCtx *Context) DRPCDialPod(
+	target string,
+	remoteInstanceID base.SQLInstanceID,
+	remoteLocality roachpb.Locality,
+	class rpcbase.ConnectionClass,
+) *DRPCConnection {
+	return rpcCtx.DRPCDialNode(target, roachpb.NodeID(remoteInstanceID), remoteLocality, class)
+}
+
 // grpcDialNodeInternal connects to the remote node and sets up the async heartbeater.
 // This intentionally takes no `context.Context`; it uses one derived from rpcCtx.masterCtx.
 func (rpcCtx *Context) grpcDialNodeInternal(
@@ -2050,8 +2101,37 @@ func (rpcCtx *Context) grpcDialNodeInternal(
 	}
 
 	// Slow path. Race to create a peer.
-	conns := &rpcCtx.peers
+	return rpcDialNodeInternal(
+		rpcCtx, k, newGRPCPeerOptions(rpcCtx, k, remoteLocality),
+	)
+}
+
+// drpcDialNodeInternal connects to the remote node and sets up the async heartbeater.
+// This intentionally takes no `context.Context`; it uses one derived from rpcCtx.masterCtx.
+func (rpcCtx *Context) drpcDialNodeInternal(
+	target string,
+	remoteNodeID roachpb.NodeID,
+	remoteLocality roachpb.Locality,
+	class rpcbase.ConnectionClass,
+) *DRPCConnection {
+	k := peerKey{TargetAddr: target, NodeID: remoteNodeID, Class: class}
+	if p, ok := rpcCtx.drpcPeers.get(k); ok {
+		// There's a cached peer, so we have a cached connection, use it.
+		return p.c
+	}
+
+	// Slow path. Race to create a peer.
+	return rpcDialNodeInternal(
+		rpcCtx, k, newDRPCPeerOptions(rpcCtx, k, remoteLocality),
+	)
+}
 
+// rpcDialNodeInternal is used to dial a new connection to the peer if one
+// doesn't already exist.
+func rpcDialNodeInternal[Conn rpcConn](
+	rpcCtx *Context, k peerKey, peerOpts *peerOptions[Conn],
+) *Connection[Conn] {
+	conns := peerOpts.peers
 	conns.mu.Lock()
 	defer conns.mu.Unlock()
 
@@ -2060,12 +2140,11 @@ func (rpcCtx *Context) grpcDialNodeInternal(
 	}
 
 	// Won race. Actually create a peer.
-
 	if conns.mu.m == nil {
-		conns.mu.m = map[peerKey]*peer[*grpc.ClientConn]{}
+		conns.mu.m = map[peerKey]*peer[Conn]{}
 	}
 
-	p := newPeer(rpcCtx, k, newGRPCPeerOptions(rpcCtx, k, remoteLocality))
+	p := newPeer(rpcCtx, k, peerOpts)
 	// (Asynchronously) Start the probe (= heartbeat loop). The breaker is healthy
 	// right now (it was just created) but the call to `.Probe` will launch the
 	// probe[1] regardless.

pkg/rpc/drpc.go

@@ -12,8 +12,11 @@ import (
 	"net"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/errors"
 	"storj.io/drpc"
 	"storj.io/drpc/drpcconn"
@@ -29,6 +32,14 @@ import (
 // Default idle connection timeout for DRPC connections in the pool.
 var defaultDRPCConnIdleTimeout = 5 * time.Minute
 
+type drpcCloseNotifier struct {
+	conn drpc.Conn
+}
+
+func (d *drpcCloseNotifier) CloseNotify(ctx context.Context) <-chan struct{} {
+	return d.conn.Closed()
+}
+
 func dialDRPC(
 	rpcCtx *Context,
 ) func(ctx context.Context, target string, _ rpcbase.ConnectionClass) (drpc.Conn, error) {
@@ -174,3 +185,82 @@ func (srv *drpcOffServer) Serve(_ context.Context, lis net.Listener) error {
 func (srv *drpcOffServer) Register(interface{}, drpc.Description) error {
 	return nil
 }
+
+// newDRPDCPeerOptions creates peerOptions for DRPC peers.
+func newDRPCPeerOptions(
+	rpcCtx *Context, k peerKey, locality roachpb.Locality,
+) *peerOptions[drpc.Conn] {
+	pm, _ := rpcCtx.metrics.acquire(k, locality)
+	return &peerOptions[drpc.Conn]{
+		locality: locality,
+		peers:    &rpcCtx.drpcPeers,
+		connOptions: &ConnectionOptions[drpc.Conn]{
+			dial: func(ctx context.Context, target string, _ rpcbase.ConnectionClass) (drpc.Conn, error) {
+				// TODO(server): could use connection class instead of empty key here.
+				pool := drpcpool.New[struct{}, drpcpool.Conn](drpcpool.Options{
+					Expiration: defaultDRPCConnIdleTimeout,
+				})
+				pooledConn := pool.Get(ctx /* unused */, struct{}{}, func(ctx context.Context,
+					_ struct{}) (drpcpool.Conn, error) {
+
+					netConn, err := drpcmigrate.DialWithHeader(ctx, "tcp", target, drpcmigrate.DRPCHeader)
+					if err != nil {
+						return nil, err
+					}
+
+					opts := drpcconn.Options{
+						Manager: drpcmanager.Options{
+							Reader: drpcwire.ReaderOptions{
+								MaximumBufferSize: math.MaxInt,
+							},
+							Stream: drpcstream.Options{
+								MaximumBufferSize: 0, // unlimited
+							},
+							SoftCancel: true, // don't close the transport when stream context is canceled
+						},
+					}
+					var conn *drpcconn.Conn
+					if rpcCtx.ContextOptions.Insecure {
+						conn = drpcconn.NewWithOptions(netConn, opts)
+					} else {
+						tlsConfig, err := rpcCtx.GetClientTLSConfig()
+						if err != nil {
+							return nil, err
+						}
+						// Clone TLS config to avoid modifying a cached TLS config.
+						tlsConfig = tlsConfig.Clone()
+						// TODO(server): remove this hack which is necessary at least in
+						// testing to get TestDRPCSelectQuery to pass.
+						tlsConfig.InsecureSkipVerify = true
+						tlsConn := tls.Client(netConn, tlsConfig)
+						conn = drpcconn.NewWithOptions(tlsConn, opts)
+					}
+
+					return conn, nil
+				})
+				// `pooledConn.Close` doesn't tear down any of the underlying TCP
+				// connections but simply marks the pooledConn handle as returning
+				// errors. When we "close" this conn, we want to tear down all of
+				// the connections in the pool (in effect mirroring the behavior of
+				// gRPC where a single conn is shared).
+				return &closeEntirePoolConn{
+					Conn: pooledConn,
+					pool: pool,
+				}, nil
+			},
+			connEquals: func(a, b drpc.Conn) bool {
+				return a == b
+			},
+			newBatchStreamClient: func(ctx context.Context, cc drpc.Conn) (BatchStreamClient, error) {
+				return kvpb.NewDRPCInternalClientAdapter(cc).BatchStream(ctx)
+			},
+			newCloseNotifier: func(_ *stop.Stopper, cc drpc.Conn) closeNotifier {
+				return &drpcCloseNotifier{conn: cc}
+			},
+		},
+		newHeartbeatClient: func(cc drpc.Conn) RPCHeartbeatClient {
+			return NewDRPCHeartbeatClientAdapter(cc)
+		},
+		pm: pm,
+	}
+}

pkg/rpc/nodedialer/nodedialer.go

@@ -106,6 +106,28 @@ func (n *Dialer) Dial(
 	return conn, err
 }
 
+// DRPCDial returns a drpc connection to the given node. It logs whenever the
+// node first becomes unreachable or reachable. This method is similar to
+// Dial, but it dials a DRPC connection instead of a gRPC connection.
+func (n *Dialer) DRPCDial(
+	ctx context.Context, nodeID roachpb.NodeID, class rpcbase.ConnectionClass,
+) (drpc.Conn, error) {
+	if n == nil || n.resolver == nil {
+		return nil, errors.New("no node dialer configured")
+	}
+	// Don't trip the breaker if we're already canceled.
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		return nil, errors.Wrap(ctxErr, "dial")
+	}
+	addr, locality, err := n.resolver(nodeID)
+	if err != nil {
+		err = errors.Wrapf(err, "failed to resolve n%d", nodeID)
+		return nil, err
+	}
+	conn, _, err := n.drpcDial(ctx, nodeID, addr, locality, true, class)
+	return conn, err
+}
+
 // DialNoBreaker is like Dial, but will not check the circuit breaker before
 // trying to connect. This function should only be used when there is good
 // reason to believe that the node is reachable.
@@ -123,6 +145,23 @@ func (n *Dialer) DialNoBreaker(
 	return conn, err
 }
 
+// DRPCDialNoBreaker is like DRPCDial, but will not check the circuit breaker
+// before trying to connect. This function should only be used when there is
+// good reason to believe that the node is reachable.
+func (n *Dialer) DRPCDialNoBreaker(
+	ctx context.Context, nodeID roachpb.NodeID, class rpcbase.ConnectionClass,
+) (drpc.Conn, error) {
+	if n == nil || n.resolver == nil {
+		return nil, errors.New("no node dialer configured")
+	}
+	addr, locality, err := n.resolver(nodeID)
+	if err != nil {
+		return nil, err
+	}
+	conn, _, err := n.drpcDial(ctx, nodeID, addr, locality, false, class)
+	return conn, err
+}
+
 // DialInternalClient is a specialization of DialClass for callers that
 // want a kvpb.InternalClient. This supports an optimization to bypass the
 // network for the local node.
@@ -213,6 +252,41 @@ func (n *Dialer) dial(
 	return conn, pool, dconn, drpcStreamPool, nil
 }
 
+// drcpDial performs the dialing of the remote connection. If checkBreaker
+// is set (which it usually is), circuit breakers for the peer will be
+// checked. This method is similar to dial, but it dials a DRPC
+// connection instead of a gRPC connection.
+func (n *Dialer) drpcDial(
+	ctx context.Context,
+	nodeID roachpb.NodeID,
+	addr net.Addr,
+	locality roachpb.Locality,
+	checkBreaker bool,
+	class rpcbase.ConnectionClass,
+) (drpc.Conn, *rpc.DRPCBatchStreamPool, error) {
+	const ctxWrapMsg = "drpcDial"
+	// Don't trip the breaker if we're already canceled.
+	if ctxErr := ctx.Err(); ctxErr != nil {
+		return nil, nil, errors.Wrap(ctxErr, ctxWrapMsg)
+	}
+	rpcConn := n.rpcContext.DRPCDialNode(addr.String(), nodeID, locality, class)
+	connect := rpcConn.ConnectEx
+	if !checkBreaker {
+		connect = rpcConn.ConnectNoBreaker
+	}
+	_, dconn, err := connect(ctx)
+	if err != nil {
+		// If we were canceled during the dial, don't trip the breaker.
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return nil, nil, errors.Wrap(ctxErr, ctxWrapMsg)
+		}
+		err = errors.Wrapf(err, "failed to connect to n%d at %v", nodeID, addr)
+		return nil, nil, err
+	}
+	drpcStreamPool := rpcConn.DRPCBatchStreamPool()
+	return dconn, drpcStreamPool, nil
+}
+
 // ConnHealth returns nil if we have an open connection of the request
 // class to the given node that succeeded on its most recent heartbeat.
 // Returns circuit.ErrBreakerOpen if the breaker is tripped, otherwise