Repository: pingcap/tidb (Go)

Issue:
Unique key lock not consistent when row is changed or not
## Bug Report

Please answer these questions before submitting your issue. Thanks!

### 1. Minimal reproduce step (Required)
```
mysql> show create table t;
+-------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Table | Create Table                                                                                                                                                         |
+-------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| t     | CREATE TABLE `t` (
  `i` varchar(20) NOT NULL,
  PRIMARY KEY (`i`) /*T![clustered_index] NONCLUSTERED */
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin |
+-------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------+
1 row in set (0.00 sec)

mysql> insert into t values ('a');
Query OK, 0 rows affected (0.00 sec)
Rows matched: 1  Changed: 0  Warnings: 0

mysql> begin;
Query OK, 0 rows affected (0.00 sec)

mysql> select * from t;
+---+
| i |
+---+
| a |
+---+
1 row in set (0.00 sec)

mysql> update t set i = 'a'; -- insert into t values ('a') in another session
Query OK, 0 rows affected (0.00 sec)
Rows matched: 1  Changed: 0  Warnings: 0

mysql> update t set i = 'b'; -- insert into t values ('b') in another session
Query OK, 1 row affected (0.00 sec)
Rows matched: 1  Changed: 1  Warnings: 0

mysql> rollback;
Query OK, 0 rows affected (0.00 sec)
```

<!-- a step by step guide for reproducing the bug. -->

### 2. What did you expect to see? (Required)
The first insert in another session should wait lock.

### 3. What did you see instead (Required)
ERROR 1062 (23000): Duplicate entry 'a' for key 'PRIMARY'

### 4. What is your TiDB version? (Required)

01d02568c25946684b6c4cc4cc5343f6d2068a99

<!-- Paste the output of SELECT tidb_version() -->




Discussion:
It is caused by https://github.com/pingcap/tidb/blob/f0717dfe07ae4842c7731b23368e664801d2bd42/sessionctx/variable/session.go#L251, it only collect row, so the unique index is missing to lock.
After this for such execution plans:

```
update uk = uk where sk = 1; // The uk value is not changed
The plan:
update
    select lock
        index look up 
```

The unique key could be locked but https://github.com/pingcap/tidb/pull/25730 would not work so the LOCK record could accumulate and lead to performance issues acquiring the pessimistic locks.
@zyguan 
We may need to find out another solution to avoid the performance issues, maybe it's needed to revert the fix first before that.
> After this for such execution plans:
> 
> ```
> update uk = uk where sk = 1; // The uk value is not changed
> The plan:
> update
>     select lock
>         index look up 
> ```
> 
> The unique key could be locked but #25730 would not work so the LOCK record could accumulate and lead to performance issues acquiring the pessimistic locks.

#25730 doesn't work because it's not a point plan. A LOCK record will be left on row key as well (not only unique key) if nothing updated. But it's usually not a problem since row values are more likely to be changed. Maybe we can convert lock records to put records as #25730 did. However, it seems more difficult to construct uk values here.

![2022-08-09_212546](https://user-images.githubusercontent.com/6850317/183658347-e6ba79a0-04a8-48e3-9392-b1d8bffa2748.png)

I also notice that #36498 doesn't work for the following case...
```sql
/* s0 */ drop table if exists t;
-- s0 >> 0 rows affected
/* s0 */ create table t (k varchar(10), v int, primary key (k), key idx(k));
-- s0 >> 0 rows affected
/* s0 */ insert into t values ('a', 10);
-- s0 >> 1 rows affected
/* s1 */ begin pessimistic;
-- s1 >> 0 rows affected
/* s1 */ update t force index(idx) set v = 11 where k = 'a';
-- s1 >> 1 rows affected
/* s2 */ insert into t values ('a', 100);
-- s2 >> E1062: Duplicate entry 'a' for key 'PRIMARY'
/* s1 */ commit;
-- s1 >> 0 rows affected
```
There could be some cases the accumulating LOCK records may introduce performance issues, actually it's found when we're doing some simulating benchmark tests(/cc @MyonKeminta ).
Compared with the not-locked issue, the performance issue may introduce much more production risks, I prefer not to add more locks before we solve the accumulating problem completely. Also the not lock issue exist in all the previous versions and it seems to be not a big problem.
> There could be some cases the accumulating LOCK records may introduce performance issues, actually it's found when we're doing some simulating benchmark tests(/cc @MyonKeminta ).
> Compared with the not-locked issue, the performance issue may introduce much more production risks, I prefer not to add more locks before we solve the accumulating problem completely. Also the not lock issue exist in all the previous versions and it seems to be not a big problem.

Agreed.
It seems an ultimate solution to the lock accumulation problem is necessary for us. According to our previous research, it will take us a lot of efforts, but considering our future, I think it worth.



Patch:
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
 }
 


Test patch:
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()