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

zrpc/internal/config.go

@@ -24,4 +24,6 @@ type (
 		Prometheus bool     `json:",default=true"`
 		Breaker    bool     `json:",default=true"`
 	}
+
+	ServerSpecifiedTimeoutConf = serverinterceptors.ServerSpecifiedTimeoutConf
 )

zrpc/internal/serverinterceptors/timeoutinterceptor.go

@@ -14,11 +14,23 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+type (
+	// ServerSpecifiedTimeoutConf defines specified timeout for gRPC method.
+	ServerSpecifiedTimeoutConf struct {
+		FullMethod string
+		Timeout    time.Duration
+	}
+
+	specifiedTimeoutCache map[string]time.Duration
+)
+
 // UnaryTimeoutInterceptor returns a func that sets timeout to incoming unary requests.
-func UnaryTimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor {
+func UnaryTimeoutInterceptor(timeout time.Duration, specifiedTimeouts ...ServerSpecifiedTimeoutConf) grpc.UnaryServerInterceptor {
+	cache := cacheSpecifiedTimeout(specifiedTimeouts)
 	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo,
 		handler grpc.UnaryHandler) (any, error) {
-		ctx, cancel := context.WithTimeout(ctx, timeout)
+		t := getTimeoutByUnaryServerInfo(info, timeout, cache)
+		ctx, cancel := context.WithTimeout(ctx, t)
 		defer cancel()
 
 		var resp any
@@ -59,3 +71,28 @@ func UnaryTimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor
 		}
 	}
 }
+
+func cacheSpecifiedTimeout(specifiedTimeouts []ServerSpecifiedTimeoutConf) specifiedTimeoutCache {
+	cache := make(specifiedTimeoutCache, len(specifiedTimeouts))
+	for _, st := range specifiedTimeouts {
+		if st.FullMethod != "" {
+			cache[st.FullMethod] = st.Timeout
+		}
+	}
+
+	return cache
+}
+
+func getTimeoutByUnaryServerInfo(info *grpc.UnaryServerInfo, defaultTimeout time.Duration, specifiedTimeout specifiedTimeoutCache) time.Duration {
+	if ts, ok := info.Server.(TimeoutStrategy); ok {
+		return ts.GetTimeoutByFullMethod(info.FullMethod, defaultTimeout)
+	} else if v, ok := specifiedTimeout[info.FullMethod]; ok {
+		return v
+	}
+
+	return defaultTimeout
+}
+
+type TimeoutStrategy interface {
+	GetTimeoutByFullMethod(fullMethod string, defaultTimeout time.Duration) time.Duration
+}

zrpc/internal/serverinterceptors/timeoutinterceptor_test.go

@@ -12,6 +12,11 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+var (
+	deadlineExceededErr = status.Error(codes.DeadlineExceeded, context.DeadlineExceeded.Error())
+	canceledErr         = status.Error(codes.Canceled, context.Canceled.Error())
+)
+
 func TestUnaryTimeoutInterceptor(t *testing.T) {
 	interceptor := UnaryTimeoutInterceptor(time.Millisecond * 10)
 	_, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
@@ -68,7 +73,7 @@ func TestUnaryTimeoutInterceptor_timeoutExpire(t *testing.T) {
 		return nil, nil
 	})
 	wg.Wait()
-	assert.EqualValues(t, status.Error(codes.DeadlineExceeded, context.DeadlineExceeded.Error()), err)
+	assert.EqualValues(t, deadlineExceededErr, err)
 }
 
 func TestUnaryTimeoutInterceptor_cancel(t *testing.T) {
@@ -88,5 +93,171 @@ func TestUnaryTimeoutInterceptor_cancel(t *testing.T) {
 	})
 
 	wg.Wait()
-	assert.EqualValues(t, status.Error(codes.Canceled, context.Canceled.Error()), err)
+	assert.EqualValues(t, canceledErr, err)
+}
+
+type tempServer struct {
+	timeout time.Duration
+}
+
+func (s *tempServer) run(duration time.Duration) {
+	time.Sleep(duration)
+}
+func (s *tempServer) GetTimeoutByFullMethod(fullMethod string, defaultTimeout time.Duration) time.Duration {
+	if fullMethod == "/" {
+		return defaultTimeout
+	}
+
+	return s.timeout
+}
+
+func TestUnaryTimeoutInterceptor_TimeoutStrategy(t *testing.T) {
+	type args struct {
+		interceptorTimeout time.Duration
+		contextTimeout     time.Duration
+		serverTimeout      time.Duration
+		runTime            time.Duration
+
+		fullMethod string
+	}
+	var tests = []struct {
+		name    string
+		args    args
+		wantErr error
+	}{
+		{
+			name: "do not timeout with interceptor timeout",
+			args: args{
+				interceptorTimeout: time.Second,
+				contextTimeout:     time.Second * 5,
+				serverTimeout:      time.Second * 3,
+				runTime:            time.Millisecond * 50,
+				fullMethod:         "/",
+			},
+			wantErr: nil,
+		},
+		{
+			name: "do not timeout with timeout strategy",
+			args: args{
+				interceptorTimeout: time.Second,
+				contextTimeout:     time.Second * 5,
+				serverTimeout:      time.Second * 3,
+				runTime:            time.Second * 2,
+				fullMethod:         "/2s",
+			},
+			wantErr: nil,
+		},
+		{
+			name: "timeout with interceptor timeout",
+			args: args{
+				interceptorTimeout: time.Second,
+				contextTimeout:     time.Second * 5,
+				serverTimeout:      time.Second * 3,
+				runTime:            time.Second * 2,
+				fullMethod:         "/",
+			},
+			wantErr: deadlineExceededErr,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			interceptor := UnaryTimeoutInterceptor(tt.args.interceptorTimeout)
+			ctx, cancel := context.WithTimeout(context.Background(), tt.args.contextTimeout)
+			defer cancel()
+
+			svr := &tempServer{timeout: tt.args.serverTimeout}
+
+			_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
+				Server:     svr,
+				FullMethod: tt.args.fullMethod,
+			}, func(ctx context.Context, req interface{}) (interface{}, error) {
+				svr.run(tt.args.runTime)
+				return nil, nil
+			})
+			t.Logf("error: %+v", err)
+
+			assert.EqualValues(t, tt.wantErr, err)
+		})
+	}
+}
+
+func TestUnaryTimeoutInterceptor_SpecifiedTimeout(t *testing.T) {
+	type args struct {
+		interceptorTimeout time.Duration
+		contextTimeout     time.Duration
+		method             string
+		methodTimeout      time.Duration
+		runTime            time.Duration
+	}
+	var tests = []struct {
+		name    string
+		args    args
+		wantErr error
+	}{
+		{
+			name: "do not timeout without set timeout for full method",
+			args: args{
+				interceptorTimeout: time.Second,
+				contextTimeout:     time.Second * 5,
+				method:             "/run",
+				runTime:            time.Millisecond * 50,
+			},
+			wantErr: nil,
+		},
+		{
+			name: "do not timeout with set timeout for full method",
+			args: args{
+				interceptorTimeout: time.Second,
+				contextTimeout:     time.Second * 5,
+				method:             "/run/do_not_timeout",
+				methodTimeout:      time.Second * 3,
+				runTime:            time.Second * 2,
+			},
+			wantErr: nil,
+		},
+		{
+			name: "timeout with set timeout for full method",
+			args: args{
+				interceptorTimeout: time.Second,
+				contextTimeout:     time.Second * 5,
+				method:             "/run/timeout",
+				methodTimeout:      time.Millisecond * 100,
+				runTime:            time.Millisecond * 500,
+			},
+			wantErr: deadlineExceededErr,
+		},
+	}
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			var specifiedTimeouts []ServerSpecifiedTimeoutConf
+			if tt.args.methodTimeout > 0 {
+				specifiedTimeouts = []ServerSpecifiedTimeoutConf{
+					{
+						FullMethod: tt.args.method,
+						Timeout:    tt.args.methodTimeout,
+					},
+				}
+			}
+
+			interceptor := UnaryTimeoutInterceptor(tt.args.interceptorTimeout, specifiedTimeouts...)
+			ctx, cancel := context.WithTimeout(context.Background(), tt.args.contextTimeout)
+			defer cancel()
+
+			_, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
+				FullMethod: tt.args.method,
+			}, func(ctx context.Context, req interface{}) (interface{}, error) {
+				time.Sleep(tt.args.runTime)
+				return nil, nil
+			})
+			t.Logf("error: %+v", err)
+
+			assert.EqualValues(t, tt.wantErr, err)
+		})
+	}
 }