server: address disconnect monitor review comments

Author: King-Dylan

pkg/server/conn.go

@@ -2090,27 +2090,37 @@ func setResourceGroupTaggerForMultiStmtPrefetch(snapshot kv.Snapshot, sqls strin
 	}
 }
 
+// setSQLKillerConnectionAlive installs a connection-liveness probe on the
+// session SQLKiller and starts a background monitor for the current statement.
+// The returned cleanup is idempotent and must be called when the statement is
+// done to stop the monitor and clear the probe.
 func (cc *clientConn) setSQLKillerConnectionAlive() func() {
 	fn := func() bool {
 		if cc.bufReadConn != nil {
+			// IsAlive returns 0 only when the connection is known dead. Treat
+			// unknown states as alive so we do not interrupt queries
+			// conservatively when the liveness check itself cannot run.
 			return cc.bufReadConn.IsAlive() != 0
 		}
 		return true
 	}
 	cc.ctx.GetSessionVars().SQLKiller.IsConnectionAlive.Store(&fn)
 	stopMonitor := make(chan struct{})
-	go cc.monitorConnectionAlive(fn, stopMonitor)
+	doneMonitor := make(chan struct{})
+	go cc.monitorConnectionAlive(fn, stopMonitor, doneMonitor)
 
 	var clearOnce sync.Once
 	return func() {
 		clearOnce.Do(func() {
 			close(stopMonitor)
+			<-doneMonitor
 			cc.ctx.GetSessionVars().SQLKiller.IsConnectionAlive.Store(nil)
 		})
 	}
 }
 
-func (cc *clientConn) monitorConnectionAlive(isAlive func() bool, stop <-chan struct{}) {
+func (cc *clientConn) monitorConnectionAlive(isAlive func() bool, stop <-chan struct{}, done chan<- struct{}) {
+	defer close(done)
 	checkInterval := time.Second
 	failpoint.Inject("mockConnectionAliveMonitorInterval", func(val failpoint.Value) {
 		if interval, ok := val.(int); ok {
@@ -2123,6 +2133,11 @@ func (cc *clientConn) monitorConnectionAlive(isAlive func() bool, stop <-chan st
 		select {
 		case <-ticker.C:
 			if !isAlive() {
+				select {
+				case <-stop:
+					return
+				default:
+				}
 				cc.ctx.GetSessionVars().SQLKiller.SendKillSignal(sqlkiller.QueryInterrupted)
 				cc.cancelDispatch()
 				return
@@ -2187,6 +2202,7 @@ func (cc *clientConn) handleStmt(
 	monitoringConnectionAlive := shouldMonitorConnectionAliveDuringExecute(stmt, cc.ctx.GetSessionVars())
 	if monitoringConnectionAlive {
 		clearConnectionAlive = cc.setSQLKillerConnectionAlive()
+		defer clearConnectionAlive()
 	}
 	rs, err := cc.ctx.ExecuteStmt(ctx, stmt)
 	if rs == nil || err != nil {
@@ -2224,6 +2240,7 @@ func (cc *clientConn) handleStmt(
 		}
 		if !monitoringConnectionAlive {
 			clearConnectionAlive = cc.setSQLKillerConnectionAlive()
+			defer clearConnectionAlive()
 		}
 		cc.ctx.GetSessionVars().SQLKiller.SetFinishFunc(
 			func() {
@@ -2233,7 +2250,6 @@ func (cc *clientConn) handleStmt(
 		cc.ctx.GetSessionVars().SQLKiller.InWriteResultSet.Store(true)
 		defer cc.ctx.GetSessionVars().SQLKiller.InWriteResultSet.Store(false)
 		defer cc.ctx.GetSessionVars().SQLKiller.ClearFinishFunc()
-		defer clearConnectionAlive()
 		if retryable, err := cc.writeResultSet(ctx, rs, false, status, 0); err != nil {
 			return retryable, err
 		}

pkg/server/conn_stmt.go

@@ -316,6 +316,7 @@ func (cc *clientConn) executePreparedStmtAndWriteResult(ctx context.Context, stm
 		monitoringConnectionAlive = shouldMonitorConnectionAliveDuringExecute(planCacheStmt.PreparedAst.Stmt, vars)
 		if monitoringConnectionAlive {
 			clearConnectionAlive = cc.setSQLKillerConnectionAlive()
+			defer clearConnectionAlive()
 		}
 	}
 	rs, err := (&cc.ctx).ExecuteStmt(ctx, execStmt)
@@ -326,8 +327,8 @@ func (cc *clientConn) executePreparedStmtAndWriteResult(ctx context.Context, stm
 	if rs != nil {
 		if !monitoringConnectionAlive {
 			clearConnectionAlive = cc.setSQLKillerConnectionAlive()
+			defer clearConnectionAlive()
 		}
-		defer clearConnectionAlive()
 		defer func() {
 			if !lazy {
 				rs.Close()

pkg/server/tests/commontest/tidb_test.go

@@ -3508,99 +3508,100 @@ func TestClientDisconnectKillsAutocommitInsert(t *testing.T) {
 	ts := servertestkit.CreateTidbTestSuite(t)
 	enableFastConnectionAliveMonitor(t)
 
-	ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) {
-		dbt.MustExec("drop table if exists issue57531_insert")
-		dbt.MustExec("create table issue57531_insert (a int primary key, b int)")
-
-		conn, err := dbt.GetDB().Conn(context.Background())
-		require.NoError(t, err)
-		defer func() {
-			_ = conn.Close()
-		}()
-		netConn := getRawNetConn(t, conn)
-
-		done := make(chan error, 1)
-		go func() {
-			_, err := conn.ExecContext(context.Background(), "insert into issue57531_insert values (1, sleep(300))")
-			done <- err
-		}()
-
-		require.Eventually(t, func() bool {
-			return processlistCountByInfo(t, dbt, "insert into issue57531_insert%") == 1
-		}, 5*time.Second, 50*time.Millisecond)
-
-		require.NoError(t, netConn.Close())
-
-		var execErr error
-		require.Eventually(t, func() bool {
-			select {
-			case execErr = <-done:
-				return true
-			default:
-				return false
-			}
-		}, 5*time.Second, 50*time.Millisecond)
-		require.Error(t, execErr)
-
-		require.Eventually(t, func() bool {
-			return processlistCountByInfo(t, dbt, "insert into issue57531_insert%") == 0
-		}, 5*time.Second, 50*time.Millisecond)
-
-		var cnt int
-		err = dbt.GetDB().QueryRowContext(context.Background(), "select count(*) from issue57531_insert").Scan(&cnt)
-		require.NoError(t, err)
-		require.Equal(t, 0, cnt)
-	})
+	for _, prepared := range []bool{false, true} {
+		name := "query"
+		if prepared {
+			name = "prepared"
+		}
+		t.Run(name, func(t *testing.T) {
+			ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) {
+				tableName := "issue57531_insert_" + name
+				dbt.MustExec("drop table if exists " + tableName)
+				dbt.MustExec("create table " + tableName + " (a int primary key, b int)")
+				runClientDisconnectAutocommitInsert(t, dbt, tableName, fmt.Sprintf("insert into %s values (1, sleep(300))", tableName), prepared)
+			})
+		})
+	}
 }
 
 func TestClientDisconnectCancelsAutocommitInsertPrewrite(t *testing.T) {
 	ts := servertestkit.CreateTidbTestSuite(t)
 	enableFastConnectionAliveMonitor(t)
 
-	ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) {
-		dbt.MustExec("drop table if exists issue57531_prewrite")
-		dbt.MustExec("create table issue57531_prewrite (a int primary key, b int)")
-		testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/store/mockstore/unistore/rpcPrewriteResult", `return("notLeader")`)
+	for _, prepared := range []bool{false, true} {
+		name := "query"
+		if prepared {
+			name = "prepared"
+		}
+		t.Run(name, func(t *testing.T) {
+			ts.RunTests(t, nil, func(dbt *testkit.DBTestKit) {
+				tableName := "issue57531_prewrite_" + name
+				dbt.MustExec("drop table if exists " + tableName)
+				dbt.MustExec("create table " + tableName + " (a int primary key, b int)")
+				testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/store/mockstore/unistore/rpcPrewriteResult", `return("notLeader")`)
+				runClientDisconnectAutocommitInsert(t, dbt, tableName, fmt.Sprintf("insert into %s values (1, 1)", tableName), prepared)
+			})
+		})
+	}
+}
 
-		conn, err := dbt.GetDB().Conn(context.Background())
+func runClientDisconnectAutocommitInsert(t *testing.T, dbt *testkit.DBTestKit, tableName, insertSQL string, prepared bool) {
+	conn, err := dbt.GetDB().Conn(context.Background())
+	require.NoError(t, err)
+	defer func() {
+		_ = conn.Close()
+	}()
+
+	var stmt *sql.Stmt
+	if prepared {
+		stmt, err = conn.PrepareContext(context.Background(), insertSQL)
 		require.NoError(t, err)
 		defer func() {
-			_ = conn.Close()
+			_ = stmt.Close()
 		}()
-		netConn := getRawNetConn(t, conn)
+	}
+	netConn := getRawNetConn(t, conn)
 
-		done := make(chan error, 1)
-		go func() {
-			_, err := conn.ExecContext(context.Background(), "insert into issue57531_prewrite values (1, 1)")
-			done <- err
-		}()
+	done := make(chan error, 1)
+	go func() {
+		var execErr error
+		if prepared {
+			_, execErr = stmt.ExecContext(context.Background())
+		} else {
+			_, execErr = conn.ExecContext(context.Background(), insertSQL)
+		}
+		done <- execErr
+	}()
 
-		require.Eventually(t, func() bool {
-			return processlistCountByInfo(t, dbt, "insert into issue57531_prewrite%") == 1
-		}, 5*time.Second, 50*time.Millisecond)
+	pattern := fmt.Sprintf("insert into %s%%", tableName)
+	require.Eventually(t, func() bool {
+		return processlistCountByInfo(t, dbt, pattern) == 1
+	}, 5*time.Second, 50*time.Millisecond)
+	processID, ok := processlistIDByInfo(t, dbt, pattern)
+	require.True(t, ok)
+	cleanupProcessByID(t, dbt.GetDB(), processID)
 
-		require.NoError(t, netConn.Close())
+	require.NoError(t, netConn.Close())
 
-		var execErr error
-		require.Eventually(t, func() bool {
-			select {
-			case execErr = <-done:
-				return true
-			default:
-				return false
-			}
-		}, 5*time.Second, 50*time.Millisecond)
-		require.Error(t, execErr)
+	var execErr error
+	require.Eventually(t, func() bool {
+		select {
+		case execErr = <-done:
+			return true
+		default:
+			return false
+		}
+	}, 5*time.Second, 50*time.Millisecond)
+	require.Error(t, execErr)
 
-		require.Eventually(t, func() bool {
-			return processlistCountByInfo(t, dbt, "insert into issue57531_prewrite%") == 0
-		}, 5*time.Second, 50*time.Millisecond)
+	require.Eventually(t, func() bool {
+		return processlistCountByInfo(t, dbt, pattern) == 0
+	}, 5*time.Second, 50*time.Millisecond)
 
-		var cnt int
-		err = dbt.GetDB().QueryRowContext(context.Background(), "select count(*) from issue57531_prewrite").Scan(&cnt)
-		require.NoError(t, err)
-		require.Equal(t, 0, cnt)
-	})
+	var cnt int
+	err = dbt.GetDB().QueryRowContext(context.Background(), "select count(*) from "+tableName).Scan(&cnt)
+	require.NoError(t, err)
+	require.Equal(t, 0, cnt)
 }
 
 func enableFastConnectionAliveMonitor(t *testing.T) {
@@ -3642,6 +3643,46 @@ func processlistCountByInfo(t *testing.T, dbt *testkit.DBTestKit, pattern string
 	return cnt
 }
 
+func processlistIDByInfo(t *testing.T, dbt *testkit.DBTestKit, pattern string) (uint64, bool) {
+	var id uint64
+	err := dbt.GetDB().QueryRowContext(
+		context.Background(),
+		"select id from information_schema.processlist where info like ? limit 1",
+		pattern,
+	).Scan(&id)
+	if err == sql.ErrNoRows {
+		return 0, false
+	}
+	require.NoError(t, err)
+	return id, true
+}
+
+func cleanupProcessByID(t *testing.T, db *sql.DB, processID uint64) {
+	t.Cleanup(func() {
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+
+		conn, err := db.Conn(ctx)
+		if err != nil {
+			return
+		}
+		defer func() {
+			_ = conn.Close()
+		}()
+
+		var cnt int
+		err = conn.QueryRowContext(
+			ctx,
+			"select count(*) from information_schema.processlist where id = ?",
+			processID,
+		).Scan(&cnt)
+		if err != nil || cnt == 0 {
+			return
+		}
+		_, _ = conn.ExecContext(ctx, fmt.Sprintf("kill query %d", processID))
+	})
+}
+
 func TestCloseConnForUndeterminedError(t *testing.T) {
 	cfg := util2.NewTestConfig()
 	cfg.Host = "127.0.0.1" // No network interface listening for mysql traffic

pkg/util/sqlkiller/sqlkiller.go

@@ -220,7 +220,7 @@ func (killer *SQLKiller) HandleSignal() error {
 		} else if now.Sub(*lastCheckTime) > checkConnectionAliveDur {
 			killer.lastCheckTime.Store(&now)
 			if !(*fn)() {
-				atomic.CompareAndSwapUint32(&killer.Signal, 0, QueryInterrupted)
+				killer.sendKillSignal(QueryInterrupted)
 			}
 		}
 	}
@@ -238,7 +238,7 @@ func (killer *SQLKiller) HandleSignal() error {
 func (killer *SQLKiller) CheckConnectionAlive() {
 	fn := killer.IsConnectionAlive.Load()
 	if fn != nil && !(*fn)() {
-		atomic.CompareAndSwapUint32(&killer.Signal, 0, QueryInterrupted)
+		killer.sendKillSignal(QueryInterrupted)
 	}
 }