Merge #148611

148611: *: create DRPC RPC clients conditionally r=cthumuluru-crdb a=cthumuluru-crdb

Create DRPC RPC clients when DRPC is enabled. Although DRPC is currently disabled by default, it will be eventually controlled via `rpc.experimental_drpc.enabled` cluster setting.

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

Co-authored-by: Chandra Thumuluru <[email protected]>

Author: craig[bot]

pkg/acceptance/localcluster/cluster.go

@@ -485,7 +485,11 @@ func (n *Node) StatusClient(ctx context.Context) serverpb.RPCStatusClient {
 		}
 		return serverpb.NewGRPCStatusClientAdapter(conn)
 	}
-	return nil // This should never happen
+	conn, err := n.rpcCtx.DRPCUnvalidatedDial(n.RPCAddr(), roachpb.Locality{}).Connect(ctx)
+	if err != nil {
+		log.Fatalf(context.Background(), "failed to initialize status client: %s", err)
+	}
+	return serverpb.NewDRPCStatusClientAdapter(conn)
 }
 
 func (n *Node) logDir() string {

pkg/blobs/blobspb/rpc_clients.go

@@ -25,5 +25,9 @@ func DialBlobClient(
 		}
 		return NewGRPCBlobClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := nd.DRPCDial(ctx, nodeID, class)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCBlobClientAdapter(conn), nil
 }

pkg/cli/rpc_clients.go

@@ -107,13 +107,12 @@ func newClientConn(ctx context.Context, cfg server.Config) (rpcConn, func(), err
 			return nil, nil, errors.Wrap(err, "failed to connect to the node")
 		}
 		return &grpcConn{conn: cc}, finish, nil
-	} else {
-		dc, finish, err := rpc.NewDRPCClientConn(ctx, ccfg)
-		if err != nil {
-			return nil, nil, errors.Wrap(err, "failed to connect to the node")
-		}
-		return &drpcConn{conn: dc}, finish, nil
 	}
+	dc, finish, err := rpc.NewDRPCClientConn(ctx, ccfg)
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "failed to connect to the node")
+	}
+	return &drpcConn{conn: dc}, finish, nil
 }
 
 // dialAdminClient dials a client connection and returns an AdminClient and a

pkg/gossip/BUILD.bazel

@@ -41,6 +41,7 @@ go_library(
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_logtags//:logtags",
         "@com_github_cockroachdb_redact//:redact",
+        "@io_storj_drpc//:drpc",
         "@org_golang_google_grpc//:grpc",
     ],
 )

pkg/gossip/client.go

@@ -22,6 +22,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"google.golang.org/grpc"
+	drpc "storj.io/drpc"
 )
 
 // client is a client-side RPC connection to a gossip peer node.
@@ -415,6 +416,21 @@ func (c *client) dial(ctx context.Context, rpcCtx *rpc.Context) (*grpc.ClientCon
 	return conn.Connect(ctx)
 }
 
+// dials the peer node and returns a DRPC connection to the peer node.
+func (c *client) drpcDial(ctx context.Context, rpcCtx *rpc.Context) (drpc.Conn, error) {
+	var conn *rpc.DRPCConnection
+	if c.peerID != 0 {
+		conn = rpcCtx.DRPCDialNode(c.addr.String(), c.peerID, c.locality, rpcbase.SystemClass)
+	} else {
+		// TODO(baptist): Use this as a temporary connection for getting
+		// onto gossip and then replace with a validated connection.
+		log.Infof(ctx, "unvalidated bootstrap gossip dial to %s", c.addr)
+		conn = rpcCtx.DRPCUnvalidatedDial(c.addr.String(), c.locality)
+	}
+
+	return conn.Connect(ctx)
+}
+
 // dialGossipClient establishes a DRPC connection if enabled; otherwise,
 // it falls back to gRPC. The established connection is used to create a
 // RPCGossipClient.
@@ -428,5 +444,9 @@ func (c *client) dialGossipClient(
 		}
 		return NewGRPCGossipClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := c.drpcDial(ctx, rpcCtx)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCGossipClientAdapter(conn), nil
 }

pkg/keyvisualizer/keyvispb/rpc_clients.go

@@ -25,5 +25,9 @@ func DialKeyVisualizerClient(
 		}
 		return NewGRPCKeyVisualizerClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := nd.DRPCDial(ctx, nodeID, class)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCKeyVisualizerClientAdapter(conn), nil
 }

pkg/kv/kvclient/kvtenant/BUILD.bazel

@@ -46,6 +46,7 @@ go_library(
         "//pkg/util/uuid",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_errors//errorspb",
+        "@io_storj_drpc//:drpc",
         "@org_golang_google_grpc//:grpc",
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//status",

pkg/kv/kvclient/kvtenant/connector.go

@@ -51,6 +51,7 @@ import (
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
+	"storj.io/drpc"
 )
 
 func init() {
@@ -993,6 +994,19 @@ func (c *connector) dialAddrs(ctx context.Context) (*client, error) {
 					RPCTimeSeriesClient:       tspb.NewGRPCTimeSeriesClientAdapter(conn),
 				}, nil
 			}
+			conn, err := c.drpcDialAddr(ctx, addr)
+			if err != nil {
+				log.Warningf(ctx, "error dialing tenant KV address %s: %v", addr, err)
+				continue
+			}
+			return &client{
+				RPCTenantServiceClient:    kvpb.NewDRPCTenantServiceClientAdapter(conn),
+				RPCTenantSpanConfigClient: kvpb.NewDRPCTenantSpanConfigClientAdapter(conn),
+				RPCTenantUsageClient:      kvpb.NewDRPCTenantUsageClientAdapter(conn),
+				RPCStatusClient:           serverpb.NewDRPCStatusClientAdapter(conn),
+				RPCAdminClient:            serverpb.NewDRPCAdminClientAdapter(conn),
+				RPCTimeSeriesClient:       tspb.NewDRPCTimeSeriesClientAdapter(conn),
+			}, nil
 		}
 	}
 	return nil, errors.Wrap(ctx.Err(), "dial addrs")
@@ -1009,6 +1023,17 @@ func (c *connector) dialAddr(ctx context.Context, addr string) (conn *grpc.Clien
 	return conn, err
 }
 
+func (c *connector) drpcDialAddr(ctx context.Context, addr string) (conn drpc.Conn, err error) {
+	if c.rpcDialTimeout == 0 {
+		return c.rpcContext.DRPCUnvalidatedDial(addr, roachpb.Locality{}).Connect(ctx)
+	}
+	err = timeutil.RunWithTimeout(ctx, "dial addr", c.rpcDialTimeout, func(ctx context.Context) error {
+		conn, err = c.rpcContext.DRPCUnvalidatedDial(addr, roachpb.Locality{}).Connect(ctx)
+		return err
+	})
+	return conn, err
+}
+
 func (c *connector) tryForgetClient(ctx context.Context, client *client) {
 	if ctx.Err() != nil {
 		// Error (may be) due to context. Don't forget client.

pkg/kv/kvserver/closedts/ctpb/rpc_clients.go

@@ -25,5 +25,9 @@ func DialSideTransportClient(
 		}
 		return NewGRPCSideTransportClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := nd.DRPCDial(ctx, nodeID, class)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCSideTransportClientAdapter(conn), nil
 }

pkg/kv/kvserver/rpc_clients.go

@@ -25,7 +25,11 @@ func DialMultiRaftClient(
 		}
 		return NewGRPCMultiRaftClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := nd.DRPCDial(ctx, nodeID, class)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCMultiRaftClientAdapter(conn), nil
 }
 
 // DialPerReplicaClient establishes a DRPC connection if enabled; otherwise,
@@ -41,7 +45,11 @@ func DialPerReplicaClient(
 		}
 		return NewGRPCPerReplicaClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := nd.DRPCDial(ctx, nodeID, class)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCPerReplicaClientAdapter(conn), nil
 }
 
 // DialPerStoreClient establishes a DRPC connection if enabled; otherwise,
@@ -57,5 +65,9 @@ func DialPerStoreClient(
 		}
 		return NewGRPCPerStoreClientAdapter(conn), nil
 	}
-	return nil, nil
+	conn, err := nd.DRPCDial(ctx, nodeID, class)
+	if err != nil {
+		return nil, err
+	}
+	return NewDRPCPerStoreClientAdapter(conn), nil
 }