From 16ca7bc026edfb95f5e7a4be040d504794a734c5 Mon Sep 17 00:00:00 2001
From: xiezhengyao <xiezhengyao@bytedance.com>
Date: Tue, 21 Jan 2025 21:22:03 +0800
Subject: [PATCH] optimize: the graceful shutdown logic to avoid errors when a
 connection in idle state receives new requests after being closed

---
 connection_impl.go    |  1 +
 connection_onevent.go |  6 ------
 netpoll_server.go     | 12 ++++++++----
 3 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/connection_impl.go b/connection_impl.go
index 2d4f3f48..1908b7e6 100644
--- a/connection_impl.go
+++ b/connection_impl.go
@@ -54,6 +54,7 @@ type connection struct {
 	maxSize         int       // The maximum size of data between two Release().
 	bookSize        int       // The size of data that can be read at once.
 	state           connState // Connection state should be changed sequentially.
+	wasIdle         bool      // It's for double checking a conn is really idle, avoiding closing a conn that may receive traffic immediately.
 }
 
 var (
diff --git a/connection_onevent.go b/connection_onevent.go
index 8efcd98d..1a0c93df 100644
--- a/connection_onevent.go
+++ b/connection_onevent.go
@@ -39,11 +39,6 @@ func disableGopool() error {
 
 // ------------------------------------ implement OnPrepare, OnRequest, CloseCallback ------------------------------------
 
-type gracefulExit interface {
-	isIdle() (yes bool)
-	Close() (err error)
-}
-
 // onEvent is the collection of event processing.
 // OnPrepare, OnRequest, CloseCallback share the lock processing,
 // which is a CAS lock and can only be cleared by OnRequest.
@@ -310,7 +305,6 @@ func (c *connection) register() (err error) {
 	return nil
 }
 
-// isIdle implements gracefulExit.
 func (c *connection) isIdle() (yes bool) {
 	return c.isUnlock(processing) &&
 		c.inputBuffer.IsEmpty() &&
diff --git a/netpoll_server.go b/netpoll_server.go
index 61ed09c2..368a35eb 100644
--- a/netpoll_server.go
+++ b/netpoll_server.go
@@ -66,9 +66,13 @@ func (s *server) Close(ctx context.Context) error {
 	for {
 		activeConn := 0
 		s.connections.Range(func(key, value interface{}) bool {
-			conn, ok := value.(gracefulExit)
+			conn, ok := value.(*connection)
 			if !ok || conn.isIdle() {
-				value.(Connection).Close()
+				if conn.wasIdle {
+					value.(Connection).Close()
+				} else {
+					conn.wasIdle = true
+				}
 			} else {
 				activeConn++
 			}
@@ -83,8 +87,8 @@ func (s *server) Close(ctx context.Context) error {
 		waitTime := time.Millisecond * time.Duration(activeConn)
 		if waitTime > time.Second { // max wait time is 1000 ms
 			waitTime = time.Millisecond * 1000
-		} else if waitTime < time.Millisecond*50 { // min wait time is 50 ms
-			waitTime = time.Millisecond * 50
+		} else if waitTime < time.Millisecond*200 { // min wait time is 200 ms
+			waitTime = time.Millisecond * 200
 		}
 		select {
 		case <-ctx.Done():