Fix the problem that statement being oom-killed within DoneAggressive…

…Locking causing the transaction still in aggressive locking state (tikv#1355) (tikv#1487)

 

Co-authored-by: MyonKeminta <[email protected]>

txnkv/transaction/txn.go

@@ -735,6 +735,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()
@@ -754,7 +766,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
@@ -763,6 +774,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 {
@@ -787,7 +808,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.