rpc: remove DRPC dependency from Peer and Connection generics

Although `Peer` and `Connection` are generic and intended to support both
gRPC and DRPC connections, the current implementation has a hardcoded
dependency on DRPC. As a result, a DRPC connection is always dialed,
regardless of the intended type. This PR removes the direct dependency
on DRPC, allowing the appropriate connection type (gRPC or DRPC) to be
used based on the generic parameter.

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

Author: cthumuluru-crdb

pkg/kv/kvpb/api.proto

@@ -3755,6 +3755,11 @@ service Cluster {
   rpc Join (JoinNodeRequest) returns (JoinNodeResponse) {}
 }
 
+// RangeFeed service implemented by nodes for KV API requests.
+service RangeFeed {
+  rpc MuxRangeFeed (stream RangeFeedRequest) returns (stream MuxRangeFeedEvent) {}
+}
+
 // Batch and RangeFeed service implemented by nodes for KV API requests.
 service Internal {
   rpc Batch (BatchRequest) returns (BatchResponse) {}
@@ -3882,9 +3887,4 @@ message ScanStats {
 
 // UsedFollowerRead indicates whether at least some reads were served by the
 // follower replicas.
-message UsedFollowerRead {}
-
-// RangeFeed service implemented by nodes for KV API requests.
-service RangeFeed {
-  rpc MuxRangeFeed       (stream RangeFeedRequest)   returns (stream MuxRangeFeedEvent)       {}
-}
+message UsedFollowerRead {}

pkg/rpc/connection.go

@@ -15,7 +15,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
-	"storj.io/drpc"
 )
 
 // rpcConn defines a lightweight interface that both grpc.ClientConn and drpc.Conn
@@ -91,9 +90,6 @@ type Connection[Conn rpcConn] struct {
 	//
 	// The pool is only initialized once the rpc connection is resolved.
 	batchStreamPool streamPool[*kvpb.BatchRequest, *kvpb.BatchResponse, Conn]
-	// TODO(server): remove this once the code is consolidated to use generic
-	// Connection and Pool for drpc.Conn.
-	drpcBatchStreamPool DRPCBatchStreamPool
 }
 
 // newConnectionToNodeID makes a Connection for the given node, class, and nontrivial Signal
@@ -104,15 +100,13 @@ func newConnectionToNodeID[Conn rpcConn](
 	breakerSignal func() circuit.Signal,
 	connOptions *ConnectionOptions[Conn],
 ) *Connection[Conn] {
-	drpcConnEquals := func(a, b drpc.Conn) bool { return a == b }
 	c := &Connection[Conn]{
 		breakerSignalFn: breakerSignal,
 		k:               k,
 		connFuture: connFuture[Conn]{
 			ready: make(chan struct{}),
 		},
-		batchStreamPool:     makeStreamPool(opts.Stopper, connOptions.newBatchStreamClient, connOptions.connEquals),
-		drpcBatchStreamPool: makeStreamPool(opts.Stopper, newDRPCBatchStream, drpcConnEquals),
+		batchStreamPool: makeStreamPool(opts.Stopper, connOptions.newBatchStreamClient, connOptions.connEquals),
 	}
 	return c
 }
@@ -124,15 +118,15 @@ func newConnectionToNodeID[Conn rpcConn](
 // block but fall back to defErr in this case.
 func (c *Connection[Conn]) waitOrDefault(
 	ctx context.Context, defErr error, sig circuit.Signal,
-) (Conn, drpc.Conn, error) {
+) (Conn, error) {
 	// Check the circuit breaker first. If it is already tripped now, we
 	// want it to take precedence over connFuture below (which is closed in
 	// the common case of a connection going bad after having been healthy
 	// for a while).
 	var nilConn Conn
 	select {
 	case <-sig.C():
-		return nilConn, nil, sig.Err()
+		return nilConn, sig.Err()
 	default:
 	}
 
@@ -143,26 +137,26 @@ func (c *Connection[Conn]) waitOrDefault(
 		select {
 		case <-c.connFuture.C():
 		case <-sig.C():
-			return nilConn, nil, sig.Err()
+			return nilConn, sig.Err()
 		case <-ctx.Done():
-			return nilConn, nil, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
+			return nilConn, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
 		}
 	} else {
 		select {
 		case <-c.connFuture.C():
 		case <-sig.C():
-			return nilConn, nil, sig.Err()
+			return nilConn, sig.Err()
 		case <-ctx.Done():
-			return nilConn, nil, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
+			return nilConn, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
 		default:
-			return nilConn, nil, defErr
+			return nilConn, defErr
 		}
 	}
 
 	// Done waiting, c.connFuture has resolved, return the result. Note that this
 	// conn could be unhealthy (or there may not even be a conn, i.e. Err() !=
 	// nil), if that's what the caller wanted (ConnectNoBreaker).
-	return c.connFuture.Conn(), c.connFuture.DRPCConn(), c.connFuture.Err()
+	return c.connFuture.Conn(), c.connFuture.Err()
 }
 
 // Connect returns the underlying rpc connection after it has been validated,
@@ -172,18 +166,10 @@ func (c *Connection[Conn]) waitOrDefault(
 // an error. In rare cases, this behavior is undesired and ConnectNoBreaker may
 // be used instead.
 func (c *Connection[Conn]) Connect(ctx context.Context) (Conn, error) {
-	cc, _, err := c.waitOrDefault(ctx, nil /* defErr */, c.breakerSignalFn())
+	cc, err := c.waitOrDefault(ctx, nil /* defErr */, c.breakerSignalFn())
 	return cc, err
 }
 
-// ConnectEx is similar to Connect but it addition to gRPC connection, it also
-// returns underlying drpc connection after it has been validated.
-// TODO(server): remove this once the code is consolidated to use generic
-// Connection and Pool for drpc.Conn.
-func (c *Connection[Conn]) ConnectEx(ctx context.Context) (Conn, drpc.Conn, error) {
-	return c.waitOrDefault(ctx, nil /* defErr */, c.breakerSignalFn())
-}
-
 type neverTripSignal struct{}
 
 func (s *neverTripSignal) Err() error {
@@ -202,7 +188,7 @@ func (s *neverTripSignal) IsTripped() bool {
 // that it will latch onto (or start) an existing connection attempt even if
 // previous attempts have not succeeded. This may be preferable to Connect
 // if the caller is already certain that a peer is available.
-func (c *Connection[Conn]) ConnectNoBreaker(ctx context.Context) (Conn, drpc.Conn, error) {
+func (c *Connection[Conn]) ConnectNoBreaker(ctx context.Context) (Conn, error) {
 	// For ConnectNoBreaker we don't use the default Signal but pass a dummy one
 	// that never trips. (The probe tears down the Conn on quiesce so we don't rely
 	// on the Signal for that).
@@ -226,7 +212,7 @@ func (c *Connection[Conn]) ConnectNoBreaker(ctx context.Context) (Conn, drpc.Con
 // latest heartbeat. Returns ErrNotHeartbeated if the peer was just contacted for
 // the first time and the first heartbeat has not occurred yet.
 func (c *Connection[Conn]) Health() error {
-	_, _, err := c.waitOrDefault(context.Background(), ErrNotHeartbeated, c.breakerSignalFn())
+	_, err := c.waitOrDefault(context.Background(), ErrNotHeartbeated, c.breakerSignalFn())
 	return err
 }
 
@@ -241,17 +227,9 @@ func (c *Connection[Conn]) BatchStreamPool() *streamPool[*kvpb.BatchRequest, *kv
 	return &c.batchStreamPool
 }
 
-func (c *Connection[Conn]) DRPCBatchStreamPool() *DRPCBatchStreamPool {
-	if !c.connFuture.Resolved() {
-		panic("DRPCBatchStreamPool called on unresolved connection")
-	}
-	return &c.drpcBatchStreamPool
-}
-
 type connFuture[Conn rpcConn] struct {
 	ready chan struct{}
 	cc    Conn
-	dc    drpc.Conn
 	err   error
 }
 
@@ -279,14 +257,6 @@ func (s *connFuture[Conn]) Conn() Conn {
 	return s.cc
 }
 
-// DRPCConn must only be called after C() has been closed.
-func (s *connFuture[Conn]) DRPCConn() drpc.Conn {
-	if s.err != nil {
-		return nil
-	}
-	return s.dc
-}
-
 func (s *connFuture[Conn]) Resolved() bool {
 	select {
 	case <-s.ready:
@@ -298,12 +268,12 @@ func (s *connFuture[Conn]) Resolved() bool {
 
 // Resolve is idempotent. Only the first call has any effect.
 // Not thread safe.
-func (s *connFuture[Conn]) Resolve(cc Conn, dc drpc.Conn, err error) {
+func (s *connFuture[Conn]) Resolve(cc Conn, err error) {
 	select {
 	case <-s.ready:
 		// Already resolved, noop.
 	default:
-		s.cc, s.dc, s.err = cc, dc, err
+		s.cc, s.err = cc, err
 		close(s.ready)
 	}
 }

pkg/rpc/drpc.go

@@ -41,7 +41,7 @@ func (d *drpcCloseNotifier) CloseNotify(ctx context.Context) <-chan struct{} {
 	return d.conn.Closed()
 }
 
-// TODO DB Server: unexport this once dial methods are added in rpccontext.
+// TODO(server): unexport this once dial methods are added in rpccontext.
 func DialDRPC(
 	rpcCtx *Context,
 ) func(ctx context.Context, target string, _ rpcbase.ConnectionClass) (drpc.Conn, error) {
@@ -211,64 +211,12 @@ func newDRPCPeerOptions(
 		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
-			},
+			dial: DialDRPC(rpcCtx),
 			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)
+				return kvpb.NewDRPCKVBatchClientAdapter(cc).BatchStream(ctx)
 			},
 			newCloseNotifier: func(_ *stop.Stopper, cc drpc.Conn) closeNotifier {
 				return &drpcCloseNotifier{conn: cc}