Add exitAggressiveLockingIfInapplicable function and call within mutex of lockKeys (#749)

* Add exitAggressiveLockingIfInapplicable function and call within mutex of lockKeys

Signed-off-by: MyonKeminta <MyonKeminta@users.noreply.github.com>

* fix fmt

Signed-off-by: MyonKeminta <MyonKeminta@users.noreply.github.com>

* Add test

Signed-off-by: MyonKeminta <MyonKeminta@users.noreply.github.com>

---------

Signed-off-by: MyonKeminta <MyonKeminta@users.noreply.github.com>
Co-authored-by: MyonKeminta <MyonKeminta@users.noreply.github.com>

integration_tests/2pc_test.go

@@ -1496,6 +1496,30 @@ func (s *testCommitterSuite) TestAggressiveLockingLoadValueOptionChanges() {
 	}
 }
 
+func (s *testCommitterSuite) TestAggressiveLockingExitIfInapplicable() {
+	txn := s.begin()
+	txn.SetPessimistic(true)
+	txn.StartAggressiveLocking()
+	s.True(txn.IsInAggressiveLockingMode())
+	lockCtx := &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
+	s.NoError(txn.LockKeys(context.Background(), lockCtx, []byte("k1")))
+
+	txn.RetryAggressiveLocking(context.Background())
+	lockCtx = &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
+	s.NoError(txn.LockKeys(context.Background(), lockCtx, []byte("k2")))
+	s.True(txn.IsInAggressiveLockingMode())
+
+	lockCtx = &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
+	s.NoError(txn.LockKeys(context.Background(), lockCtx, []byte("k3"), []byte("k4")))
+	s.False(txn.IsInAggressiveLockingMode())
+	s.checkIsKeyLocked([]byte("k1"), false)
+	s.checkIsKeyLocked([]byte("k2"), true)
+	s.checkIsKeyLocked([]byte("k3"), true)
+	s.checkIsKeyLocked([]byte("k4"), true)
+
+	s.NoError(txn.Rollback())
+}
+
 // TestElapsedTTL tests that elapsed time is correct even if ts physical time is greater than local time.
 func (s *testCommitterSuite) TestElapsedTTL() {
 	key := []byte("key")

txnkv/transaction/txn.go

@@ -861,6 +861,19 @@ func (txn *KVTxn) collectAggressiveLockingStats(lockCtx *tikv.LockCtx, keys int,
 	lockCtx.Stats.AggressiveLockDerivedCount += filteredAggressiveLockedKeysCount
 }
 
+func (txn *KVTxn) exitAggressiveLockingIfInapplicable(ctx context.Context, keys [][]byte) error {
+	if len(keys) > 1 && txn.IsInAggressiveLockingMode() {
+		// Only allow fair locking if it only needs to lock one key. Considering that it's possible that a
+		// statement causes multiple calls to `LockKeys` (which means some keys may have been locked in fair
+		// locking mode), here we exit fair locking mode by calling DoneFairLocking instead of cancelling.
+		// Then the previously-locked keys during execution in this statement (if any) will be turned into the state
+		// as if they were locked in normal way.
+		// Note that the issue https://github.com/pingcap/tidb/issues/35682 also exists here.
+		txn.DoneAggressiveLocking(ctx)
+	}
+	return nil
+}
+
 // LockKeys tries to lock the entries with the keys in KV store.
 // lockCtx is the context for lock, lockCtx.lockWaitTime in ms
 func (txn *KVTxn) LockKeys(ctx context.Context, lockCtx *tikv.LockCtx, keysInput ...[]byte) error {
@@ -889,6 +902,12 @@ func (txn *KVTxn) lockKeys(ctx context.Context, lockCtx *tikv.LockCtx, fn func()
 	startTime := time.Now()
 	txn.mu.Lock()
 	defer txn.mu.Unlock()
+
+	err = txn.exitAggressiveLockingIfInapplicable(ctx, keysInput)
+	if err != nil {
+		return err
+	}
+
 	defer func() {
 		if txn.isInternal() {
 			metrics.TxnCmdHistogramWithLockKeysInternal.Observe(time.Since(startTime).Seconds())