From cb8d9d413aa94a94ff43212140f2f8cd402c063a Mon Sep 17 00:00:00 2001
From: masonchen2014 <32946302+masonchen2014@users.noreply.github.com>
Date: Sat, 24 Jul 2021 21:51:46 +0800
Subject: [PATCH] simplify timeoutinterceptor (#840)

Co-authored-by: chenmusheng <chenmusheng@laoyuegou.com>
---
 zrpc/client_test.go                           | 10 +++++++-
 .../clientinterceptors/timeoutinterceptor.go  | 23 +------------------
 .../timeoutinterceptor_test.go                | 19 ---------------
 zrpc/internal/mock/depositserver.go           |  2 ++
 4 files changed, 12 insertions(+), 42 deletions(-)

diff --git a/zrpc/client_test.go b/zrpc/client_test.go
index 90baa997..dc1f08e1 100644
--- a/zrpc/client_test.go
+++ b/zrpc/client_test.go
@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/tal-tech/go-zero/core/logx"
@@ -58,6 +59,13 @@ func TestDepositServer_Deposit(t *testing.T) {
 			codes.OK,
 			"",
 		},
+		{
+			"valid request with long handling time",
+			2000.00,
+			nil,
+			codes.DeadlineExceeded,
+			fmt.Sprintf("context deadline exceeded"),
+		},
 	}
 
 	directClient := MustNewClient(
@@ -79,7 +87,7 @@ func TestDepositServer_Deposit(t *testing.T) {
 			func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
 				invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
 				return invoker(ctx, method, req, reply, cc, opts...)
-			}))
+			}), WithTimeout(1000*time.Millisecond))
 	assert.Nil(t, err)
 	clients := []Client{
 		directClient,
diff --git a/zrpc/internal/clientinterceptors/timeoutinterceptor.go b/zrpc/internal/clientinterceptors/timeoutinterceptor.go
index 81cb88aa..95880f1f 100644
--- a/zrpc/internal/clientinterceptors/timeoutinterceptor.go
+++ b/zrpc/internal/clientinterceptors/timeoutinterceptor.go
@@ -17,27 +17,6 @@ func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
 
 		ctx, cancel := context.WithTimeout(ctx, timeout)
 		defer cancel()
-
-		// create channel with buffer size 1 to avoid goroutine leak
-		done := make(chan error, 1)
-		panicChan := make(chan interface{}, 1)
-		go func() {
-			defer func() {
-				if p := recover(); p != nil {
-					panicChan <- p
-				}
-			}()
-
-			done <- invoker(ctx, method, req, reply, cc, opts...)
-		}()
-
-		select {
-		case p := <-panicChan:
-			panic(p)
-		case err := <-done:
-			return err
-		case <-ctx.Done():
-			return ctx.Err()
-		}
+		return invoker(ctx, method, req, reply, cc, opts...)
 	}
 }
diff --git a/zrpc/internal/clientinterceptors/timeoutinterceptor_test.go b/zrpc/internal/clientinterceptors/timeoutinterceptor_test.go
index 12941b17..d45a8511 100644
--- a/zrpc/internal/clientinterceptors/timeoutinterceptor_test.go
+++ b/zrpc/internal/clientinterceptors/timeoutinterceptor_test.go
@@ -49,25 +49,6 @@ func TestTimeoutInterceptor_timeout(t *testing.T) {
 	assert.Nil(t, err)
 }
 
-func TestTimeoutInterceptor_timeoutExpire(t *testing.T) {
-	const timeout = time.Millisecond * 10
-	interceptor := TimeoutInterceptor(timeout)
-	ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond)
-	defer cancel()
-	var wg sync.WaitGroup
-	wg.Add(1)
-	cc := new(grpc.ClientConn)
-	err := interceptor(ctx, "/foo", nil, nil, cc,
-		func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
-			opts ...grpc.CallOption) error {
-			defer wg.Done()
-			time.Sleep(time.Millisecond * 50)
-			return nil
-		})
-	wg.Wait()
-	assert.Equal(t, context.DeadlineExceeded, err)
-}
-
 func TestTimeoutInterceptor_panic(t *testing.T) {
 	timeouts := []time.Duration{0, time.Millisecond * 10}
 	for _, timeout := range timeouts {
diff --git a/zrpc/internal/mock/depositserver.go b/zrpc/internal/mock/depositserver.go
index 658a1b77..6f7a6275 100644
--- a/zrpc/internal/mock/depositserver.go
+++ b/zrpc/internal/mock/depositserver.go
@@ -2,6 +2,7 @@ package mock
 
 import (
 	"context"
+	"time"
 
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -16,5 +17,6 @@ func (*DepositServer) Deposit(ctx context.Context, req *DepositRequest) (*Deposi
 		return nil, status.Errorf(codes.InvalidArgument, "cannot deposit %v", req.GetAmount())
 	}
 
+	time.Sleep(time.Duration(req.GetAmount()) * time.Millisecond)
 	return &DepositResponse{Ok: true}, nil
 }