瀏覽代碼

add zrpc client interceptor

kevin 4 年之前
父節點
當前提交
dbca20e3df
共有 5 個文件被更改,包括 42 次插入27 次删除
  1. 4 4
      doc/goctl.md
  2. 1 1
      doc/sharedcalls.md
  3. 5 0
      zrpc/client.go
  4. 30 21
      zrpc/internal/client.go
  5. 2 1
      zrpc/internal/client_test.go

+ 4 - 4
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生成(业务剥离中,暂未开放)
 

+ 1 - 1
doc/sharedcalls.md

@@ -106,7 +106,7 @@ func main() {
 
     // 没有拿到结果,则调用makeCall方法去获取资源,注意此处仍然是锁住的,可以保证只有一个goroutine可以调用makecall
     c := g.makeCall(key, fn)
-  // 返回调用结果
+    // 返回调用结果
     return c.val, c.err
   }
   ```

+ 5 - 0
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()
 }

+ 30 - 21
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) Conn() *grpc.ClientConn {
-	return c.conn
-}
-
-func WithDialOption(opt grpc.DialOption) ClientOption {
-	return func(options *ClientOptions) {
-		options.DialOptions = append(options.DialOptions, opt)
-	}
+func (c *client) AddInterceptor(interceptor grpc.UnaryClientInterceptor) {
+	c.interceptors = append(c.interceptors, interceptor)
 }
 
-func WithTimeout(timeout time.Duration) ClientOption {
-	return func(options *ClientOptions) {
-		options.Timeout = timeout
-	}
+func (c *client) Conn() *grpc.ClientConn {
+	return c.conn
 }
 
-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
+	}
 }

+ 2 - 1
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)
 }