From a6335a5aaf934be56a32aecac66c8b36021fdf67 Mon Sep 17 00:00:00 2001 From: MyonKeminta <9948422+MyonKeminta@users.noreply.github.com> Date: Thu, 30 May 2024 16:20:09 +0800 Subject: [PATCH] Fix the problem that statement being oom-killed within DoneAggressiveLocking causing the transaction still in aggressive locking state (#1355) * Quick fix the problem that statement being oom-killed within DoneAggressiveLocking causing the transaction still in aggressive locking state Signed-off-by: MyonKeminta * Add comments to explain the change Signed-off-by: MyonKeminta --------- Signed-off-by: MyonKeminta Co-authored-by: MyonKeminta --- txnkv/transaction/txn.go | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 0ee2b5f2..775b17be 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -957,6 +957,18 @@ func (txn *KVTxn) CancelAggressiveLocking(ctx context.Context) { if txn.aggressiveLockingContext == nil { panic("Trying to cancel aggressive locking while it's not started") } + + // Unset `aggressiveLockingContext` in a defer block to ensure it must be executed even it panicked on the half way. + // It's because that if it's panicked by an OOM-kill of TiDB, it can then be recovered and the user can still + // continue using the transaction's state. + // The usage of `defer` can be removed once we have other way to avoid the panicking. + // See: https://github.com/pingcap/tidb/issues/53540#issuecomment-2138089140 + // Currently the problem only exists in `DoneAggressiveLocking`, but we do the same to `CancelAggressiveLocking` + // to the two function consistent, and prevent for new panics that might be introduced in the future. + defer func() { + txn.aggressiveLockingContext = nil + }() + txn.cleanupAggressiveLockingRedundantLocks(context.Background()) if txn.aggressiveLockingContext.assignedPrimaryKey { txn.resetPrimary() @@ -976,7 +988,6 @@ func (txn *KVTxn) CancelAggressiveLocking(ctx context.Context) { txn.asyncPessimisticRollback(context.Background(), keys, forUpdateTS) txn.lockedCnt -= len(keys) } - txn.aggressiveLockingContext = nil } // DoneAggressiveLocking finishes the current aggressive locking. The locked keys will be moved to the membuffer as if @@ -985,6 +996,16 @@ func (txn *KVTxn) DoneAggressiveLocking(ctx context.Context) { if txn.aggressiveLockingContext == nil { panic("Trying to finish aggressive locking while it's not started") } + + // Unset `aggressiveLockingContext` in a defer block to ensure it must be executed even it panicked on the half way. + // It's because that if it's panicked by an OOM-kill of TiDB, it can then be recovered and the user can still + // continue using the transaction's state. + // The usage of `defer` can be removed once we have other way to avoid the panicking. + // See: https://github.com/pingcap/tidb/issues/53540#issuecomment-2138089140 + defer func() { + txn.aggressiveLockingContext = nil + }() + txn.cleanupAggressiveLockingRedundantLocks(context.Background()) if txn.forUpdateTSChecks == nil { @@ -1009,7 +1030,6 @@ func (txn *KVTxn) DoneAggressiveLocking(ctx context.Context) { txn.committer.maxLockedWithConflictTS = txn.aggressiveLockingContext.maxLockedWithConflictTS } } - txn.aggressiveLockingContext = nil } // IsInAggressiveLockingMode checks if the transaction is currently in aggressive locking mode.