Repository: zeromicro/go-zero
Pull request #1198 (2021-11-02 10:01:47)

feat: rpc dial can be set with NonBlock mode
1. use WithBlock() by default
2. NonBlock can be set in config or WithNonBlock() option


Patch:
diff --git a/zrpc/client.go b/zrpc/client.go
index 242d1924..77b444cf 100644
--- a/zrpc/client.go
+++ b/zrpc/client.go
@@ -14,6 +14,8 @@ import (
 var (
 	// WithDialOption is an alias of internal.WithDialOption.
 	WithDialOption = internal.WithDialOption
+	// WithNonBlock sets the dialing to be nonblock.
+	WithNonBlock = internal.WithNonBlock
 	// WithTimeout is an alias of internal.WithTimeout.
 	WithTimeout = internal.WithTimeout
 	// WithRetry is an alias of internal.WithRetry.
@@ -57,6 +59,9 @@ func NewClient(c RpcClientConf, options ...ClientOption) (Client, error) {
 			Token: c.Token,
 		})))
 	}
+	if c.NonBlock {
+		opts = append(opts, WithNonBlock())
+	}
 	if c.Timeout > 0 {
 		opts = append(opts, WithTimeout(time.Duration(c.Timeout)*time.Millisecond))
 	}
diff --git a/zrpc/config.go b/zrpc/config.go
index 07b72724..b3d33c89 100644
--- a/zrpc/config.go
+++ b/zrpc/config.go
@@ -28,6 +28,7 @@ type (
 		Target    string          `json:",optional"`
 		App       string          `json:",optional"`
 		Token     string          `json:",optional"`
+		NonBlock  bool            `json:",optional"`
 		Retry     bool            `json:",optional"` // grpc auto retry
 		Timeout   int64           `json:",default=2000"`
 	}
diff --git a/zrpc/internal/client.go b/zrpc/internal/client.go
index b23ced5e..2366fb73 100644
--- a/zrpc/internal/client.go
+++ b/zrpc/internal/client.go
@@ -35,6 +35,7 @@ type (
 
 	// A ClientOptions is a client options.
 	ClientOptions struct {
+		NonBlock    bool
 		Timeout     time.Duration
 		Secure      bool
 		Retry       bool
@@ -75,8 +76,11 @@ func (c *client) buildDialOptions(opts ...ClientOption) []grpc.DialOption {
 		options = append([]grpc.DialOption(nil), grpc.WithInsecure())
 	}
 
+	if !cliOpts.NonBlock {
+		options = append(options, grpc.WithBlock())
+	}
+
 	options = append(options,
-		grpc.WithBlock(),
 		WithUnaryClientInterceptors(
 			clientinterceptors.UnaryTracingInterceptor,
 			clientinterceptors.DurationInterceptor,
@@ -122,6 +126,13 @@ func WithDialOption(opt grpc.DialOption) ClientOption {
 	}
 }
 
+// WithNonBlock sets the dialing to be nonblock.
+func WithNonBlock() ClientOption {
+	return func(options *ClientOptions) {
+		options.NonBlock = true
+	}
+}
+
 // WithTimeout returns a func to customize a ClientOptions with given timeout.
 func WithTimeout(timeout time.Duration) ClientOption {
 	return func(options *ClientOptions) {


Test patch:
diff --git a/zrpc/internal/client_test.go b/zrpc/internal/client_test.go
index a548d97c..643f3076 100644
--- a/zrpc/internal/client_test.go
+++ b/zrpc/internal/client_test.go
@@ -31,6 +31,13 @@ func TestWithRetry(t *testing.T) {
 	assert.True(t, options.Retry)
 }
 
+func TestWithNonBlock(t *testing.T) {
+	var options ClientOptions
+	opt := WithNonBlock()
+	opt(&options)
+	assert.True(t, options.NonBlock)
+}
+
 func TestWithUnaryClientInterceptor(t *testing.T) {
 	var options ClientOptions
 	opt := WithUnaryClientInterceptor(func(ctx context.Context, method string, req, reply interface{},