diff --git a/zrpc/client.go b/zrpc/client.go index 700f5619..3c1ff363 100644 --- a/zrpc/client.go +++ b/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() } diff --git a/zrpc/internal/client.go b/zrpc/internal/client.go index 8528b9f8..08019566 100644 --- a/zrpc/internal/client.go +++ b/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)) + } +} diff --git a/zrpc/internal/client_test.go b/zrpc/internal/client_test.go index 6d22c348..af6eed78 100644 --- a/zrpc/internal/client_test.go +++ b/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")