feat: support the specified timeout of rpc methods (#2742)

Co-authored-by: hanzijian <hanzijian@52tt.com>
Co-authored-by: Kevin Wan <wanjunfeng@gmail.com>

zrpc/internal/clientinterceptors/timeoutinterceptor.go

@@ -11,13 +11,36 @@ import (
 func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
 	return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
 		invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
-		if timeout <= 0 {
+		t := getTimeoutByCallOptions(opts, timeout)
+		if t <= 0 {
 			return invoker(ctx, method, req, reply, cc, opts...)
 		}
 
-		ctx, cancel := context.WithTimeout(ctx, timeout)
+		ctx, cancel := context.WithTimeout(ctx, t)
 		defer cancel()
 
 		return invoker(ctx, method, req, reply, cc, opts...)
 	}
 }
+
+func getTimeoutByCallOptions(callOptions []grpc.CallOption, defaultTimeout time.Duration) time.Duration {
+	for _, callOption := range callOptions {
+		if o, ok := callOption.(TimeoutCallOption); ok {
+			return o.timeout
+		}
+	}
+
+	return defaultTimeout
+}
+
+type TimeoutCallOption struct {
+	grpc.EmptyCallOption
+
+	timeout time.Duration
+}
+
+func WithTimeoutCallOption(timeout time.Duration) grpc.CallOption {
+	return TimeoutCallOption{
+		timeout: timeout,
+	}
+}

zrpc/internal/clientinterceptors/timeoutinterceptor_test.go

@@ -66,3 +66,74 @@ func TestTimeoutInterceptor_panic(t *testing.T) {
 		})
 	}
 }
+
+func TestTimeoutInterceptor_TimeoutCallOption(t *testing.T) {
+	type args struct {
+		interceptorTimeout time.Duration
+		callOptionTimeout  time.Duration
+		runTime            time.Duration
+	}
+	var tests = []struct {
+		name    string
+		args    args
+		wantErr error
+	}{
+		{
+			name: "do not timeout without call option timeout",
+			args: args{
+				interceptorTimeout: time.Second,
+				runTime:            time.Millisecond * 50,
+			},
+			wantErr: nil,
+		},
+		{
+			name: "timeout without call option timeout",
+			args: args{
+				interceptorTimeout: time.Second,
+				runTime:            time.Second * 2,
+			},
+			wantErr: context.DeadlineExceeded,
+		},
+		{
+			name: "do not timeout with call option timeout",
+			args: args{
+				interceptorTimeout: time.Second,
+				callOptionTimeout:  time.Second * 3,
+				runTime:            time.Second * 2,
+			},
+			wantErr: nil,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			interceptor := TimeoutInterceptor(tt.args.interceptorTimeout)
+
+			cc := new(grpc.ClientConn)
+			var co []grpc.CallOption
+			if tt.args.callOptionTimeout > 0 {
+				co = append(co, WithTimeoutCallOption(tt.args.callOptionTimeout))
+			}
+
+			err := interceptor(context.Background(), "/foo", nil, nil, cc,
+				func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
+					opts ...grpc.CallOption) error {
+					timer := time.NewTimer(tt.args.runTime)
+					defer timer.Stop()
+
+					select {
+					case <-timer.C:
+						return nil
+					case <-ctx.Done():
+						return ctx.Err()
+					}
+				}, co...,
+			)
+			t.Logf("error: %+v", err)
+
+			assert.EqualValues(t, tt.wantErr, err)
+		})
+	}
+}