From 767267a1845a4d9bfab7835142bc9beb9f491489 Mon Sep 17 00:00:00 2001
From: Hangjie Mo <mohangjie1995@gmail.com>
Date: Tue, 27 Jun 2023 10:32:34 +0800
Subject: [PATCH] server: exit tidb-server directly when all connections not in
 txn (#44953)

close pingcap/tidb#44951
---
 server/conn.go              |  7 +++++++
 server/conn_test.go         | 30 +++++++++++++++++++++++++++++-
 server/server.go            |  9 +++------
 tidb-server/main.go         |  6 +++---
 util/signal/signal_posix.go |  4 ++--
 5 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/server/conn.go b/server/conn.go
index e2c3836eb54e3..c7378c32ed835 100644
--- a/server/conn.go
+++ b/server/conn.go
@@ -1160,6 +1160,13 @@ func (cc *clientConn) Run(ctx context.Context) {
 			return
 		}
 
+		// Should check InTxn() to avoid execute `begin` stmt.
+		if cc.server.inShutdownMode.Load() {
+			if !cc.ctx.GetSessionVars().InTxn() {
+				return
+			}
+		}
+
 		if !cc.CompareAndSwapStatus(connStatusReading, connStatusDispatching) {
 			return
 		}
diff --git a/server/conn_test.go b/server/conn_test.go
index 609a32576d726..0acb7f9342f30 100644
--- a/server/conn_test.go
+++ b/server/conn_test.go
@@ -763,7 +763,7 @@ func TestConnExecutionTimeout(t *testing.T) {
 }
 
 func TestShutDown(t *testing.T) {
-	store := testkit.CreateMockStore(t)
+	store, dom := testkit.CreateMockStoreAndDomain(t)
 
 	cc := &clientConn{}
 	se, err := session.CreateSession4Test(store)
@@ -775,6 +775,34 @@ func TestShutDown(t *testing.T) {
 	// assert ErrQueryInterrupted
 	err = cc.handleQuery(context.Background(), "select 1")
 	require.Equal(t, exeerrors.ErrQueryInterrupted, err)
+
+	cfg := newTestConfig()
+	cfg.Port = 0
+	cfg.Status.StatusPort = 0
+	drv := NewTiDBDriver(store)
+	srv, err := NewServer(cfg, drv)
+	require.NoError(t, err)
+	srv.SetDomain(dom)
+
+	cc = &clientConn{server: srv}
+	cc.setCtx(tc)
+
+	// test in txn
+	srv.clients[dom.NextConnID()] = cc
+	cc.getCtx().GetSessionVars().SetInTxn(true)
+
+	waitTime := 100 * time.Millisecond
+	begin := time.Now()
+	srv.DrainClients(waitTime, waitTime)
+	require.Greater(t, time.Since(begin), waitTime)
+
+	// test not in txn
+	srv.clients[dom.NextConnID()] = cc
+	cc.getCtx().GetSessionVars().SetInTxn(false)
+
+	begin = time.Now()
+	srv.DrainClients(waitTime, waitTime)
+	require.Less(t, time.Since(begin), waitTime)
 }
 
 type snapshotCache interface {
diff --git a/server/server.go b/server/server.go
index c2baaaf2b6465..0594785500870 100644
--- a/server/server.go
+++ b/server/server.go
@@ -34,7 +34,6 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io"
-	"math/rand"
 	"net"
 	"net/http"         //nolint:goimports
 	_ "net/http/pprof" // #nosec G108 for pprof
@@ -327,9 +326,6 @@ func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
 		}
 	}
 
-	// Init rand seed for randomBuf()
-	rand.Seed(time.Now().UTC().UnixNano())
-
 	variable.RegisterStatistics(s)
 
 	return s, nil
@@ -549,8 +545,6 @@ func (s *Server) closeListener() {
 	metrics.ServerEventCounter.WithLabelValues(metrics.EventClose).Inc()
 }
 
-var gracefulCloseConnectionsTimeout = 15 * time.Second
-
 // Close closes the server.
 func (s *Server) Close() {
 	s.startShutdown()
@@ -917,6 +911,9 @@ func (s *Server) DrainClients(drainWait time.Duration, cancelWait time.Duration)
 	go func() {
 		defer close(allDone)
 		for _, conn := range conns {
+			if !conn.getCtx().GetSessionVars().InTxn() {
+				continue
+			}
 			select {
 			case <-conn.quit:
 			case <-quitWaitingForConns:
diff --git a/tidb-server/main.go b/tidb-server/main.go
index 48456c371e09c..22600cb6c6266 100644
--- a/tidb-server/main.go
+++ b/tidb-server/main.go
@@ -245,9 +245,9 @@ func main() {
 	terror.RegisterFinish()
 
 	exited := make(chan struct{})
-	signal.SetupSignalHandler(func(graceful bool) {
+	signal.SetupSignalHandler(func() {
 		svr.Close()
-		cleanup(svr, storage, dom, graceful)
+		cleanup(svr, storage, dom)
 		cpuprofile.StopCPUProfiler()
 		resourcemanager.InstanceResourceManager.Stop()
 		close(exited)
@@ -854,7 +854,7 @@ func closeDomainAndStorage(storage kv.Storage, dom *domain.Domain) {
 // We should better provider a dynamic way to set this value.
 var gracefulCloseConnectionsTimeout = 15 * time.Second
 
-func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain, _ bool) {
+func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain) {
 	dom.StopAutoAnalyze()
 
 	drainClientWait := gracefulCloseConnectionsTimeout
diff --git a/util/signal/signal_posix.go b/util/signal/signal_posix.go
index 8e9e25308c0ef..8b416b93c543e 100644
--- a/util/signal/signal_posix.go
+++ b/util/signal/signal_posix.go
@@ -27,7 +27,7 @@ import (
 )
 
 // SetupSignalHandler setup signal handler for TiDB Server
-func SetupSignalHandler(shutdownFunc func(bool)) {
+func SetupSignalHandler(shutdownFunc func()) {
 	usrDefSignalChan := make(chan os.Signal, 1)
 
 	signal.Notify(usrDefSignalChan, syscall.SIGUSR1)
@@ -52,6 +52,6 @@ func SetupSignalHandler(shutdownFunc func(bool)) {
 	go func() {
 		sig := <-closeSignalChan
 		logutil.BgLogger().Info("got signal to exit", zap.Stringer("signal", sig))
-		shutdownFunc(sig != syscall.SIGHUP)
+		shutdownFunc()
 	}()
 }