From 08fa64d5539ebdae3ad9dcb3ca8a98bd51d76ec8 Mon Sep 17 00:00:00 2001
From: kevin <wanjunfeng@gmail.com>
Date: Tue, 11 Aug 2020 14:33:10 +0800
Subject: [PATCH] refactor rpcx, export WithDialOption and WithTimeout

---
 .../graceful/dns/api/svc/servicecontext.go    |  2 +-
 .../graceful/etcd/api/svc/servicecontext.go   |  2 +-
 example/tracing/edge/main.go                  |  2 +-
 example/tracing/portal/server.go              |  4 +--
 rpcx/client.go                                | 29 +++++++++++++------
 rpcx/internal/client.go                       |  4 ---
 rpcx/proxy.go                                 |  4 +--
 7 files changed, 27 insertions(+), 20 deletions(-)

diff --git a/example/graceful/dns/api/svc/servicecontext.go b/example/graceful/dns/api/svc/servicecontext.go
index 7bd478e0..42d8d4b7 100644
--- a/example/graceful/dns/api/svc/servicecontext.go
+++ b/example/graceful/dns/api/svc/servicecontext.go
@@ -3,5 +3,5 @@ package svc
 import "github.com/tal-tech/go-zero/rpcx"
 
 type ServiceContext struct {
-	Client *rpcx.RpcClient
+	Client rpcx.Client
 }
diff --git a/example/graceful/etcd/api/svc/servicecontext.go b/example/graceful/etcd/api/svc/servicecontext.go
index 7bd478e0..42d8d4b7 100644
--- a/example/graceful/etcd/api/svc/servicecontext.go
+++ b/example/graceful/etcd/api/svc/servicecontext.go
@@ -3,5 +3,5 @@ package svc
 import "github.com/tal-tech/go-zero/rpcx"
 
 type ServiceContext struct {
-	Client *rpcx.RpcClient
+	Client rpcx.Client
 }
diff --git a/example/tracing/edge/main.go b/example/tracing/edge/main.go
index ce43a02a..558ba94b 100644
--- a/example/tracing/edge/main.go
+++ b/example/tracing/edge/main.go
@@ -15,7 +15,7 @@ import (
 
 var (
 	configFile = flag.String("f", "config.json", "the config file")
-	client     *rpcx.RpcClient
+	client     rpcx.Client
 )
 
 func handle(w http.ResponseWriter, r *http.Request) {
diff --git a/example/tracing/portal/server.go b/example/tracing/portal/server.go
index 47198344..6c5e0c4d 100644
--- a/example/tracing/portal/server.go
+++ b/example/tracing/portal/server.go
@@ -20,11 +20,11 @@ type (
 	}
 
 	PortalServer struct {
-		userRpc *rpcx.RpcClient
+		userRpc rpcx.Client
 	}
 )
 
-func NewPortalServer(client *rpcx.RpcClient) *PortalServer {
+func NewPortalServer(client rpcx.Client) *PortalServer {
 	return &PortalServer{
 		userRpc: client,
 	}
diff --git a/rpcx/client.go b/rpcx/client.go
index cb523951..7f589d56 100644
--- a/rpcx/client.go
+++ b/rpcx/client.go
@@ -10,11 +10,22 @@ import (
 	"google.golang.org/grpc"
 )
 
-type RpcClient struct {
-	client internal.Client
-}
+var (
+	WithDialOption = internal.WithDialOption
+	WithTimeout    = internal.WithTimeout
+)
+
+type (
+	Client interface {
+		Conn() *grpc.ClientConn
+	}
+
+	RpcClient struct {
+		client Client
+	}
+)
 
-func MustNewClient(c RpcClientConf, options ...internal.ClientOption) *RpcClient {
+func MustNewClient(c RpcClientConf, options ...internal.ClientOption) Client {
 	cli, err := NewClient(c, options...)
 	if err != nil {
 		log.Fatal(err)
@@ -23,20 +34,20 @@ func MustNewClient(c RpcClientConf, options ...internal.ClientOption) *RpcClient
 	return cli
 }
 
-func NewClient(c RpcClientConf, options ...internal.ClientOption) (*RpcClient, error) {
+func NewClient(c RpcClientConf, options ...internal.ClientOption) (Client, error) {
 	var opts []internal.ClientOption
 	if c.HasCredential() {
-		opts = append(opts, internal.WithDialOption(grpc.WithPerRPCCredentials(&auth.Credential{
+		opts = append(opts, WithDialOption(grpc.WithPerRPCCredentials(&auth.Credential{
 			App:   c.App,
 			Token: c.Token,
 		})))
 	}
 	if c.Timeout > 0 {
-		opts = append(opts, internal.WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
+		opts = append(opts, WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
 	}
 	opts = append(opts, options...)
 
-	var client internal.Client
+	var client Client
 	var err error
 	if len(c.Server) > 0 {
 		client, err = internal.NewDirectClient(c.Server, opts...)
@@ -52,7 +63,7 @@ func NewClient(c RpcClientConf, options ...internal.ClientOption) (*RpcClient, e
 	}, nil
 }
 
-func NewClientNoAuth(c discov.EtcdConf) (*RpcClient, error) {
+func NewClientNoAuth(c discov.EtcdConf) (Client, error) {
 	client, err := internal.NewDiscovClient(c.Hosts, c.Key)
 	if err != nil {
 		return nil, err
diff --git a/rpcx/internal/client.go b/rpcx/internal/client.go
index 9576f261..1b799787 100644
--- a/rpcx/internal/client.go
+++ b/rpcx/internal/client.go
@@ -18,10 +18,6 @@ type (
 	}
 
 	ClientOption func(options *ClientOptions)
-
-	Client interface {
-		Conn() *grpc.ClientConn
-	}
 )
 
 func WithDialOption(opt grpc.DialOption) ClientOption {
diff --git a/rpcx/proxy.go b/rpcx/proxy.go
index e4c90d0a..dfdff273 100644
--- a/rpcx/proxy.go
+++ b/rpcx/proxy.go
@@ -12,7 +12,7 @@ import (
 
 type RpcProxy struct {
 	backend     string
-	clients     map[string]*RpcClient
+	clients     map[string]Client
 	options     []internal.ClientOption
 	sharedCalls syncx.SharedCalls
 	lock        sync.Mutex
@@ -21,7 +21,7 @@ type RpcProxy struct {
 func NewRpcProxy(backend string, opts ...internal.ClientOption) *RpcProxy {
 	return &RpcProxy{
 		backend:     backend,
-		clients:     make(map[string]*RpcClient),
+		clients:     make(map[string]Client),
 		options:     opts,
 		sharedCalls: syncx.NewSharedCalls(),
 	}