unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
pingcap/tidb
mid9/10
concurrencyerror-handlingtesting-strategy

util: prevent SQLKiller state races during concurrent reset

Project: pingcap/tidb | Language: Go | Concepts: concurrency, error-handling, testing-strategy | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70557

The problem

The SQLKiller component in TiDB could enter an inconsistent state when a kill signal is sent at the same time as a reset. The kill signal and related event state (channel, description, triggered flag) were updated independently, so a concurrent Reset could clear the signal but leave the event triggered, or vice versa. This led to situations where the system would panic (crash in the statement layer) or behave unpredictably, such as logging a nil error or dropping client connections unexpectedly.

Why it's non-trivial

A naive fix might just add more atomic operations or simple locks around individual fields, but that would not guarantee that all related state transitions (signal, event, description) are observed in a consistent order. Without careful synchronization, subtle races could still leave the system in an invalid state, especially under high concurrency or in real-world failover scenarios. Logging and error reporting also need to see a consistent snapshot of state, which is tricky if fields can change between reads.

The approach

The PR introduces a mutex to serialize all updates to the kill signal and its associated event state, and ensures that logging and error construction see a consistent snapshot. It also adds deterministic, failpoint-driven unit tests to verify the fix.

1. Add a mutex to serialize writers to signal and event state:

 type SQLKiller struct {
 	killEvent struct {
 		ch   chan struct{}
 		desc string
+		// Mutex serializes Signal writers with ch, desc, and triggered.
 		sync.Mutex
 		triggered bool
 	}

This change adds a sync.Mutex to the killEvent struct, ensuring that any code updating the kill signal or related event state must hold this lock.

2. Combine signal installation and event triggering under one lock:

 func (killer *SQLKiller) SendKillSignalWithKillEventReason(killSignal killSignal, desc string) {
-	{
-		killer.killEvent.Lock()
-		killer.killEvent.desc = desc
-		killer.killEvent.Unlock()
-	}
-	killer.sendKillSignal(killSignal)
-	killer.triggerKillEvent()
+	killer.killEvent.Lock()
+	killer.killEvent.desc = desc
+	signalSent, eventDesc := killer.sendKillSignalLocked(killSignal)
+	killer.triggerKillEventLocked()
+	killer.killEvent.Unlock()
+
+	if signalSent {
+		killer.logKillSignal(killSignal, eventDesc)
+	}
 }

Now, the description, signal, and event are all set within a single critical section. Logging is done after the lock is released, using the captured event description.

3. Make logging and error construction use a consistent snapshot:

-func (killer *SQLKiller) sendKillSignal(reason killSignal) {
+func (killer *SQLKiller) sendKillSignal(reason killSignal) {
+	killer.killEvent.Lock()
+	signalSent, eventDesc := killer.sendKillSignalLocked(reason)
+	killer.killEvent.Unlock()
+
+	if signalSent {
+		killer.logKillSignal(reason, eventDesc)
+	}
+}

This ensures that the log message is constructed only after the signal and event description are set, and that both are read together while holding the lock.

4. Update error construction to take the description as an argument:

-func (killer *SQLKiller) getKillError(status killSignal) error {
+func (killer *SQLKiller) getKillError(status killSignal, desc string) error {
     ...
-	case KilledByMemArbitrator:
-		return exeerrors.ErrQueryExecStopped.GenWithStackByArgs(killer.getKillEventReason(), killer.ConnID.Load())
+	case KilledByMemArbitrator:
+		return exeerrors.ErrQueryExecStopped.GenWithStackByArgs(desc, killer.ConnID.Load())

This avoids reading a possibly stale or inconsistent description after the lock is released.

5. Serialize Reset() and add a failpoint for deterministic testing:

-func (killer *SQLKiller) Reset() {
-	if atomic.LoadUint32(&killer.Signal) != 0 {
+	killer.killEvent.Lock()
+	status := atomic.SwapUint32(&killer.Signal, UnspecifiedKillSignal)
+	// Keep this hook immediately after the atomic clear. If Reset is ever split
+	// into separate observation and clear operations, place it between them.
+	failpoint.InjectCall("afterResetKillSignalSwap")
+	killer.resetKillEventLocked()
+	killer.killEvent.Unlock()
+
+	if status != UnspecifiedKillSignal {
 		logutil.BgLogger().Warn("kill finished", zap.Uint64("conn", killer.ConnID.Load()))
 	}
-
-	atomic.StoreUint32(&killer.Signal, 0)
-	killer.resetKillEvent()
 	killer.lastCheckTime.Store(nil)
 }

The entire reset sequence is now protected by the lock, and a failpoint is injected for test control.

6. Add deterministic, failpoint-driven concurrency tests:

func TestSQLKillerConcurrentReset(t *testing.T) {
    ...
    t.Run("reset after successful kill signal CAS", func(t *testing.T) {
        ...
        testfailpoint.EnableCall(t, beforeLogFailpoint, func() {
            killer.Reset()
        })
        ...
        require.NotPanics(t, func() {
            killer.SendKillSignal(QueryInterrupted)
        })
        ...
    })

    t.Run("kill signal after reset clear", func(t *testing.T) {
        ...
        testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/util/sqlkiller/afterResetKillSignalSwap", func() {
            assertStateLockHeld(t, killer)
            go func() {
                defer close(killSent)
                killer.SendKillSignalWithKillEventReason(KilledByMemArbitrator, reason)
            }()
        })

        killer.Reset()
        <-killSent
        ...
    })
}

These tests use failpoints to deterministically interleave Reset and kill-signal operations, and then assert that the state is always consistent and no panics occur.

From the review thread

No notable review discussion on this PR.

Takeaways

  • Synchronize all related state transitions under a single lock to prevent race conditions in concurrent environments.
  • Capture all necessary state for logging and error reporting inside the critical section, and perform side effects (like logging) after releasing the lock.
  • Use deterministic, failpoint-driven tests to reliably reproduce and verify concurrency bugs and their fixes.