transaction: fix union select for update race

session/pessimistic_test.go

@@ -1486,3 +1486,15 @@ func (s *testPessimisticSuite) TestPessimisticTxnWithDDLChangeColumn(c *C) {
 
 	tk2.MustExec("drop database if exists test_db")
 }
+
+func (s *testPessimisticSuite) TestPessimisticUnionForUpdate(c *C) {
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(id int, v int, k int, primary key (id), key kk(k))")
+	tk.MustExec("insert into t select 1, 1, 1")
+	tk.MustExec("begin pessimistic")
+	tk.MustQuery("(select * from t where id between 0 and 1 for update) union all (select * from t where id between 0 and 1 for update)")
+	tk.MustExec("update t set k = 2 where k = 1")
+	tk.MustExec("commit")
+	tk.MustExec("admin check table t")
+}

store/tikv/txn.go

@@ -330,6 +330,8 @@ func (txn *tikvTxn) rollbackPessimisticLocks() error {
 
 // lockWaitTime in ms, except that kv.LockAlwaysWait(0) means always wait lock, kv.LockNowait(-1) means nowait lock
 func (txn *tikvTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keysInput ...kv.Key) error {
+	txn.mu.Lock()
+	defer txn.mu.Unlock()
 	// Exclude keys that are already locked.
 	var err error
 	keys := make([][]byte, 0, len(keysInput))
@@ -347,7 +349,6 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keysInput
 			*lockCtx.LockKeysCount += int32(len(keys))
 		}
 	}()
-	txn.mu.Lock()
 	for _, key := range keysInput {
 		if _, ok := txn.lockedMap[string(key)]; !ok {
 			keys = append(keys, key)
@@ -357,7 +358,6 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keysInput
 			lockCtx.Values[string(key)] = kv.ReturnedValue{AlreadyLocked: true}
 		}
 	}
-	txn.mu.Unlock()
 	if len(keys) == 0 {
 		return nil
 	}
@@ -423,13 +423,11 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keysInput
 			txn.committer.ttlManager.run(txn.committer, lockCtx)
 		}
 	}
-	txn.mu.Lock()
 	txn.lockKeys = append(txn.lockKeys, keys...)
 	for _, key := range keys {
 		txn.lockedMap[string(key)] = struct{}{}
 	}
 	txn.dirty = true
-	txn.mu.Unlock()
 	return nil
 }