diff --git a/executor/adapter.go b/executor/adapter.go index 1caf36fddd773..430dae026582b 100644 --- a/executor/adapter.go +++ b/executor/adapter.go @@ -733,7 +733,7 @@ func (a *ExecStmt) handlePessimisticDML(ctx context.Context, e Executor) error { if err1 != nil { return err1 } - keys = txnCtx.CollectUnchangedRowKeys(keys) + keys = txnCtx.CollectUnchangedLockKeys(keys) if len(keys) == 0 { return nil } diff --git a/executor/write.go b/executor/write.go index 2c9c8fe4107fa..397ec2669a46d 100644 --- a/executor/write.go +++ b/executor/write.go @@ -158,7 +158,21 @@ func updateRecord(ctx context.Context, sctx sessionctx.Context, h kv.Handle, old unchangedRowKey := tablecodec.EncodeRowKeyWithHandle(physicalID, h) txnCtx := sctx.GetSessionVars().TxnCtx if txnCtx.IsPessimistic { - txnCtx.AddUnchangedRowKey(unchangedRowKey) + txnCtx.AddUnchangedLockKey(unchangedRowKey) + for _, idx := range t.Indices() { + if !idx.Meta().Unique { + continue + } + ukVals, err := idx.FetchValues(oldData, nil) + if err != nil { + return false, err + } + unchangedUniqueKey, _, err := idx.GenIndexKey(sc, ukVals, h, nil) + if err != nil { + return false, err + } + txnCtx.AddUnchangedLockKey(unchangedUniqueKey) + } } return false, nil } diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go index ddf249281858a..05e2c059fb45f 100644 --- a/sessionctx/variable/session.go +++ b/sessionctx/variable/session.go @@ -169,8 +169,9 @@ type TxnCtxNoNeedToRestore struct { currentShard int64 shardRand *rand.Rand - // unchangedRowKeys is used to store the unchanged rows that needs to lock for pessimistic transaction. - unchangedRowKeys map[string]struct{} + // unchangedLockKeys is used to store the unchanged keys that needs to lock for pessimistic transaction, including + // row keys and unique keys. + unchangedLockKeys map[string]struct{} PessimisticCacheHit int @@ -239,20 +240,20 @@ func (tc *TransactionContext) updateShard() { tc.currentShard = int64(murmur3.Sum32(buf[:])) } -// AddUnchangedRowKey adds an unchanged row key in update statement for pessimistic lock. -func (tc *TransactionContext) AddUnchangedRowKey(key []byte) { - if tc.unchangedRowKeys == nil { - tc.unchangedRowKeys = map[string]struct{}{} +// AddUnchangedLockKey adds an unchanged key in update statement for pessimistic lock. +func (tc *TransactionContext) AddUnchangedLockKey(key []byte) { + if tc.unchangedLockKeys == nil { + tc.unchangedLockKeys = map[string]struct{}{} } - tc.unchangedRowKeys[string(key)] = struct{}{} + tc.unchangedLockKeys[string(key)] = struct{}{} } -// CollectUnchangedRowKeys collects unchanged row keys for pessimistic lock. -func (tc *TransactionContext) CollectUnchangedRowKeys(buf []kv.Key) []kv.Key { - for key := range tc.unchangedRowKeys { +// CollectUnchangedLockKeys collects unchanged keys for pessimistic lock. +func (tc *TransactionContext) CollectUnchangedLockKeys(buf []kv.Key) []kv.Key { + for key := range tc.unchangedLockKeys { buf = append(buf, kv.Key(key)) } - tc.unchangedRowKeys = nil + tc.unchangedLockKeys = nil return buf } diff --git a/tests/realtikvtest/pessimistictest/pessimistic_test.go b/tests/realtikvtest/pessimistictest/pessimistic_test.go index 58bb2b73a86f5..1759fff3e3bd4 100644 --- a/tests/realtikvtest/pessimistictest/pessimistic_test.go +++ b/tests/realtikvtest/pessimistictest/pessimistic_test.go @@ -503,6 +503,38 @@ func TestLockUnchangedRowKey(t *testing.T) { tk2.MustExec("rollback") } +func TestLockUnchangedUniqueKey(t *testing.T) { + store, clean := realtikvtest.CreateMockStoreAndSetup(t) + defer clean() + + tk := testkit.NewTestKit(t, store) + tk2 := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk2.MustExec("use test") + + // ref https://github.com/pingcap/tidb/issues/36438 + tk.MustExec("drop table if exists t") + tk.MustExec("create table t (i varchar(10), unique key(i))") + tk.MustExec("insert into t values ('a')") + tk.MustExec("begin pessimistic") + tk.MustExec("update t set i = 'a'") + + errCh := make(chan error, 1) + go func() { + _, err := tk2.Exec("insert into t values ('a')") + errCh <- err + }() + + select { + case <-errCh: + require.Fail(t, "insert is not blocked by update") + case <-time.After(500 * time.Millisecond): + tk.MustExec("rollback") + } + + require.Error(t, <-errCh) +} + func TestOptimisticConflicts(t *testing.T) { store, clean := realtikvtest.CreateMockStoreAndSetup(t) defer clean()