فهرست منبع

fix zrpc client interceptor calling problem

kevin 4 سال پیش
والد
کامیت
cee170f3e9
3فایلهای تغییر یافته به همراه24 افزوده شده و 20 حذف شده
  1. 2 7
      zrpc/client.go
  2. 11 13
      zrpc/internal/client.go
  3. 11 0
      zrpc/internal/client_test.go

+ 2 - 7
zrpc/client.go

@@ -19,7 +19,6 @@ type (
 	ClientOption = internal.ClientOption
 
 	Client interface {
-		AddInterceptor(interceptor grpc.UnaryClientInterceptor)
 		Conn() *grpc.ClientConn
 	}
 
@@ -66,8 +65,8 @@ func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
 	}, nil
 }
 
-func NewClientNoAuth(c discov.EtcdConf) (Client, error) {
-	client, err := internal.NewClient(internal.BuildDiscovTarget(c.Hosts, c.Key))
+func NewClientNoAuth(c discov.EtcdConf, opts ...ClientOption) (Client, error) {
+	client, err := internal.NewClient(internal.BuildDiscovTarget(c.Hosts, c.Key), opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -81,10 +80,6 @@ func NewClientWithTarget(target string, opts ...ClientOption) (Client, error) {
 	return internal.NewClient(target, opts...)
 }
 
-func (rc *RpcClient) AddInterceptor(interceptor grpc.UnaryClientInterceptor) {
-	rc.client.AddInterceptor(interceptor)
-}
-
 func (rc *RpcClient) Conn() *grpc.ClientConn {
 	return rc.client.Conn()
 }

+ 11 - 13
zrpc/internal/client.go

@@ -31,8 +31,7 @@ type (
 	ClientOption func(options *ClientOptions)
 
 	client struct {
-		conn         *grpc.ClientConn
-		interceptors []grpc.UnaryClientInterceptor
+		conn *grpc.ClientConn
 	}
 )
 
@@ -46,18 +45,14 @@ func NewClient(target string, opts ...ClientOption) (*client, error) {
 	return &cli, nil
 }
 
-func (c *client) AddInterceptor(interceptor grpc.UnaryClientInterceptor) {
-	c.interceptors = append(c.interceptors, interceptor)
-}
-
 func (c *client) Conn() *grpc.ClientConn {
 	return c.conn
 }
 
 func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
-	var clientOptions ClientOptions
+	var cliOpts ClientOptions
 	for _, opt := range opts {
-		opt(&clientOptions)
+		opt(&cliOpts)
 	}
 
 	options := []grpc.DialOption{
@@ -68,14 +63,11 @@ func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 			clientinterceptors.DurationInterceptor,
 			clientinterceptors.BreakerInterceptor,
 			clientinterceptors.PrometheusInterceptor,
-			clientinterceptors.TimeoutInterceptor(clientOptions.Timeout),
+			clientinterceptors.TimeoutInterceptor(cliOpts.Timeout),
 		),
 	}
-	for _, interceptor := range c.interceptors {
-		options = append(options, WithUnaryClientInterceptors(interceptor))
-	}
 
-	return append(options, clientOptions.DialOptions...)
+	return append(options, cliOpts.DialOptions...)
 }
 
 func (c *client) dial(server string, opts ...ClientOption) error {
@@ -111,3 +103,9 @@ func WithTimeout(timeout time.Duration) ClientOption {
 		options.Timeout = timeout
 	}
 }
+
+func WithUnaryClientInterceptor(interceptor grpc.UnaryClientInterceptor) ClientOption {
+	return func(options *ClientOptions) {
+		options.DialOptions = append(options.DialOptions, WithUnaryClientInterceptors(interceptor))
+	}
+}

+ 11 - 0
zrpc/internal/client_test.go

@@ -1,6 +1,7 @@
 package internal
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -23,6 +24,16 @@ func TestWithTimeout(t *testing.T) {
 	assert.Equal(t, time.Second, options.Timeout)
 }
 
+func TestWithUnaryClientInterceptor(t *testing.T) {
+	var options ClientOptions
+	opt := WithUnaryClientInterceptor(func(ctx context.Context, method string, req, reply interface{},
+		cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
+		return nil
+	})
+	opt(&options)
+	assert.Equal(t, 1, len(options.DialOptions))
+}
+
 func TestBuildDialOptions(t *testing.T) {
 	var c client
 	agent := grpc.WithUserAgent("chrome")