@@ -9,65 +9,33 @@ import (
99 context "context"
1010
1111 "github.com/cockroachdb/cockroach/pkg/roachpb"
12+ "github.com/cockroachdb/cockroach/pkg/rpc/nodedialer"
1213 "github.com/cockroachdb/cockroach/pkg/rpc/rpcbase"
1314)
1415
1516// DialMultiRaftClient establishes a DRPC connection if enabled; otherwise,
1617// it falls back to gRPC. The established connection is used to create a
1718// RPCMultiRaftClient.
1819func DialMultiRaftClient (
19- nd rpcbase. NodeDialer , ctx context.Context , nodeID roachpb.NodeID , class rpcbase.ConnectionClass ,
20+ nd * nodedialer. Dialer , ctx context.Context , nodeID roachpb.NodeID , class rpcbase.ConnectionClass ,
2021) (RPCMultiRaftClient , error ) {
21- if ! rpcbase .TODODRPC {
22- conn , err := nd .Dial (ctx , nodeID , class )
23- if err != nil {
24- return nil , err
25- }
26- return NewGRPCMultiRaftClientAdapter (conn ), nil
27- }
28- conn , err := nd .DRPCDial (ctx , nodeID , class )
29- if err != nil {
30- return nil , err
31- }
32- return NewDRPCMultiRaftClientAdapter (conn ), nil
22+ return nodedialer .DialRPCClient (nd , ctx , nodeID , class , NewGRPCMultiRaftClientAdapter , NewDRPCMultiRaftClientAdapter )
3323}
3424
3525// DialPerReplicaClient establishes a DRPC connection if enabled; otherwise,
3626// it falls back to gRPC. The established connection is used to create a
3727// RPCPerReplicaClient.
3828func DialPerReplicaClient (
39- nd rpcbase. NodeDialer , ctx context.Context , nodeID roachpb.NodeID , class rpcbase.ConnectionClass ,
29+ nd * nodedialer. Dialer , ctx context.Context , nodeID roachpb.NodeID , class rpcbase.ConnectionClass ,
4030) (RPCPerReplicaClient , error ) {
41- if ! rpcbase .TODODRPC {
42- conn , err := nd .Dial (ctx , nodeID , class )
43- if err != nil {
44- return nil , err
45- }
46- return NewGRPCPerReplicaClientAdapter (conn ), nil
47- }
48- conn , err := nd .DRPCDial (ctx , nodeID , class )
49- if err != nil {
50- return nil , err
51- }
52- return NewDRPCPerReplicaClientAdapter (conn ), nil
31+ return nodedialer .DialRPCClient (nd , ctx , nodeID , class , NewGRPCPerReplicaClientAdapter , NewDRPCPerReplicaClientAdapter )
5332}
5433
5534// DialPerStoreClient establishes a DRPC connection if enabled; otherwise,
5635// it falls back to gRPC. The established connection is used to create a
5736// RPCPerStoreClient.
5837func DialPerStoreClient (
59- nd rpcbase. NodeDialer , ctx context.Context , nodeID roachpb.NodeID , class rpcbase.ConnectionClass ,
38+ nd * nodedialer. Dialer , ctx context.Context , nodeID roachpb.NodeID , class rpcbase.ConnectionClass ,
6039) (RPCPerStoreClient , error ) {
61- if ! rpcbase .TODODRPC {
62- conn , err := nd .Dial (ctx , nodeID , class )
63- if err != nil {
64- return nil , err
65- }
66- return NewGRPCPerStoreClientAdapter (conn ), nil
67- }
68- conn , err := nd .DRPCDial (ctx , nodeID , class )
69- if err != nil {
70- return nil , err
71- }
72- return NewDRPCPerStoreClientAdapter (conn ), nil
40+ return nodedialer .DialRPCClient (nd , ctx , nodeID , class , NewGRPCPerStoreClientAdapter , NewDRPCPerStoreClientAdapter )
7341}
0 commit comments