From 1d627103f7209b5587ea47fe4f85ca513169b050 Mon Sep 17 00:00:00 2001
From: Jiahao Huang <july2993@gmail.com>
Date: Mon, 6 Mar 2023 21:26:20 +0800
Subject: [PATCH] Close lisener in server.Close

---
 server/server.go | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/server/server.go b/server/server.go
index 87f1ac39024e5..72122bdda6a72 100644
--- a/server/server.go
+++ b/server/server.go
@@ -559,6 +559,10 @@ var gracefulCloseConnectionsTimeout = 15 * time.Second
 // Close closes the server.
 func (s *Server) Close() {
 	s.startShutdown()
+	s.rwlock.Lock() // // prevent new connections
+	defer s.rwlock.Unlock()
+	s.inShutdownMode.Store(true)
+	s.closeListener()
 }
 
 func (s *Server) registerConn(conn *clientConn) bool {
@@ -874,8 +878,6 @@ func (s *Server) DrainClients(drainWait time.Duration, cancelWait time.Duration)
 	conns := make(map[uint64]*clientConn)
 
 	s.rwlock.Lock()
-	s.inShutdownMode.Store(true)
-	s.closeListener()
 	for k, v := range s.clients {
 		conns[k] = v
 	}