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

Co-authored-by: hanzijian <hanzijian@52tt.com> Co-authored-by: Kevin Wan <wanjunfeng@gmail.com>
2023-10-25 21:01:57 +08:00
parent 2a335c7608
commit 842c4d81cc
10 changed files with 378 additions and 29 deletions
--- a/zrpc/internal/serverinterceptors/timeoutinterceptor.go
+++ b/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
+}
--- a/zrpc/internal/serverinterceptors/timeoutinterceptor_test.go
+++ b/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)
+		})
+	}
 }