util: fix data race in sqlkiller

pingcap · Aug 22, 2024 · 7fe2fd2 · 7fe2fd2
1 parent 2761e91
commit 7fe2fd2
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/pkg/executor/executor.go b/pkg/executor/executor.go
@@ -1856,7 +1856,7 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	vars.MemTracker.Killer = &vars.SQLKiller
 	vars.DiskTracker.Killer = &vars.SQLKiller
 	vars.SQLKiller.Reset()
-	vars.SQLKiller.ConnID = vars.ConnectionID
+	vars.SQLKiller.ConnID.Store(vars.ConnectionID)
 
 	isAnalyze := false
 	if execStmt, ok := s.(*ast.ExecuteStmt); ok {

diff --git a/pkg/util/sqlkiller/sqlkiller.go b/pkg/util/sqlkiller/sqlkiller.go
@@ -42,7 +42,7 @@ const (
 // SQLKiller is used to kill a query.
 type SQLKiller struct {
 	Signal killSignal
-	ConnID uint64
+	ConnID atomic.Uint64
 	// FinishFuncLock is used to ensure that Finish is not called and modified at the same time.
 	// An external call to the Finish function only allows when the main goroutine to be in the writeResultSet process.
 	// When the main goroutine exits the writeResultSet process, the Finish function will be cleared.
@@ -58,7 +58,7 @@ func (killer *SQLKiller) SendKillSignal(reason killSignal) {
 	if atomic.CompareAndSwapUint32(&killer.Signal, 0, reason) {
 		status := atomic.LoadUint32(&killer.Signal)
 		err := killer.getKillError(status)
-		logutil.BgLogger().Warn("kill initiated", zap.Uint64("connection ID", killer.ConnID), zap.String("reason", err.Error()))
+		logutil.BgLogger().Warn("kill initiated", zap.Uint64("connection ID", killer.ConnID.Load()), zap.String("reason", err.Error()))
 	}
 }
 
@@ -115,7 +115,7 @@ func (killer *SQLKiller) HandleSignal() error {
 	failpoint.Inject("randomPanic", func(val failpoint.Value) {
 		if p, ok := val.(int); ok {
 			if rand.Float64() > (float64)(p)/1000 {
-				if killer.ConnID != 0 {
+				if killer.ConnID.Load() != 0 {
 					targetStatus := rand.Int31n(5)
 					atomic.StoreUint32(&killer.Signal, uint32(targetStatus))
 				}
@@ -126,15 +126,15 @@ func (killer *SQLKiller) HandleSignal() error {
 	err := killer.getKillError(status)
 	if status == ServerMemoryExceeded {
 		logutil.BgLogger().Warn("global memory controller, NeedKill signal is received successfully",
-			zap.Uint64("conn", killer.ConnID))
+			zap.Uint64("conn", killer.ConnID.Load()))
 	}
 	return err
 }
 
 // Reset resets the SqlKiller.
 func (killer *SQLKiller) Reset() {
 	if atomic.LoadUint32(&killer.Signal) != 0 {
-		logutil.BgLogger().Warn("kill finished", zap.Uint64("conn", killer.ConnID))
+		logutil.BgLogger().Warn("kill finished", zap.Uint64("conn", killer.ConnID.Load()))
 	}
 	atomic.StoreUint32(&killer.Signal, 0)
 }