add zrpc client interceptor

zrpc/client.go

@@ -17,6 +17,7 @@ var (
 
 type (
 	Client interface {
+		AddInterceptor(interceptor grpc.UnaryClientInterceptor)
 		Conn() *grpc.ClientConn
 	}
 
@@ -78,6 +79,10 @@ func NewClientWithTarget(target string, opts ...internal.ClientOption) (Client,
 	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()
 }

zrpc/internal/client.go

@@ -31,37 +31,30 @@ type (
 	ClientOption func(options *ClientOptions)
 
 	client struct {
-		conn *grpc.ClientConn
+		conn         *grpc.ClientConn
+		interceptors []grpc.UnaryClientInterceptor
 	}
 )
 
 func NewClient(target string, opts ...ClientOption) (*client, error) {
+	var cli client
 	opts = append(opts, WithDialOption(grpc.WithBalancerName(p2c.Name)))
-	conn, err := dial(target, opts...)
-	if err != nil {
+	if err := cli.dial(target, opts...); err != nil {
 		return nil, err
 	}
 
-	return &client{conn: conn}, nil
+	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 WithDialOption(opt grpc.DialOption) ClientOption {
-	return func(options *ClientOptions) {
-		options.DialOptions = append(options.DialOptions, opt)
-	}
-}
-
-func WithTimeout(timeout time.Duration) ClientOption {
-	return func(options *ClientOptions) {
-		options.Timeout = timeout
-	}
-}
-
-func buildDialOptions(opts ...ClientOption) []grpc.DialOption {
+func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 	var clientOptions ClientOptions
 	for _, opt := range opts {
 		opt(&clientOptions)
@@ -78,12 +71,15 @@ func buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 			clientinterceptors.TimeoutInterceptor(clientOptions.Timeout),
 		),
 	}
+	for _, interceptor := range c.interceptors {
+		options = append(options, WithUnaryClientInterceptors(interceptor))
+	}
 
 	return append(options, clientOptions.DialOptions...)
 }
 
-func dial(server string, opts ...ClientOption) (*grpc.ClientConn, error) {
-	options := buildDialOptions(opts...)
+func (c *client) dial(server string, opts ...ClientOption) error {
+	options := c.buildDialOptions(opts...)
 	timeCtx, cancel := context.WithTimeout(context.Background(), dialTimeout)
 	defer cancel()
 	conn, err := grpc.DialContext(timeCtx, server, options...)
@@ -96,9 +92,22 @@ func dial(server string, opts ...ClientOption) (*grpc.ClientConn, error) {
 				service = server[pos+1:]
 			}
 		}
-		return nil, fmt.Errorf("rpc dial: %s, error: %s, make sure rpc service %q is alread started",
+		return fmt.Errorf("rpc dial: %s, error: %s, make sure rpc service %q is alread started",
 			server, err.Error(), service)
 	}
 
-	return conn, nil
+	c.conn = conn
+	return nil
+}
+
+func WithDialOption(opt grpc.DialOption) ClientOption {
+	return func(options *ClientOptions) {
+		options.DialOptions = append(options.DialOptions, opt)
+	}
+}
+
+func WithTimeout(timeout time.Duration) ClientOption {
+	return func(options *ClientOptions) {
+		options.Timeout = timeout
+	}
 }

zrpc/internal/client_test.go

@@ -24,7 +24,8 @@ func TestWithTimeout(t *testing.T) {
 }
 
 func TestBuildDialOptions(t *testing.T) {
+	var c client
 	agent := grpc.WithUserAgent("chrome")
-	opts := buildDialOptions(WithDialOption(agent))
+	opts := c.buildDialOptions(WithDialOption(agent))
 	assert.Contains(t, opts, agent)
 }