From 9e6c2ba2c02aa67fd06b69719effa753dedc2a3f Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sun, 21 Mar 2021 16:54:34 +0800
Subject: [PATCH] avoid goroutine leak after timeout (#575)

---
 core/fx/timeout.go                                     | 4 ++--
 zrpc/internal/clientinterceptors/timeoutinterceptor.go | 4 ++--
 zrpc/internal/serverinterceptors/timeoutinterceptor.go | 1 +
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/core/fx/timeout.go b/core/fx/timeout.go
index afd0748c..33f70e8a 100644
--- a/core/fx/timeout.go
+++ b/core/fx/timeout.go
@@ -26,7 +26,8 @@ func DoWithTimeout(fn func() error, timeout time.Duration, opts ...DoOption) err
 	ctx, cancel := contextx.ShrinkDeadline(parentCtx, timeout)
 	defer cancel()
 
-	done := make(chan error)
+	// create channel with buffer size 1 to avoid goroutine leak
+	done := make(chan error, 1)
 	panicChan := make(chan interface{}, 1)
 	go func() {
 		defer func() {
@@ -35,7 +36,6 @@ func DoWithTimeout(fn func() error, timeout time.Duration, opts ...DoOption) err
 			}
 		}()
 		done <- fn()
-		close(done)
 	}()
 
 	select {
diff --git a/zrpc/internal/clientinterceptors/timeoutinterceptor.go b/zrpc/internal/clientinterceptors/timeoutinterceptor.go
index 90554cfb..c27678ff 100644
--- a/zrpc/internal/clientinterceptors/timeoutinterceptor.go
+++ b/zrpc/internal/clientinterceptors/timeoutinterceptor.go
@@ -19,7 +19,8 @@ func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
 		ctx, cancel := contextx.ShrinkDeadline(ctx, timeout)
 		defer cancel()
 
-		done := make(chan error)
+		// create channel with buffer size 1 to avoid goroutine leak
+		done := make(chan error, 1)
 		panicChan := make(chan interface{}, 1)
 		go func() {
 			defer func() {
@@ -29,7 +30,6 @@ func TimeoutInterceptor(timeout time.Duration) grpc.UnaryClientInterceptor {
 			}()
 
 			done <- invoker(ctx, method, req, reply, cc, opts...)
-			close(done)
 		}()
 
 		select {
diff --git a/zrpc/internal/serverinterceptors/timeoutinterceptor.go b/zrpc/internal/serverinterceptors/timeoutinterceptor.go
index 8f607195..8db1e9c8 100644
--- a/zrpc/internal/serverinterceptors/timeoutinterceptor.go
+++ b/zrpc/internal/serverinterceptors/timeoutinterceptor.go
@@ -20,6 +20,7 @@ func UnaryTimeoutInterceptor(timeout time.Duration) grpc.UnaryServerInterceptor
 		var err error
 		var lock sync.Mutex
 		done := make(chan struct{})
+		// create channel with buffer size 1 to avoid goroutine leak
 		panicChan := make(chan interface{}, 1)
 		go func() {
 			defer func() {