rpc: refactor Peer to support DRPC and gRPC conns

Refactor Peer further to make it easy to support both gRPC and DRPC
peers. In this change, gRPC Peer continues to dial DRPC connections to keep
the changes backward compatible. As part of enabling DRPC support across
the codebase, future commits will remove the DRPC-specific code and
replace it with DRPC generic peer.

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

Author: cthumuluru-crdb

pkg/rpc/connection.go

@@ -10,6 +10,7 @@ import (
 	"io"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
 	"github.com/cockroachdb/cockroach/pkg/util/circuit"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/errors"
@@ -21,10 +22,13 @@ import (
 // must implement. It is used as a type constraint for rpc connections and allows
 // the Connection and Peer structs to work seamlessly with both gRPC and DRPC
 // connections.
-type rpcConn interface {
-	io.Closer
-	comparable
-}
+type rpcConn io.Closer
+
+// dialFunc is a generic function type used to establish an RPC connection.
+// It takes a context for cancellation/timeouts, a target string (e.g., address),
+// and a ConnectionClass to categorize the connection's purpose or priority.
+// It returns the established connection (of type Conn) or an error if dialing fails.
+type dialFunc[Conn rpcConn] func(ctx context.Context, target string, class rpcbase.ConnectionClass) (Conn, error)
 
 // heartbeatClientConstructor is a function type that creates a HeartbeatClient
 // for a given rpc connection. This allows us to use different implementations of
@@ -37,6 +41,23 @@ type closeNotifier interface {
 	CloseNotify(ctx context.Context) <-chan struct{}
 }
 
+// ConnectionOptions allow for customization of connection behaviors such as:
+//   - Establishing a new connection.
+//   - Client constuctors for clients like batch streams.
+//   - Comparing two connections for equivalence.
+type ConnectionOptions[Conn rpcConn] struct {
+	// dial function to open a new connection.
+	dial dialFunc[Conn]
+	// connEquals defines the equivalence function for two RPC connections.
+	connEquals equalsFunc[Conn]
+	// newBatchStreamClient is a constructor function for creating a new batch
+	// stream client associated with a specific RPC connection.
+	newBatchStreamClient streamConstructor[*kvpb.BatchRequest, *kvpb.BatchResponse, Conn]
+	// newCloseNotifier is a constructor function for creating a new
+	// closeNotifier associated with a specific RPC connection.
+	newCloseNotifier closeNotifierConstructor[Conn]
+}
+
 // closeNotifierConstructor is a function type that creates a closeNotifier
 // for a given rpc connection. This allows us to use different implementations of
 // closeNotifier for different types of connections (e.g., gRPC and DRPC).
@@ -81,16 +102,17 @@ func newConnectionToNodeID[Conn rpcConn](
 	opts *ContextOptions,
 	k peerKey,
 	breakerSignal func() circuit.Signal,
-	newBatchStreamClient streamConstructor[*kvpb.BatchRequest, *kvpb.BatchResponse, Conn],
+	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, newBatchStreamClient),
-		drpcBatchStreamPool: makeStreamPool(opts.Stopper, newDRPCBatchStream),
+		batchStreamPool:     makeStreamPool(opts.Stopper, connOptions.newBatchStreamClient, connOptions.connEquals),
+		drpcBatchStreamPool: makeStreamPool(opts.Stopper, newDRPCBatchStream, drpcConnEquals),
 	}
 	return c
 }
@@ -107,10 +129,10 @@ func (c *Connection[Conn]) waitOrDefault(
 	// 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 cc Conn
+	var nilConn Conn
 	select {
 	case <-sig.C():
-		return cc, nil, sig.Err()
+		return nilConn, nil, sig.Err()
 	default:
 	}
 
@@ -121,19 +143,19 @@ func (c *Connection[Conn]) waitOrDefault(
 		select {
 		case <-c.connFuture.C():
 		case <-sig.C():
-			return cc, nil, sig.Err()
+			return nilConn, nil, sig.Err()
 		case <-ctx.Done():
-			return cc, nil, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
+			return nilConn, nil, 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 cc, nil, sig.Err()
+			return nilConn, nil, sig.Err()
 		case <-ctx.Done():
-			return cc, nil, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
+			return nilConn, nil, errors.Wrapf(ctx.Err(), "while connecting to n%d at %s", c.k.NodeID, c.k.TargetAddr)
 		default:
-			return cc, nil, defErr
+			return nilConn, nil, defErr
 		}
 	}
 

pkg/rpc/context.go

@@ -2065,7 +2065,7 @@ func (rpcCtx *Context) grpcDialNodeInternal(
 		conns.mu.m = map[peerKey]*peer[*grpc.ClientConn]{}
 	}
 
-	p := rpcCtx.newPeer(k, remoteLocality)
+	p := newPeer(rpcCtx, k, newGRPCPeerOptions(rpcCtx, k, remoteLocality))
 	// (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/grpc.go

@@ -8,6 +8,9 @@ package rpc
 import (
 	"context"
 
+	"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/stop"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/connectivity"
@@ -47,3 +50,34 @@ func (g *grpcCloseNotifier) CloseNotify(ctx context.Context) <-chan struct{} {
 }
 
 type GRPCConnection = Connection[*grpc.ClientConn]
+
+// newGRPCPeerOptions creates peerOptions for gRPC peers.
+func newGRPCPeerOptions(
+	rpcCtx *Context, k peerKey, locality roachpb.Locality,
+) *peerOptions[*grpc.ClientConn] {
+	pm, lm := rpcCtx.metrics.acquire(k, locality)
+	return &peerOptions[*grpc.ClientConn]{
+		locality: locality,
+		peers:    &rpcCtx.peers,
+		connOptions: &ConnectionOptions[*grpc.ClientConn]{
+			dial: func(ctx context.Context, target string, class rpcbase.ConnectionClass) (*grpc.ClientConn, error) {
+				additionalDialOpts := []grpc.DialOption{grpc.WithStatsHandler(&statsTracker{lm})}
+				additionalDialOpts = append(additionalDialOpts, rpcCtx.testingDialOpts...)
+				return rpcCtx.grpcDialRaw(ctx, target, class, additionalDialOpts...)
+			},
+			connEquals: func(a, b *grpc.ClientConn) bool {
+				return a == b
+			},
+			newBatchStreamClient: func(ctx context.Context, cc *grpc.ClientConn) (BatchStreamClient, error) {
+				return kvpb.NewInternalClient(cc).BatchStream(ctx)
+			},
+			newCloseNotifier: func(stopper *stop.Stopper, cc *grpc.ClientConn) closeNotifier {
+				return &grpcCloseNotifier{stopper: stopper, conn: cc}
+			},
+		},
+		newHeartbeatClient: func(cc *grpc.ClientConn) RPCHeartbeatClient {
+			return NewGRPCHeartbeatClientAdapter(cc)
+		},
+		pm: pm,
+	}
+}

pkg/rpc/peer.go

@@ -14,18 +14,15 @@ import (
 	"github.com/VividCortex/ewma"
 	"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/circuit"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/netutil"
-	"github.com/cockroachdb/cockroach/pkg/util/stop"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/logtags"
 	"github.com/cockroachdb/redact"
-	"google.golang.org/grpc"
 	"google.golang.org/grpc/status"
 	"storj.io/drpc"
 )
@@ -121,15 +118,13 @@ func (p *peer[Conn]) releaseMetricsLocked() {
 // See (*peer).launch for details on the probe (heartbeat loop) itself.
 type peer[Conn rpcConn] struct {
 	peerMetrics
-	k                    peerKey
-	opts                 *ContextOptions
-	heartbeatInterval    time.Duration
-	heartbeatTimeout     time.Duration
-	dial                 func(ctx context.Context, target string, class rpcbase.ConnectionClass) (Conn, error)
-	dialDRPC             func(ctx context.Context, target string, class rpcbase.ConnectionClass) (drpc.Conn, error)
-	newHeartbeatClient   heartbeatClientConstructor[Conn]
-	newBatchStreamClient streamConstructor[*kvpb.BatchRequest, *kvpb.BatchResponse, Conn]
-	newCloseNotifier     closeNotifierConstructor[Conn]
+	k                  peerKey
+	opts               *ContextOptions
+	newHeartbeatClient heartbeatClientConstructor[Conn]
+	heartbeatInterval  time.Duration
+	heartbeatTimeout   time.Duration
+	connOptions        *ConnectionOptions[Conn]
+	drpcDial           dialFunc[drpc.Conn]
 	// b maintains connection health. This breaker's async probe is always
 	// active - it is the heartbeat loop and manages `mu.c.` (including
 	// recreating it after the connection fails and has to be redialed).
@@ -210,6 +205,14 @@ func (p *peer[Conn]) snap() PeerSnap[Conn] {
 	return p.mu.PeerSnap
 }
 
+type peerOptions[Conn rpcConn] struct {
+	locality           roachpb.Locality
+	pm                 peerMetrics
+	newHeartbeatClient heartbeatClientConstructor[Conn]
+	connOptions        *ConnectionOptions[Conn]
+	peers              *peerMap[Conn]
+}
+
 // newPeer returns circuit breaker that trips when connection (associated
 // with provided peerKey) is failed. The breaker's probe *is* the heartbeat loop
 // and is thus running at all times. The exception is a decommissioned node, for
@@ -229,39 +232,26 @@ func (p *peer[Conn]) snap() PeerSnap[Conn] {
 // map, the next attempt to dial the node will start from a blank slate. In
 // other words, even with this theoretical race, the situation will sort itself
 // out quickly.
-func (rpcCtx *Context) newPeer(k peerKey, locality roachpb.Locality) *peer[*grpc.ClientConn] {
+func newPeer[Conn rpcConn](rpcCtx *Context, k peerKey, peerOpts *peerOptions[Conn]) *peer[Conn] {
 	// Initialization here is a bit circular. The peer holds the breaker. The
 	// breaker probe references the peer because it needs to replace the one-shot
 	// Connection when it makes a new connection in the probe. And (all but the
 	// first incarnation of) the Connection also holds on to the breaker since the
 	// Connect method needs to do the short-circuiting (if a Connection is created
 	// while the breaker is tripped, we want to block in Connect only once we've
 	// seen the first heartbeat succeed).
-	pm, lm := rpcCtx.metrics.acquire(k, locality)
-	p := &peer[*grpc.ClientConn]{
-		peerMetrics:        pm,
+	p := &peer[Conn]{
+		peerMetrics:        peerOpts.pm,
 		logDisconnectEvery: log.Every(time.Minute),
 		k:                  k,
 		remoteClocks:       rpcCtx.RemoteClocks,
 		opts:               &rpcCtx.ContextOptions,
-		peers:              &rpcCtx.peers,
-		dial: func(ctx context.Context, target string, class rpcbase.ConnectionClass) (*grpc.ClientConn, error) {
-			additionalDialOpts := []grpc.DialOption{grpc.WithStatsHandler(&statsTracker{lm})}
-			additionalDialOpts = append(additionalDialOpts, rpcCtx.testingDialOpts...)
-			return rpcCtx.grpcDialRaw(ctx, target, class, additionalDialOpts...)
-		},
-		dialDRPC: dialDRPC(rpcCtx),
-		newHeartbeatClient: func(cc *grpc.ClientConn) RPCHeartbeatClient {
-			return NewGRPCHeartbeatClientAdapter(cc)
-		},
-		newBatchStreamClient: func(ctx context.Context, cc *grpc.ClientConn) (BatchStreamClient, error) {
-			return kvpb.NewInternalClient(cc).BatchStream(ctx)
-		},
-		newCloseNotifier: func(stopper *stop.Stopper, cc *grpc.ClientConn) closeNotifier {
-			return &grpcCloseNotifier{stopper: stopper, conn: cc}
-		},
-		heartbeatInterval: rpcCtx.RPCHeartbeatInterval,
-		heartbeatTimeout:  rpcCtx.RPCHeartbeatTimeout,
+		peers:              peerOpts.peers,
+		connOptions:        peerOpts.connOptions,
+		drpcDial:           dialDRPC(rpcCtx),
+		newHeartbeatClient: peerOpts.newHeartbeatClient,
+		heartbeatInterval:  rpcCtx.RPCHeartbeatInterval,
+		heartbeatTimeout:   rpcCtx.RPCHeartbeatTimeout,
 	}
 	var b *circuit.Breaker
 
@@ -275,8 +265,8 @@ func (rpcCtx *Context) newPeer(k peerKey, locality roachpb.Locality) *peer[*grpc
 		},
 	})
 	p.b = b
-	c := newConnectionToNodeID(p.opts, k, b.Signal, p.newBatchStreamClient)
-	p.mu.PeerSnap = PeerSnap[*grpc.ClientConn]{c: c}
+	c := newConnectionToNodeID(p.opts, k, b.Signal, p.connOptions)
+	p.mu.PeerSnap = PeerSnap[Conn]{c: c}
 
 	return p
 }
@@ -375,7 +365,7 @@ func (p *peer[Conn]) run(ctx context.Context, report func(error), done func()) {
 		func() {
 			p.mu.Lock()
 			defer p.mu.Unlock()
-			p.mu.c = newConnectionToNodeID(p.opts, p.k, p.mu.c.breakerSignalFn, p.newBatchStreamClient)
+			p.mu.c = newConnectionToNodeID(p.opts, p.k, p.mu.c.breakerSignalFn, p.connOptions)
 		}()
 
 		if p.snap().deleteAfter != 0 {
@@ -388,14 +378,14 @@ func (p *peer[Conn]) run(ctx context.Context, report func(error), done func()) {
 }
 
 func (p *peer[Conn]) runOnce(ctx context.Context, report func(error)) error {
-	cc, err := p.dial(ctx, p.k.TargetAddr, p.k.Class)
+	cc, err := p.connOptions.dial(ctx, p.k.TargetAddr, p.k.Class)
 	if err != nil {
 		return err
 	}
 	defer func() {
 		_ = cc.Close() // nolint:grpcconnclose
 	}()
-	dc, err := p.dialDRPC(ctx, p.k.TargetAddr, p.k.Class)
+	dc, err := p.drpcDial(ctx, p.k.TargetAddr, p.k.Class)
 	if err != nil {
 		return err
 	}
@@ -406,7 +396,7 @@ func (p *peer[Conn]) runOnce(ctx context.Context, report func(error)) error {
 	// Set up notifications on a channel when gRPC tears down, so that we
 	// can trigger another instant heartbeat for expedited circuit breaker
 	// tripping.
-	connClosedCh := p.newCloseNotifier(p.opts.Stopper, cc).CloseNotify(ctx)
+	connClosedCh := p.connOptions.newCloseNotifier(p.opts.Stopper, cc).CloseNotify(ctx)
 
 	if p.remoteClocks != nil {
 		p.remoteClocks.OnConnect(ctx, p.k.NodeID)

pkg/rpc/stream_pool.go

@@ -43,6 +43,10 @@ type result[Resp any] struct {
 // that stream pools eventually shrink when the load decreases.
 const defaultPooledStreamIdleTimeout = 10 * time.Second
 
+// equalsFunc is a generic function type used to compare two RPC connections
+// for equivalence.
+type equalsFunc[Conn rpcConn] func(a, b Conn) bool
+
 // pooledStream is a wrapper around a grpc.ClientStream that is managed by a
 // streamPool. It is responsible for sending a single request and receiving a
 // single response on the stream at a time, mimicking the behavior of a gRPC
@@ -194,6 +198,7 @@ type streamPool[Req, Resp any, Conn rpcConn] struct {
 	stopper     *stop.Stopper
 	idleTimeout time.Duration
 	newStream   streamConstructor[Req, Resp, Conn]
+	connEquals  equalsFunc[Conn]
 
 	// cc and ccCtx are set on bind, when the gRPC connection is established.
 	cc Conn
@@ -207,12 +212,13 @@ type streamPool[Req, Resp any, Conn rpcConn] struct {
 }
 
 func makeStreamPool[Req, Resp any, Conn rpcConn](
-	stopper *stop.Stopper, newStream streamConstructor[Req, Resp, Conn],
+	stopper *stop.Stopper, newStream streamConstructor[Req, Resp, Conn], connEquals equalsFunc[Conn],
 ) streamPool[Req, Resp, Conn] {
 	return streamPool[Req, Resp, Conn]{
 		stopper:     stopper,
 		idleTimeout: defaultPooledStreamIdleTimeout,
 		newStream:   newStream,
+		connEquals:  connEquals,
 	}
 }
 
@@ -280,7 +286,7 @@ func (p *streamPool[Req, Resp, Conn]) remove(s *pooledStream[Req, Resp, Conn]) b
 
 func (p *streamPool[Req, Resp, Conn]) newPooledStream() (*pooledStream[Req, Resp, Conn], error) {
 	var zero Conn
-	if p.cc == zero {
+	if p.connEquals(p.cc, zero) {
 		return nil, errors.AssertionFailedf("streamPool not bound to a client conn")
 	}