From 3c1a00c64522cff3e0ed607ea58b78c812c2dc78 Mon Sep 17 00:00:00 2001
From: crazycs <crazycs520@gmail.com>
Date: Mon, 26 Nov 2018 14:39:01 +0800
Subject: [PATCH 1/3] server: close connection when tidb server closed.

---
 server/server.go    | 14 ++++++++++++++
 tidb-server/main.go |  2 ++
 2 files changed, 16 insertions(+)

diff --git a/server/server.go b/server/server.go
index 49455fe1d4e5..fe30ad1ed2a1 100644
--- a/server/server.go
+++ b/server/server.go
@@ -351,6 +351,10 @@ func (s *Server) Kill(connectionID uint64, query bool) {
 		return
 	}
 
+	killConn(conn, query)
+}
+
+func killConn(conn *clientConn, query bool) {
 	conn.mu.RLock()
 	cancelFunc := conn.mu.cancelFunc
 	conn.mu.RUnlock()
@@ -365,6 +369,16 @@ func (s *Server) Kill(connectionID uint64, query bool) {
 	}
 }
 
+func (s *Server) KillAllConnections() {
+	s.rwlock.Lock()
+	defer s.rwlock.Unlock()
+	log.Info("[server] kill all connections.")
+
+	for _, conn := range s.clients {
+		killConn(conn, false)
+	}
+}
+
 // GracefulDown waits all clients to close.
 func (s *Server) GracefulDown() {
 	log.Info("[server] graceful shutdown.")
diff --git a/tidb-server/main.go b/tidb-server/main.go
index a1062d9fa59a..2960db0210cd 100644
--- a/tidb-server/main.go
+++ b/tidb-server/main.go
@@ -513,6 +513,8 @@ func closeDomainAndStorage() {
 func cleanup() {
 	if graceful {
 		svr.GracefulDown()
+	} else {
+		svr.KillAllConnections()
 	}
 	closeDomainAndStorage()
 }

From d4211261fda04ed450059a34dc2078b2192fb6ae Mon Sep 17 00:00:00 2001
From: crazycs <crazycs520@gmail.com>
Date: Mon, 26 Nov 2018 14:59:59 +0800
Subject: [PATCH 2/3] add comment

---
 server/server.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/server/server.go b/server/server.go
index fe30ad1ed2a1..4b93f0a8f9d0 100644
--- a/server/server.go
+++ b/server/server.go
@@ -369,6 +369,7 @@ func killConn(conn *clientConn, query bool) {
 	}
 }
 
+// KillAllConnections kills all connections when server is not gracefully shutdown.
 func (s *Server) KillAllConnections() {
 	s.rwlock.Lock()
 	defer s.rwlock.Unlock()

From 34fee85c3fa4fce57fd47d7a95cc5c289e481513 Mon Sep 17 00:00:00 2001
From: crazycs <crazycs520@gmail.com>
Date: Mon, 26 Nov 2018 19:49:15 +0800
Subject: [PATCH 3/3] address comment

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

diff --git a/server/server.go b/server/server.go
index 4b93f0a8f9d0..5e516493673b 100644
--- a/server/server.go
+++ b/server/server.go
@@ -355,18 +355,18 @@ func (s *Server) Kill(connectionID uint64, query bool) {
 }
 
 func killConn(conn *clientConn, query bool) {
+	if !query {
+		// Mark the client connection status as WaitShutdown, when the goroutine detect
+		// this, it will end the dispatch loop and exit.
+		atomic.StoreInt32(&conn.status, connStatusWaitShutdown)
+	}
+
 	conn.mu.RLock()
 	cancelFunc := conn.mu.cancelFunc
 	conn.mu.RUnlock()
 	if cancelFunc != nil {
 		cancelFunc()
 	}
-
-	if !query {
-		// Mark the client connection status as WaitShutdown, when the goroutine detect
-		// this, it will end the dispatch loop and exit.
-		atomic.StoreInt32(&conn.status, connStatusWaitShutdown)
-	}
 }
 
 // KillAllConnections kills all connections when server is not gracefully shutdown.