diff --git a/doc/goctl.md b/doc/goctl.md index 519fb606..edbc7149 100644 --- a/doc/goctl.md +++ b/doc/goctl.md @@ -239,10 +239,10 @@ src 示例代码如下 ``` 结构体中不需要提供Id,CreateTime,UpdateTime三个字段,会自动生成 -结构体中每个tag有两个可选标签 c 和 o -c 是改字段的注释 -o 是改字段需要生产的操作函数 可以取得get,find,set 分别表示生成返回单个对象的查询方法,返回多个对象的查询方法,设置该字段方法 -生成的目标文件会覆盖该简单go文件 +结构体中每个tag有两个可选标签 c 和 o +c 是该字段的注释 +o 是该字段需要生产的操作函数 可以取得get,find,set 分别表示生成返回单个对象的查询方法,返回多个对象的查询方法,设置该字段方法 +生成的目标文件会覆盖该简单go文件 ## goctl rpc生成(业务剥离中,暂未开放) diff --git a/doc/sharedcalls.md b/doc/sharedcalls.md index 6eed1529..d640a28b 100644 --- a/doc/sharedcalls.md +++ b/doc/sharedcalls.md @@ -106,7 +106,7 @@ func main() { // 没有拿到结果,则调用makeCall方法去获取资源,注意此处仍然是锁住的,可以保证只有一个goroutine可以调用makecall c := g.makeCall(key, fn) - // 返回调用结果 + // 返回调用结果 return c.val, c.err } ``` diff --git a/zrpc/client.go b/zrpc/client.go index eda789d6..5e47be60 100644 --- a/zrpc/client.go +++ b/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() } diff --git a/zrpc/internal/client.go b/zrpc/internal/client.go index 53addf18..9936d019 100644 --- a/zrpc/internal/client.go +++ b/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 + } } diff --git a/zrpc/internal/client_test.go b/zrpc/internal/client_test.go index 00b4875c..6d22c348 100644 --- a/zrpc/internal/client_test.go +++ b/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) }