From 5ec654e27888281158655d9b8ae0d7ea6372e78c Mon Sep 17 00:00:00 2001
From: ekexium <eke@fastmail.com>
Date: Tue, 26 Jul 2022 13:15:10 +0800
Subject: [PATCH] executor: do not acqurie pessimistic lock for non-unique
 index keys (#36229)

close pingcap/tidb#36235
---
 session/tidb_test.go     | 25 ++++++++++++++-----------
 session/txn.go           |  9 ++++++---
 tablecodec/tablecodec.go | 13 +++++++++++++
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/session/tidb_test.go b/session/tidb_test.go
index e2808080b7c1c..f2c07acbf64b3 100644
--- a/session/tidb_test.go
+++ b/session/tidb_test.go
@@ -76,7 +76,8 @@ func TestParseErrorWarn(t *testing.T) {
 
 func TestKeysNeedLock(t *testing.T) {
 	rowKey := tablecodec.EncodeRowKeyWithHandle(1, kv.IntHandle(1))
-	indexKey := tablecodec.EncodeIndexSeekKey(1, 1, []byte{1})
+	uniqueIndexKey := tablecodec.EncodeIndexSeekKey(1, 1, []byte{1})
+	nonUniqueIndexKey := tablecodec.EncodeIndexSeekKey(1, 2, []byte{1})
 	uniqueValue := make([]byte, 8)
 	uniqueUntouched := append(uniqueValue, '1')
 	nonUniqueVal := []byte{'0'}
@@ -90,18 +91,20 @@ func TestKeysNeedLock(t *testing.T) {
 	}{
 		{rowKey, rowVal, true},
 		{rowKey, deleteVal, true},
-		{indexKey, nonUniqueVal, false},
-		{indexKey, nonUniqueUntouched, false},
-		{indexKey, uniqueValue, true},
-		{indexKey, uniqueUntouched, false},
-		{indexKey, deleteVal, false},
+		{nonUniqueIndexKey, nonUniqueVal, false},
+		{nonUniqueIndexKey, nonUniqueUntouched, false},
+		{uniqueIndexKey, uniqueValue, true},
+		{uniqueIndexKey, uniqueUntouched, false},
+		{uniqueIndexKey, deleteVal, false},
 	}
 
 	for _, test := range tests {
-		require.Equal(t, test.need, keyNeedToLock(test.key, test.val, 0))
-	}
+		need := keyNeedToLock(test.key, test.val, 0)
+		require.Equal(t, test.need, need)
 
-	flag := kv.KeyFlags(1)
-	require.True(t, flag.HasPresumeKeyNotExists())
-	require.True(t, keyNeedToLock(indexKey, deleteVal, flag))
+		flag := kv.KeyFlags(1)
+		need = keyNeedToLock(test.key, test.val, flag)
+		require.True(t, flag.HasPresumeKeyNotExists())
+		require.True(t, need)
+	}
 }
diff --git a/session/txn.go b/session/txn.go
index ce36621ba8afe..0620d59d1cf3c 100644
--- a/session/txn.go
+++ b/session/txn.go
@@ -447,9 +447,12 @@ func keyNeedToLock(k, v []byte, flags kv.KeyFlags) bool {
 	if tablecodec.IsUntouchedIndexKValue(k, v) {
 		return false
 	}
-	isNonUniqueIndex := tablecodec.IsIndexKey(k) && len(v) == 1
-	// Put row key and unique index need to lock.
-	return !isNonUniqueIndex
+
+	if !tablecodec.IsIndexKey(k) {
+		return true
+	}
+
+	return tablecodec.IndexKVIsUnique(v)
 }
 
 func getBinlogMutation(ctx sessionctx.Context, tableID int64) *binlog.TableMutation {
diff --git a/tablecodec/tablecodec.go b/tablecodec/tablecodec.go
index 90adccbc5fa64..3e34a10e30d22 100644
--- a/tablecodec/tablecodec.go
+++ b/tablecodec/tablecodec.go
@@ -1580,3 +1580,16 @@ func decodeIndexKvGeneral(key, value []byte, colsLen int, hdStatus HandleStatus,
 	}
 	return resultValues, nil
 }
+
+// IndexKVIsUnique uses to judge if an index is unique, it can handle the KV committed by txn already, it doesn't consider the untouched flag.
+func IndexKVIsUnique(value []byte) bool {
+	if len(value) <= MaxOldEncodeValueLen {
+		return len(value) == 8
+	}
+	if getIndexVersion(value) == 1 {
+		segs := SplitIndexValueForClusteredIndexVersion1(value)
+		return segs.CommonHandle != nil
+	}
+	segs := SplitIndexValue(value)
+	return segs.IntHandle != nil || segs.CommonHandle != nil
+}