From bc0bf3c4c7c130a5f1ae83af723ce5b8e12a546c Mon Sep 17 00:00:00 2001
From: Sainnhe Park <i@sainnhe.dev>
Date: Sat, 11 Jan 2025 13:27:03 +0800
Subject: [PATCH] fix(utils): wait for goroutine locks in Stop()

---
 pkg/remote/trans/netpollmux/server_handler.go | 5 +----
 pkg/remote/trans/nphttp2/server_handler.go    | 5 -----
 server/server.go                              | 4 ++++
 3 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/pkg/remote/trans/netpollmux/server_handler.go b/pkg/remote/trans/netpollmux/server_handler.go
index 3f3f846a00..2a314cda3c 100644
--- a/pkg/remote/trans/netpollmux/server_handler.go
+++ b/pkg/remote/trans/netpollmux/server_handler.go
@@ -27,7 +27,6 @@ import (
 
 	"github.com/cloudwego/netpoll"
 
-	goroutinelock "github.com/cloudwego/kitex/internal/utils/goroutine_lock"
 	"github.com/cloudwego/kitex/pkg/endpoint"
 	"github.com/cloudwego/kitex/pkg/gofunc"
 	"github.com/cloudwego/kitex/pkg/kerrors"
@@ -366,9 +365,7 @@ func (t *svrTransHandler) GracefulShutdown(ctx context.Context) error {
 			}
 			return true
 		})
-		// 4. waiting for goroutine locks to be released
-		goroutinelock.GoroutineWg.Wait()
-		// 5. waiting all crrst packets received by client
+		// 4. waiting all crrst packets received by client
 		deadline := time.Now().Add(defaultExitWaitGracefulShutdownTime)
 		ticker := time.NewTicker(defaultExitWaitGracefulShutdownTime / 10)
 		defer ticker.Stop()
diff --git a/pkg/remote/trans/nphttp2/server_handler.go b/pkg/remote/trans/nphttp2/server_handler.go
index 6d2f095fcc..601689a63a 100644
--- a/pkg/remote/trans/nphttp2/server_handler.go
+++ b/pkg/remote/trans/nphttp2/server_handler.go
@@ -31,7 +31,6 @@ import (
 
 	"github.com/cloudwego/netpoll"
 
-	goroutinelock "github.com/cloudwego/kitex/internal/utils/goroutine_lock"
 	"github.com/cloudwego/kitex/pkg/endpoint"
 	"github.com/cloudwego/kitex/pkg/gofunc"
 	"github.com/cloudwego/kitex/pkg/kerrors"
@@ -408,7 +407,6 @@ func (t *svrTransHandler) GracefulShutdown(ctx context.Context) error {
 			}
 			t.mu.Unlock()
 			if activeNums == 0 {
-				goroutinelock.GoroutineWg.Wait()
 				return nil
 			}
 		case <-graceTimer.C:
@@ -418,9 +416,6 @@ func (t *svrTransHandler) GracefulShutdown(ctx context.Context) error {
 				svrTrans := elem.Value.(*SvrTrans)
 				svrTrans.tr.Close()
 			}
-			// Goroutine Locks must wait after all connections are closed, otherwise new connections might be created
-			// while waiting for goroutine locks.
-			goroutinelock.GoroutineWg.Wait()
 			t.mu.Unlock()
 			return nil
 		}
diff --git a/server/server.go b/server/server.go
index 025e493173..e4984d092b 100644
--- a/server/server.go
+++ b/server/server.go
@@ -30,6 +30,7 @@ import (
 	"github.com/cloudwego/localsession/backup"
 
 	internal_server "github.com/cloudwego/kitex/internal/server"
+	goroutinelock "github.com/cloudwego/kitex/internal/utils/goroutine_lock"
 	"github.com/cloudwego/kitex/pkg/acl"
 	"github.com/cloudwego/kitex/pkg/diagnosis"
 	"github.com/cloudwego/kitex/pkg/discovery"
@@ -319,6 +320,9 @@ func (s *server) Stop() (err error) {
 			}
 			s.svr = nil
 		}
+		// Goroutine Locks must wait after all connections are closed, otherwise new connections might be created while
+		// waiting for goroutine locks.
+		goroutinelock.GoroutineWg.Wait()
 	})
 	return
 }