fix zrpc client interceptor calling problem

master
kevin 4 years ago
parent 907efd92c9
commit cee170f3e9

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

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

@ -1,6 +1,7 @@
package internal package internal
import ( import (
"context"
"testing" "testing"
"time" "time"
@ -23,6 +24,16 @@ func TestWithTimeout(t *testing.T) {
assert.Equal(t, time.Second, options.Timeout) 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) { func TestBuildDialOptions(t *testing.T) {
var c client var c client
agent := grpc.WithUserAgent("chrome") agent := grpc.WithUserAgent("chrome")

Loading…
Cancel
Save