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

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't 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

@@ -135,6 +135,7 @@ func planCachePreprocess(ctx context.Context, sctx sessionctx.Context, isNonPrep
 		// Cached plan in prepared struct does NOT have a "cache key" with
 		// schema version like prepared plan cache key
 		stmt.PointGet.pointPlan = nil
+		stmt.PointGet.planCacheKey = ""
 		stmt.PointGet.columnNames = nil
 		stmt.PointGet.pointPlanHints = nil
 		stmt.PointGet.Executor = nil
@@ -219,7 +220,7 @@ func GetPlanFromPlanCache(ctx context.Context, sctx sessionctx.Context,
 	if stmtCtx.UseCache() {
 		var cachedVal *PlanCacheValue
 		var hit, isPointPlan bool
-		if stmt.PointGet.pointPlan != nil { // if it's PointGet Plan, no need to use paramTypes
+		if stmt.PointGet.pointPlan != nil && stmt.PointGet.planCacheKey == cacheKey { // if it's PointGet Plan, no need to use paramTypes
 			cachedVal = &PlanCacheValue{
 				Plan:          stmt.PointGet.pointPlan,
 				OutputColumns: stmt.PointGet.columnNames,
@@ -322,6 +323,7 @@ func generateNewPlan(ctx context.Context, sctx sessionctx.Context, isNonPrepared
 			stmt.PointGet.pointPlan = p
 			stmt.PointGet.columnNames = names
 			stmt.PointGet.pointPlanHints = stmtCtx.StmtHints.Clone()
+			stmt.PointGet.planCacheKey = cacheKey
 		}
 	}
 	sessVars.FoundInPlanCache = false

pkg/planner/core/plan_cache_test.go

@@ -1907,3 +1907,24 @@ func TestInstancePlanCacheAcrossSession(t *testing.T) {
 	tk2.MustQuery(`execute st using @a`).Sort().Check(testkit.Rows(`1`, `2`, `3`))
 	tk2.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
 }
+
+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(`commit`)
+}

pkg/planner/core/plan_cache_utils.go

@@ -25,6 +25,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/bindinfo"
+	"github.com/pingcap/tidb/pkg/config"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
@@ -330,11 +331,7 @@ func NewPlanCacheKey(sctx sessionctx.Context, stmt *PlanCacheStmt) (key, binding
 	}
 
 	// this variable might affect the plan
-	if vars.ForeignKeyChecks {
-		hash = append(hash, '1')
-	} else {
-		hash = append(hash, '0')
-	}
+	hash = append(hash, bool2Byte(vars.ForeignKeyChecks))
 
 	// "limit ?" can affect the cached plan: "limit 1" and "limit 10000" should use different plans.
 	if len(stmt.limits) > 0 {
@@ -386,9 +383,23 @@ func NewPlanCacheKey(sctx sessionctx.Context, stmt *PlanCacheStmt) (key, binding
 			hash = codec.EncodeInt(hash, id)
 		}
 	}
+
+	// txn status
+	hash = append(hash, '|')
+	hash = append(hash, bool2Byte(vars.InTxn()))
+	hash = append(hash, bool2Byte(vars.IsAutocommit()))
+	hash = append(hash, bool2Byte(config.GetGlobalConfig().PessimisticTxn.PessimisticAutoCommit.Load()))
+
 	return string(hash), binding, true, "", nil
 }
 
+func bool2Byte(flag bool) byte {
+	if flag {
+		return '1'
+	}
+	return '0'
+}
+
 // PlanCacheValue stores the cached Statement and StmtNode.
 type PlanCacheValue struct {
 	Plan          base.Plan          // not-read-only, session might update it before reusing
@@ -515,6 +526,9 @@ type PointGetExecutorCache struct {
 	pointPlanHints *hint.StmtHints
 	columnNames    types.NameSlice
 
+	// the cache key for this statement, have to check whether the cache key changes before reusing this plan for safety.
+	planCacheKey string
+
 	ColumnInfos any
 	// Executor is only used for point get scene.
 	// Notice that we should only cache the PointGetExecutor that have a snapshot with MaxTS in it.

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

@@ -756,10 +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()
-	require.Equal(t, float64(1), cnt)
+	require.Equal(t, float64(0), cnt)
 	tk.MustExec("insert into t1 values(1)")
 	// Cached plan is invalid now, it is not chosen and removed.
 	tk.MustQuery("execute stmt1 using @p0").Check(testkit.Rows(
@@ -790,6 +792,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()