server: revert "#48905" for release-6.5 (#49119)

close #49098

server/BUILD.bazel

@@ -116,7 +116,6 @@ go_library(
         "@org_golang_google_grpc//channelz/service",
         "@org_golang_google_grpc//keepalive",
         "@org_golang_google_grpc//peer",
-        "@org_uber_go_atomic//:atomic",
         "@org_uber_go_zap//:zap",
     ],
 )

server/conn.go

@@ -172,7 +172,6 @@ func newClientConn(s *Server) *clientConn {
 		status:       connStatusDispatching,
 		lastActive:   time.Now(),
 		authPlugin:   mysql.AuthNativePassword,
-		quit:         make(chan struct{}),
 		ppEnabled:    s.cfg.ProxyProtocol.Networks != "",
 	}
 }
@@ -216,8 +215,6 @@ type clientConn struct {
 		sync.RWMutex
 		cancelFunc context.CancelFunc
 	}
-	// quit is close once clientConn quit Run().
-	quit       chan struct{}
 	extensions *extension.SessionExtensions
 
 	// Proxy Protocol Enabled
@@ -1096,12 +1093,6 @@ func (cc *clientConn) Run(ctx context.Context) {
 			terror.Log(err)
 			metrics.PanicCounter.WithLabelValues(metrics.LabelSession).Inc()
 		}
-		if atomic.LoadInt32(&cc.status) != connStatusShutdown {
-			err := cc.Close()
-			terror.Log(err)
-		}
-
-		close(cc.quit)
 	}()
 
 	// Usually, client connection status changes between [dispatching] <=> [reading].
@@ -1110,13 +1101,6 @@ func (cc *clientConn) Run(ctx context.Context) {
 	// The client connection would detect the events when it fails to change status
 	// by CAS operation, it would then take some actions accordingly.
 	for {
-		// Close connection between txn when we are going to shutdown server.
-		if cc.server.inShutdownMode.Load() {
-			if !cc.ctx.GetSessionVars().InTxn() {
-				return
-			}
-		}
-
 		if !atomic.CompareAndSwapInt32(&cc.status, connStatusDispatching, connStatusReading) ||
 			// The judge below will not be hit by all means,
 			// But keep it stayed as a reminder and for the code reference for connStatusWaitShutdown.
@@ -1126,7 +1110,6 @@ func (cc *clientConn) Run(ctx context.Context) {
 
 		cc.alloc.Reset()
 		// close connection when idle time is more than wait_timeout
-		// default 28800(8h), FIXME: should not block at here when we kill the connection.
 		waitTimeout := cc.getSessionVarsWaitTimeout(ctx)
 		cc.pkt.setReadTimeout(time.Duration(waitTimeout) * time.Second)
 		start := time.Now()
@@ -1213,6 +1196,22 @@ func (cc *clientConn) Run(ctx context.Context) {
 	}
 }
 
+// ShutdownOrNotify will Shutdown this client connection, or do its best to notify.
+func (cc *clientConn) ShutdownOrNotify() bool {
+	if (cc.ctx.Status() & mysql.ServerStatusInTrans) > 0 {
+		return false
+	}
+	// If the client connection status is reading, it's safe to shutdown it.
+	if atomic.CompareAndSwapInt32(&cc.status, connStatusReading, connStatusShutdown) {
+		return true
+	}
+	// If the client connection status is dispatching, we can't shutdown it immediately,
+	// so set the status to WaitShutdown as a notification, the loop in clientConn.Run
+	// will detect it and then exit.
+	atomic.StoreInt32(&cc.status, connStatusWaitShutdown)
+	return false
+}
+
 func errStrForLog(err error, enableRedactLog bool) string {
 	if enableRedactLog {
 		// currently, only ErrParse is considered when enableRedactLog because it may contain sensitive information like

server/conn_test.go

@@ -778,6 +778,31 @@ func TestShutDown(t *testing.T) {
 	require.Equal(t, executor.ErrQueryInterrupted, err)
 }
 
+func TestShutdownOrNotify(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	se, err := session.CreateSession4Test(store)
+	require.NoError(t, err)
+	tc := &TiDBContext{
+		Session: se,
+		stmts:   make(map[int]*TiDBStatement),
+	}
+	cc := &clientConn{
+		connectionID: 1,
+		server: &Server{
+			capability: defaultCapability,
+		},
+		status: connStatusWaitShutdown,
+	}
+	cc.setCtx(tc)
+	require.False(t, cc.ShutdownOrNotify())
+	cc.status = connStatusReading
+	require.True(t, cc.ShutdownOrNotify())
+	require.Equal(t, connStatusShutdown, cc.status)
+	cc.status = connStatusDispatching
+	require.False(t, cc.ShutdownOrNotify())
+	require.Equal(t, connStatusWaitShutdown, cc.status)
+}
+
 type snapshotCache interface {
 	SnapCacheHitCount() int
 }

server/http_status.go

@@ -539,7 +539,7 @@ func (s *Server) handleStatus(w http.ResponseWriter, req *http.Request) {
 	// If the server is in the process of shutting down, return a non-200 status.
 	// It is important not to return status{} as acquiring the s.ConnectionCount()
 	// acquires a lock that may already be held by the shutdown process.
-	if !s.health.Load() {
+	if s.inShutdownMode {
 		w.WriteHeader(http.StatusInternalServerError)
 		return
 	}