planner: fix the issue of reusing wrong point-plan for "select ... for update" (#54661)

pkg/ddl/tests/metadatalock/mdl_test.go

@@ -934,7 +934,11 @@ func TestMDLPreparePlanCacheExecute(t *testing.T) {
 
 	tk.MustQuery("select * from t2")
 	tk.MustExec(`set @a = 2, @b=4;`)
-	tk.MustExec(`execute stmt_test_1 using @a, @b;`)
+	tk.MustExec(`execute stmt_test_1 using @a, @b;`) // can't reuse the prior plan created outside this txn.
+	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("0"))
+	tk.MustExec(`execute stmt_test_1 using @a, @b;`) // can't reuse the prior plan since this table becomes dirty.
+	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("0"))
+	tk.MustExec(`execute stmt_test_1 using @a, @b;`) // can reuse the prior plan now.
 	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("1"))
 	// The plan is from cache, the metadata lock should be added to block the DDL.
 	ch <- struct{}{}

pkg/planner/core/plan_cache.go

@@ -15,6 +15,7 @@
 package core
 
 import (
+	"bytes"
 	"context"
 
 	"github.com/pingcap/errors"
@@ -143,6 +144,7 @@ func planCachePreprocess(ctx context.Context, sctx sessionctx.Context, isNonPrep
 		stmt.PointGet.Plan = nil
 		stmt.PointGet.Executor = nil
 		stmt.PointGet.ColumnInfos = nil
+		stmt.PointGet.planCacheKey = nil
 		// If the schema version has changed we need to preprocess it again,
 		// if this time it failed, the real reason for the error is schema changed.
 		// Example:
@@ -225,7 +227,7 @@ func GetPlanFromSessionPlanCache(ctx context.Context, sctx sessionctx.Context,
 		}
 	}
 
-	if stmtCtx.UseCache && stmt.PointGet.Plan != nil { // special code path for fast point plan
+	if stmtCtx.UseCache && stmt.PointGet.Plan != nil && bytes.Equal(stmt.PointGet.planCacheKey.Hash(), cacheKey.Hash()) { // special code path for fast point plan
 		if plan, names, ok, err := getCachedPointPlan(stmt, sessVars, stmtCtx); ok {
 			return plan, names, err
 		}
@@ -347,7 +349,7 @@ func generateNewPlan(ctx context.Context, sctx sessionctx.Context, isNonPrepared
 	if err != nil {
 		return nil, nil, err
 	}
-	err = tryCachePointPlan(ctx, sctx.GetPlanCtx(), stmt, p, names)
+	err = tryCachePointPlan(ctx, sctx.GetPlanCtx(), stmt, p, names, cacheKey)
 	if err != nil {
 		return nil, nil, err
 	}
@@ -802,7 +804,7 @@ func CheckPreparedPriv(sctx sessionctx.Context, stmt *PlanCacheStmt, is infosche
 // tryCachePointPlan will try to cache point execution plan, there may be some
 // short paths for these executions, currently "point select" and "point update"
 func tryCachePointPlan(_ context.Context, sctx PlanContext,
-	stmt *PlanCacheStmt, p Plan, names types.NameSlice) error {
+	stmt *PlanCacheStmt, p Plan, names types.NameSlice, cacheKey kvcache.Key) error {
 	if !sctx.GetSessionVars().StmtCtx.UseCache {
 		return nil
 	}
@@ -825,6 +827,7 @@ func tryCachePointPlan(_ context.Context, sctx PlanContext,
 		// just cache point plan now
 		stmt.PointGet.Plan = p
 		stmt.PointGet.ColumnNames = names
+		stmt.PointGet.planCacheKey = cacheKey
 		stmt.NormalizedPlan, stmt.PlanDigest = NormalizePlan(p)
 		sctx.GetSessionVars().StmtCtx.SetPlan(p)
 		sctx.GetSessionVars().StmtCtx.SetPlanDigest(stmt.NormalizedPlan, stmt.PlanDigest)
@@ -846,6 +849,7 @@ func IsPointGetPlanShortPathOK(sctx sessionctx.Context, is infoschema.InfoSchema
 	if stmt.SchemaVersion != is.SchemaMetaVersion() {
 		stmt.PointGet.Plan = nil
 		stmt.PointGet.ColumnInfos = nil
+		stmt.PointGet.planCacheKey = nil
 		return false, nil
 	}
 	// maybe we'd better check cached plan type here, current

pkg/planner/core/plan_cache_test.go

@@ -1824,6 +1824,29 @@ partition by hash (a) partitions 3`)
 	tk.MustQuery(`show warnings`).Check(testkit.Rows("Warning 1105 skip prepared plan-cache: query accesses partitioned tables is un-cacheable if tidb_partition_pruning_mode = 'static'"))
 }
 
+func TestIssue54652(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec(`use test`)
+	tk.MustExec(`create table t (pk int, a int, primary key(pk))`)
+	tk.MustExec(`set autocommit=on`)
+	tk.MustQuery(`select @@autocommit`).Check(testkit.Rows("1"))
+	tk.MustExec(`set @pk=1`)
+
+	tk.MustExec(`prepare st from 'select * from t where pk=? for update'`)
+	tk.MustExec(`execute st using @pk`)
+	tk.MustExec(`execute st using @pk`)
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+
+	tk.MustExec(`begin`)
+	tk.MustExec(`execute st using @pk`)
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) // can't reuse since it's in txn now.
+	tk.MustExec(`execute st using @pk`)
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) // can reuse since it's in txn now.
+	tk.MustExec(`commit`)
+}
+
 func TestIndexRange(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

pkg/planner/core/plan_cache_utils.go

@@ -26,6 +26,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/bindinfo"
+	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/kv"
@@ -256,6 +257,11 @@ type planCacheKey struct {
 	TiDBSuperReadOnly        bool
 	exprBlacklistTS          int64 // expr-pushdown-blacklist can affect query optimization, so we need to consider it in plan cache.
 
+	// status related to Txn
+	inTxn          bool
+	autoCommit     bool
+	pessAutoCommit bool
+
 	memoryUsage int64 // Do not include in hash
 	hash        []byte
 }
@@ -279,6 +285,9 @@ func hashInt64Uint64Map(b []byte, m map[int64]uint64) []byte {
 
 // Hash implements Key interface.
 func (key *planCacheKey) Hash() []byte {
+	if key == nil {
+		return nil
+	}
 	if len(key.hash) == 0 {
 		if key.hash == nil {
 			key.hash = make([]byte, 0, len(key.stmtText)*2)
@@ -307,10 +316,20 @@ func (key *planCacheKey) Hash() []byte {
 		key.hash = append(key.hash, hack.Slice(strconv.FormatBool(key.restrictedReadOnly))...)
 		key.hash = append(key.hash, hack.Slice(strconv.FormatBool(key.TiDBSuperReadOnly))...)
 		key.hash = codec.EncodeInt(key.hash, key.exprBlacklistTS)
+		key.hash = append(key.hash, bool2Byte(key.inTxn))
+		key.hash = append(key.hash, bool2Byte(key.autoCommit))
+		key.hash = append(key.hash, bool2Byte(key.pessAutoCommit))
 	}
 	return key.hash
 }
 
+func bool2Byte(flag bool) byte {
+	if flag {
+		return '1'
+	}
+	return '0'
+}
+
 const emptyPlanCacheKeySize = int64(unsafe.Sizeof(planCacheKey{}))
 
 // MemoryUsage return the memory usage of planCacheKey
@@ -380,6 +399,9 @@ func NewPlanCacheKey(sessionVars *variable.SessionVars, stmtText, stmtDB string,
 		restrictedReadOnly:       variable.RestrictedReadOnly.Load(),
 		TiDBSuperReadOnly:        variable.VarTiDBSuperReadOnly.Load(),
 		exprBlacklistTS:          exprBlacklistTS,
+		inTxn:                    sessionVars.InTxn(),
+		autoCommit:               sessionVars.IsAutocommit(),
+		pessAutoCommit:           config.GetGlobalConfig().PessimisticTxn.PessimisticAutoCommit.Load(),
 	}
 	for k, v := range sessionVars.IsolationReadEngines {
 		key.isolationReadEngines[k] = v
@@ -508,6 +530,9 @@ type PlanCacheStmt struct {
 		// If the current plan is not PointGet or does not use MaxTS optimization, this value should be nil here.
 		Executor any
 		Plan     any // the cached PointGet Plan
+
+		// the cache key for point-get statement, have to check whether the cache key changes before reusing this plan for safety.
+		planCacheKey kvcache.Key
 	}
 
 	// below fields are for PointGet short path

pkg/planner/core/tests/prepare/prepare_test.go

@@ -756,9 +756,12 @@ func TestPlanCacheUnionScan(t *testing.T) {
 	tk.MustQuery("execute stmt1 using @p0").Check(testkit.Rows())
 	tk.MustExec("begin")
 	tk.MustQuery("execute stmt1 using @p0").Check(testkit.Rows())
+	cnt := pb.GetCounter().GetValue()
+	require.Equal(t, float64(0), cnt) // can't reuse the plan created outside the txn
+	tk.MustQuery("execute stmt1 using @p0").Check(testkit.Rows())
 	err := counter.Write(pb)
 	require.NoError(t, err)
-	cnt := pb.GetCounter().GetValue()
+	cnt = pb.GetCounter().GetValue()
 	require.Equal(t, float64(1), cnt)
 	tk.MustExec("insert into t1 values(1)")
 	// Cached plan is invalid now, it is not chosen and removed.
@@ -790,6 +793,8 @@ func TestPlanCacheUnionScan(t *testing.T) {
 	tk.MustQuery("execute stmt2 using @p0").Check(testkit.Rows())
 	tk.MustExec("begin")
 	tk.MustQuery("execute stmt2 using @p0").Check(testkit.Rows())
+	require.Equal(t, float64(3), cnt) // can't reuse the plan created outside the txn
+	tk.MustQuery("execute stmt2 using @p0").Check(testkit.Rows())
 	err = counter.Write(pb)
 	require.NoError(t, err)
 	cnt = pb.GetCounter().GetValue()