executor: optimize cursor read point get by reading through pessimistic lock cache

executor/point_get.go

@@ -237,19 +237,32 @@ func (e *PointGetExecutor) Next(ctx context.Context, req *chunk.Chunk) error {
 				return err
 			}
 
+			// lockNonExistIdxKey indicates the key will be locked regardless of its existence.
+			lockNonExistIdxKey := !e.ctx.GetSessionVars().IsPessimisticReadConsistency()
+			// Non-exist keys are also locked if the isolation level is not read consistency,
+			// lock it before read here, then it's able to read from pessimistic lock cache.
+			if lockNonExistIdxKey {
+				err = e.lockKeyIfNeeded(ctx, e.idxKey)
+				if err != nil {
+					return err
+				}
+			}
+
 			e.handleVal, err = e.get(ctx, e.idxKey)
 			if err != nil {
 				if !kv.ErrNotExist.Equal(err) {
 					return err
 				}
 			}
 
-			// try lock the index key if isolation level is not read consistency
 			// also lock key if read consistency read a value
-			if !e.ctx.GetSessionVars().IsPessimisticReadConsistency() || len(e.handleVal) > 0 {
-				err = e.lockKeyIfNeeded(ctx, e.idxKey)
-				if err != nil {
-					return err
+			// TODO: pessimistic lock support lock-if-exist.
+			if lockNonExistIdxKey || len(e.handleVal) > 0 {
+				if !lockNonExistIdxKey {
+					err = e.lockKeyIfNeeded(ctx, e.idxKey)
+					if err != nil {
+						return err
+					}
 				}
 				// Change the unique index LOCK into PUT record.
 				if e.lock && len(e.handleVal) > 0 {
@@ -377,7 +390,7 @@ func (e *PointGetExecutor) lockKeyIfNeeded(ctx context.Context, key []byte) erro
 			return err
 		}
 		lockCtx.IterateValuesNotLocked(func(k, v []byte) {
-			seVars.TxnCtx.SetPessimisticLockCache(kv.Key(k), v)
+			seVars.TxnCtx.SetPessimisticLockCache(k, v)
 		})
 		if len(e.handleVal) > 0 {
 			seVars.TxnCtx.SetPessimisticLockCache(e.idxKey, e.handleVal)

server/conn_test.go

@@ -829,13 +829,13 @@ func TestPrefetchPointKeys(t *testing.T) {
 
 	tk.MustExec("begin pessimistic")
 	tk.MustExec("update prefetch set c = c + 1 where a = 2 and b = 2")
-	require.Equal(t, 1, tk.Session().GetSessionVars().TxnCtx.PessimisticCacheHit)
+	require.Equal(t, 2, tk.Session().GetSessionVars().TxnCtx.PessimisticCacheHit)
 	err = cc.handleQuery(ctx, query)
 	require.NoError(t, err)
 	txn, err = tk.Session().Txn(false)
 	require.NoError(t, err)
 	require.True(t, txn.Valid())
-	require.Equal(t, 5, tk.Session().GetSessionVars().TxnCtx.PessimisticCacheHit)
+	require.Equal(t, 6, tk.Session().GetSessionVars().TxnCtx.PessimisticCacheHit)
 	tk.MustExec("commit")
 	tk.MustQuery("select * from prefetch").Check(testkit.Rows("1 1 3", "2 2 6", "3 3 5"))
 }