diff --git a/internal/rpc/rpc_client.go b/internal/rpc/rpc_client.go index 375fb16..df7a838 100644 --- a/internal/rpc/rpc_client.go +++ b/internal/rpc/rpc_client.go @@ -17,8 +17,8 @@ import ( ) type RPCClient struct { - apiConfig *configs.APIConfig - nodeClients []pb.NodeServiceClient + apiConfig *configs.APIConfig + conns []*grpc.ClientConn } func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { @@ -26,8 +26,6 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { return nil, errors.New("api config should not be nil") } - nodeClients := []pb.NodeServiceClient{} - conns := []*grpc.ClientConn{} for _, endpoint := range apiConfig.RPC.Endpoints { conn, err := grpc.Dial(endpoint, grpc.WithInsecure()) @@ -40,22 +38,14 @@ func NewRPCClient(apiConfig *configs.APIConfig) (*RPCClient, error) { return nil, errors.New("[RPC]no available endpoints") } - // node clients - for _, conn := range conns { - nodeClients = append(nodeClients, pb.NewNodeServiceClient(conn)) - } - return &RPCClient{ - apiConfig: apiConfig, - nodeClients: nodeClients, + apiConfig: apiConfig, + conns: conns, }, nil } func (this *RPCClient) NodeRPC() pb.NodeServiceClient { - if len(this.nodeClients) > 0 { - return this.nodeClients[rands.Int(0, len(this.nodeClients)-1)] - } - return nil + return pb.NewNodeServiceClient(this.pickConn()) } func (this *RPCClient) Context() context.Context { @@ -79,3 +69,11 @@ func (this *RPCClient) Context() context.Context { ctx = metadata.AppendToOutgoingContext(ctx, "nodeId", this.apiConfig.NodeId, "token", token) return ctx } + +// 随机选择一个连接 +func (this *RPCClient) pickConn() *grpc.ClientConn { + if len(this.conns) == 0 { + return nil + } + return this.conns[rands.Int(0, len(this.conns)-1)] +}