feat: support WithStreamClientInterceptor for zrpc clients (#1907)

* feat: support WithStreamClientInterceptor for zrpc clients

* fix: data race

zrpc/client.go

@@ -15,6 +15,8 @@ var (
 	WithDialOption = internal.WithDialOption
 	// WithNonBlock sets the dialing to be nonblock.
 	WithNonBlock = internal.WithNonBlock
+	// WithStreamClientInterceptor is an alias of internal.WithStreamClientInterceptor.
+	WithStreamClientInterceptor = internal.WithStreamClientInterceptor
 	// WithTimeout is an alias of internal.WithTimeout.
 	WithTimeout = internal.WithTimeout
 	// WithTransportCredentials return a func to make the gRPC calls secured with given credentials.

zrpc/internal/client.go

@@ -131,6 +131,13 @@ func WithNonBlock() ClientOption {
 	}
 }
 
+// WithStreamClientInterceptor returns a func to customize a ClientOptions with given interceptor.
+func WithStreamClientInterceptor(interceptor grpc.StreamClientInterceptor) ClientOption {
+	return func(options *ClientOptions) {
+		options.DialOptions = append(options.DialOptions, WithStreamClientInterceptors(interceptor))
+	}
+}
+
 // WithTimeout returns a func to customize a ClientOptions with given timeout.
 func WithTimeout(timeout time.Duration) ClientOption {
 	return func(options *ClientOptions) {

zrpc/internal/client_test.go

@@ -31,6 +31,17 @@ func TestWithNonBlock(t *testing.T) {
 	assert.True(t, options.NonBlock)
 }
 
+func TestWithStreamClientInterceptor(t *testing.T) {
+	var options ClientOptions
+	opt := WithStreamClientInterceptor(func(ctx context.Context, desc *grpc.StreamDesc,
+		cc *grpc.ClientConn, method string, streamer grpc.Streamer,
+		opts ...grpc.CallOption) (grpc.ClientStream, error) {
+		return nil, nil
+	})
+	opt(&options)
+	assert.Equal(t, 1, len(options.DialOptions))
+}
+
 func TestWithTransportCredentials(t *testing.T) {
 	var options ClientOptions
 	opt := WithTransportCredentials(nil)

zrpc/internal/rpcserver.go

@@ -71,7 +71,7 @@ func (s *rpcServer) Start(register RegisterFn) error {
 		WithStreamServerInterceptors(streamInterceptors...))
 	server := grpc.NewServer(options...)
 	register(server)
-	// we need to make sure all others are wrapped up
+	// we need to make sure all others are wrapped up,
 	// so we do graceful stop at shutdown phase instead of wrap up phase
 	waitForCalled := proc.AddWrapUpListener(func() {
 		server.GracefulStop()