Keep the "graceful" behavior

When graceful, keep waiting all connection drained.
When not, try 15s like before.
Note after #37441 the "graceful" or not checking is changed.

server/server.go

@@ -555,12 +555,11 @@ 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()
-	drainClientWait := time.Second * 15
-	cancelClientWait := time.Second * 1
-	s.drainClients(drainClientWait, cancelClientWait)
 }
 
 func (s *Server) registerConn(conn *clientConn) bool {
@@ -867,9 +866,9 @@ func (s *Server) KillAllConnections() {
 	s.KillSysProcesses()
 }
 
-// drainWait drain all connections in drainWait.
+// DrainClients drain all connections in drainWait.
 // After drainWait duration, we kill all connections still not quit explicitly and wait for cancelWait.
-func (s *Server) drainClients(drainWait time.Duration, cancelWait time.Duration) error {
+func (s *Server) DrainClients(drainWait time.Duration, cancelWait time.Duration) {
 	logger := logutil.BgLogger()
 	logger.Info("start drain clients")
 
@@ -912,7 +911,6 @@ func (s *Server) drainClients(drainWait time.Duration, cancelWait time.Duration)
 		logger.Warn("some sessions do not quit in cancel wait time")
 	}
 
-	return nil
 }
 
 // ServerID implements SessionManager interface.

tidb-server/main.go

@@ -253,9 +253,9 @@ func main() {
 	terror.RegisterFinish()
 
 	exited := make(chan struct{})
-	signal.SetupSignalHandler(func(_ bool) {
+	signal.SetupSignalHandler(func(graceful bool) {
 		svr.Close()
-		cleanup(svr, storage, dom)
+		cleanup(svr, storage, dom, graceful)
 		cpuprofile.StopCPUProfiler()
 		resourcemanager.InstanceResourceManager.Stop()
 		close(exited)
@@ -857,8 +857,20 @@ func closeDomainAndStorage(storage kv.Storage, dom *domain.Domain) {
 	terror.Log(errors.Trace(err))
 }
 
-func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain) {
+var gracefulCloseConnectionsTimeout = 15 * time.Second
+
+func cleanup(svr *server.Server, storage kv.Storage, dom *domain.Domain, graceful bool) {
 	dom.StopAutoAnalyze()
+
+	var drainClientWait time.Duration
+	if graceful {
+		drainClientWait = 1<<63 - 1
+	} else {
+		drainClientWait = gracefulCloseConnectionsTimeout
+	}
+	cancelClientWait := time.Second * 1
+	svr.DrainClients(drainClientWait, cancelClientWait)
+
 	// Kill sys processes such as auto analyze. Otherwise, tidb-server cannot exit until auto analyze is finished.
 	// See https://github.com/pingcap/tidb/issues/40038 for details.
 	svr.KillSysProcesses()