Merge #148011

148011: cli,kv,server: consolidate `StatusClient` RPC client creation r=cthumuluru-crdb a=cthumuluru-crdb

This commit consolidates `StatusClient` RPC client creation logic in cli, kv, and server packages. It is a continuation of the work done in #147606.

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

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

Author: craig[bot]

pkg/acceptance/localcluster/BUILD.bazel

@@ -14,6 +14,7 @@ go_library(
         "//pkg/config/zonepb",
         "//pkg/roachpb",
         "//pkg/rpc",
+        "//pkg/rpc/rpcbase",
         "//pkg/server/serverpb",
         "//pkg/settings/cluster",
         "//pkg/testutils",

pkg/acceptance/localcluster/cluster.go

@@ -28,6 +28,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
@@ -477,11 +478,14 @@ func (n *Node) StatusClient(ctx context.Context) serverpb.StatusClient {
 		return existingClient
 	}
 
-	conn, err := n.rpcCtx.GRPCUnvalidatedDial(n.RPCAddr(), roachpb.Locality{}).Connect(ctx)
-	if err != nil {
-		log.Fatalf(context.Background(), "failed to initialize status client: %s", err)
+	if !rpcbase.TODODRPC {
+		conn, err := n.rpcCtx.GRPCUnvalidatedDial(n.RPCAddr(), roachpb.Locality{}).Connect(ctx)
+		if err != nil {
+			log.Fatalf(context.Background(), "failed to initialize status client: %s", err)
+		}
+		return serverpb.NewStatusClient(conn)
 	}
-	return serverpb.NewStatusClient(conn)
+	return nil // This should never happen
 }
 
 func (n *Node) logDir() string {

pkg/cli/BUILD.bazel

@@ -46,7 +46,7 @@ go_library(
         "node.go",
         "nodelocal.go",
         "prefixer.go",
-        "rpc_client.go",
+        "rpc_clients.go",
         "rpc_node_shutdown.go",
         "sql_client.go",
         "sql_shell_cmd.go",

pkg/cli/debug.go

@@ -1037,12 +1037,13 @@ func runDebugGossipValues(cmd *cobra.Command, args []string) error {
 			return errors.Wrap(err, "failed to parse provided file as gossip.InfoStatus")
 		}
 	} else {
-		status, finish, err := getStatusClient(ctx, serverCfg)
+		conn, finish, err := newClientConn(ctx, serverCfg)
 		if err != nil {
 			return err
 		}
 		defer finish()
 
+		status := conn.NewStatusClient()
 		gossipInfo, err = status.Gossip(ctx, &serverpb.GossipRequest{})
 		if err != nil {
 			return errors.Wrap(err, "failed to retrieve gossip from server")

pkg/cli/debug_list_files.go

@@ -41,13 +41,14 @@ func runDebugListFiles(cmd *cobra.Command, _ []string) error {
 	defer cancel()
 
 	// Connect to the node pointed to in the command line.
-	status, finish, err := getStatusClient(ctx, serverCfg)
+	conn, finish, err := newClientConn(ctx, serverCfg)
 	if err != nil {
 		return err
 	}
 	defer finish()
 
 	// Retrieve the details for the head node.
+	status := conn.NewStatusClient()
 	firstNodeDetails, err := status.Details(ctx, &serverpb.DetailsRequest{NodeId: "local"})
 	if err != nil {
 		return err

pkg/cli/haproxy.go

@@ -216,12 +216,13 @@ func runGenHAProxyCmd(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	c, finish, err := getStatusClient(ctx, serverCfg)
+	conn, finish, err := newClientConn(ctx, serverCfg)
 	if err != nil {
 		return err
 	}
 	defer finish()
 
+	c := conn.NewStatusClient()
 	nodeStatuses, err := c.Nodes(ctx, &serverpb.NodesRequest{})
 	if err != nil {
 		return err

pkg/cli/node.go

@@ -348,15 +348,14 @@ func runDecommissionNode(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	conn, finish, err := getClientGRPCConn(ctx, serverCfg)
+	conn, finish, err := newClientConn(ctx, serverCfg)
 	if err != nil {
-		return errors.Wrap(err, "failed to connect to the node")
+		return err
 	}
 	defer finish()
 
-	s := serverpb.NewStatusClient(conn)
-
-	localNodeID, err := getLocalNodeID(ctx, s)
+	statusClient := conn.NewStatusClient()
+	localNodeID, err := getLocalNodeID(ctx, statusClient)
 	if err != nil {
 		return err
 	}
@@ -366,12 +365,12 @@ func runDecommissionNode(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if err := expectNodesDecommissioned(ctx, s, nodeIDs, false /* expDecommissioned */); err != nil {
+	if err := expectNodesDecommissioned(ctx, statusClient, nodeIDs, false /* expDecommissioned */); err != nil {
 		return err
 	}
 
-	c := serverpb.NewAdminClient(conn)
-	if err := runDecommissionNodeImpl(ctx, c, nodeCtx.nodeDecommissionWait,
+	adminClient := conn.NewAdminClient()
+	if err := runDecommissionNodeImpl(ctx, adminClient, nodeCtx.nodeDecommissionWait,
 		nodeCtx.nodeDecommissionChecks, nodeCtx.nodeDecommissionDryRun,
 		nodeIDs, localNodeID,
 	); err != nil {
@@ -831,15 +830,14 @@ func runRecommissionNode(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	conn, finish, err := getClientGRPCConn(ctx, serverCfg)
+	conn, finish, err := newClientConn(ctx, serverCfg)
 	if err != nil {
-		return errors.Wrap(err, "failed to connect to the node")
+		return err
 	}
 	defer finish()
 
-	s := serverpb.NewStatusClient(conn)
-
-	localNodeID, err := getLocalNodeID(ctx, s)
+	statusClient := conn.NewStatusClient()
+	localNodeID, err := getLocalNodeID(ctx, statusClient)
 	if err != nil {
 		return err
 	}
@@ -849,16 +847,16 @@ func runRecommissionNode(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if err := expectNodesDecommissioned(ctx, s, nodeIDs, true /* expDecommissioned */); err != nil {
+	if err := expectNodesDecommissioned(ctx, statusClient, nodeIDs, true /* expDecommissioned */); err != nil {
 		return err
 	}
 
-	c := serverpb.NewAdminClient(conn)
+	adminClient := conn.NewAdminClient()
 	req := &serverpb.DecommissionRequest{
 		NodeIDs:          nodeIDs,
 		TargetMembership: livenesspb.MembershipStatus_ACTIVE,
 	}
-	resp, err := c.Decommission(ctx, req)
+	resp, err := adminClient.Decommission(ctx, req)
 	if err != nil {
 		cause := errors.UnwrapAll(err)
 		// If it's a specific illegal membership transition error, we try to
@@ -922,21 +920,22 @@ func runDrain(cmd *cobra.Command, args []string) (err error) {
 	}
 
 	// Establish a RPC connection.
-	c, finish, err := getAdminClient(ctx, serverCfg)
+	conn, finish, err := newClientConn(ctx, serverCfg)
 	if err != nil {
 		return err
 	}
 	defer finish()
 
-	if _, _, err := doDrain(ctx, c, targetNode); err != nil {
+	adminClient := conn.NewAdminClient()
+	if _, _, err := doDrain(ctx, adminClient, targetNode); err != nil {
 		return err
 	}
 
 	// Report "ok" if there was no error.
 	fmt.Println("drain ok")
 
 	if drainCtx.shutdown {
-		if _, err := doShutdown(ctx, c, targetNode); err != nil {
+		if _, err := doShutdown(ctx, adminClient, targetNode); err != nil {
 			return err
 		}
 		// Report "ok" if there was no error.

pkg/cli/rpc_client.go

@@ -0,0 +1,92 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package cli
+
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/rpc"
+	"github.com/cockroachdb/cockroach/pkg/server"
+	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/ts/tspb"
+	"github.com/cockroachdb/errors"
+	"google.golang.org/grpc"
+)
+
+// rpcConn defines a common interface for creating RPC clients. It hides the
+// underlying RPC connection (gRPC or DRPC), making it easy to swap
+// them without changing the caller code.
+type rpcConn interface {
+	NewStatusClient() serverpb.StatusClient
+	NewAdminClient() serverpb.AdminClient
+	NewInitClient() serverpb.InitClient
+	NewTimeSeriesClient() tspb.TimeSeriesClient
+	NewInternalClient() kvpb.InternalClient
+}
+
+// grpcConn is an implementation of rpcConn that provides methods to create
+// various RPC clients. This allows the CLI to interact with the server using
+// gRPC without exposing the underlying connection details.
+type grpcConn struct {
+	conn *grpc.ClientConn
+}
+
+func (c *grpcConn) NewStatusClient() serverpb.StatusClient {
+	return serverpb.NewStatusClient(c.conn)
+}
+
+func (c *grpcConn) NewAdminClient() serverpb.AdminClient {
+	return serverpb.NewAdminClient(c.conn)
+}
+
+func (c *grpcConn) NewInitClient() serverpb.InitClient {
+	return serverpb.NewInitClient(c.conn)
+}
+
+func (c *grpcConn) NewTimeSeriesClient() tspb.TimeSeriesClient {
+	return tspb.NewTimeSeriesClient(c.conn)
+}
+
+func (c *grpcConn) NewInternalClient() kvpb.InternalClient {
+	return kvpb.NewInternalClient(c.conn)
+}
+
+func makeRPCClientConfig(cfg server.Config) rpc.ClientConnConfig {
+	var knobs rpc.ContextTestingKnobs
+	if sknobs := cfg.TestingKnobs.Server; sknobs != nil {
+		knobs = sknobs.(*server.TestingKnobs).ContextTestingKnobs
+	}
+	return rpc.MakeClientConnConfigFromBaseConfig(
+		*cfg.Config,
+		cfg.Config.User,
+		cfg.BaseConfig.Tracer,
+		cfg.BaseConfig.Settings,
+		nil, /* clock */
+		knobs,
+	)
+}
+
+func newClientConn(ctx context.Context, cfg server.Config) (rpcConn, func(), error) {
+	ccfg := makeRPCClientConfig(cfg)
+	cc, finish, err := rpc.NewClientConn(ctx, ccfg)
+	return &grpcConn{conn: cc}, finish, errors.Wrap(err, "failed to connect to the node")
+}
+
+func getClientGRPCConn(ctx context.Context, cfg server.Config) (*grpc.ClientConn, func(), error) {
+	ccfg := makeRPCClientConfig(cfg)
+	return rpc.NewClientConn(ctx, ccfg)
+}
+
+// getAdminClient returns an AdminClient and a closure that must be invoked
+// to free associated resources.
+func getAdminClient(ctx context.Context, cfg server.Config) (serverpb.AdminClient, func(), error) {
+	conn, finish, err := getClientGRPCConn(ctx, cfg)
+	if err != nil {
+		return nil, nil, errors.Wrap(err, "failed to connect to the node")
+	}
+	return serverpb.NewAdminClient(conn), finish, nil
+}