transaction: fix union select for update race (#19006)

* transaction: fix LockKeys race

* do not update delta for lock keys

* fix more race

* fix another race

Co-authored-by: ti-srebot <66930949+ti-srebot@users.noreply.github.com>

executor/executor.go

@@ -942,13 +942,6 @@ func (e *SelectLockExec) Next(ctx context.Context, req *chunk.Chunk) error {
 		lockWaitTime = kv.LockNoWait
 	}
 
-	if len(e.keys) > 0 {
-		// This operation is only for schema validator check.
-		for id := range e.tblID2Handle {
-			e.ctx.GetSessionVars().TxnCtx.UpdateDeltaForTable(id, 0, 0, map[int64]int64{})
-		}
-	}
-
 	return doLockKeys(ctx, e.ctx, newLockCtx(e.ctx.GetSessionVars(), lockWaitTime), e.keys...)
 }
 
@@ -972,7 +965,7 @@ func newLockCtx(seVars *variable.SessionVars, lockWaitTime int64) *kv.LockCtx {
 func doLockKeys(ctx context.Context, se sessionctx.Context, lockCtx *kv.LockCtx, keys ...kv.Key) error {
 	sctx := se.GetSessionVars().StmtCtx
 	if !sctx.InUpdateStmt && !sctx.InDeleteStmt {
-		se.GetSessionVars().TxnCtx.ForUpdate = true
+		atomic.StoreUint32(&se.GetSessionVars().TxnCtx.ForUpdate, 1)
 	}
 	// Lock keys only once when finished fetching all results.
 	txn, err := se.Txn(true)

kv/kv.go

@@ -250,7 +250,7 @@ type LockCtx struct {
 	LockWaitTime          int64
 	WaitStartTime         time.Time
 	PessimisticLockWaited *int32
-	LockKeysDuration      *time.Duration
+	LockKeysDuration      *int64
 	LockKeysCount         *int32
 	ReturnValues          bool
 	Values                map[string]ReturnedValue

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")
+}

session/session.go

@@ -627,7 +627,7 @@ func (s *session) retry(ctx context.Context, maxCnt uint) (err error) {
 
 	connID := s.sessionVars.ConnectionID
 	s.sessionVars.RetryInfo.Retrying = true
-	if s.sessionVars.TxnCtx.ForUpdate {
+	if atomic.LoadUint32(&s.sessionVars.TxnCtx.ForUpdate) == 1 {
 		err = ErrForUpdateCantRetry.GenWithStackByArgs(connID)
 		return err
 	}

sessionctx/stmtctx/stmtctx.go

@@ -146,10 +146,10 @@ type StatementContext struct {
 	planNormalized        string
 	planDigest            string
 	Tables                []TableEntry
-	PointExec             bool       // for point update cached execution, Constant expression need to set "paramMarker"
-	lockWaitStartTime     *time.Time // LockWaitStartTime stores the pessimistic lock wait start time
+	PointExec             bool  // for point update cached execution, Constant expression need to set "paramMarker"
+	lockWaitStartTime     int64 // LockWaitStartTime stores the pessimistic lock wait start time
 	PessimisticLockWaited int32
-	LockKeysDuration      time.Duration
+	LockKeysDuration      int64
 	LockKeysCount         int32
 	TblInfo2UnionScan     map[*model.TableInfo]bool
 	TaskID                uint64 // unique ID for an execution of a statement
@@ -510,7 +510,7 @@ func (sc *StatementContext) GetExecDetails() execdetails.ExecDetails {
 	var details execdetails.ExecDetails
 	sc.mu.Lock()
 	details = sc.mu.execDetails
-	details.LockKeysDuration = sc.LockKeysDuration
+	details.LockKeysDuration = time.Duration(atomic.LoadInt64(&sc.LockKeysDuration))
 	sc.mu.Unlock()
 	return details
 }
@@ -651,11 +651,12 @@ func (sc *StatementContext) SetFlagsFromPBFlag(flags uint64) {
 
 // GetLockWaitStartTime returns the statement pessimistic lock wait start time
 func (sc *StatementContext) GetLockWaitStartTime() time.Time {
-	if sc.lockWaitStartTime == nil {
-		curTime := time.Now()
-		sc.lockWaitStartTime = &curTime
+	startTime := atomic.LoadInt64(&sc.lockWaitStartTime)
+	if startTime == 0 {
+		startTime = time.Now().UnixNano()
+		atomic.StoreInt64(&sc.lockWaitStartTime, startTime)
 	}
-	return *sc.lockWaitStartTime
+	return time.Unix(0, startTime)
 }
 
 //CopTasksDetails collects some useful information of cop-tasks during execution.

sessionctx/variable/session.go

@@ -163,11 +163,11 @@ type TransactionContext struct {
 	// CreateTime For metrics.
 	CreateTime     time.Time
 	StatementCount int
-	ForUpdate      bool
 	CouldRetry     bool
 	IsPessimistic  bool
 	Isolation      string
 	LockExpire     uint32
+	ForUpdate      uint32
 }
 
 // GetShard returns the shard prefix for the next `count` rowids.

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))
@@ -338,7 +340,7 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, lockCtx *kv.LockCtx, keysInput
 			if lockCtx.PessimisticLockWaited != nil {
 				if atomic.LoadInt32(lockCtx.PessimisticLockWaited) > 0 {
 					timeWaited := time.Since(lockCtx.WaitStartTime)
-					*lockCtx.LockKeysDuration = timeWaited
+					atomic.StoreInt64(lockCtx.LockKeysDuration, int64(timeWaited))
 					metrics.TiKVPessimisticLockKeysDuration.Observe(timeWaited.Seconds())
 				}
 			}
@@ -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
 }