planner: ignore close-stmt to make more queries can hit plan cache

executor/prepared.go

@@ -336,7 +336,9 @@ func (e *DeallocateExec) Next(ctx context.Context, req *chunk.Chunk) error {
 		if err != nil {
 			return err
 		}
-		e.ctx.PreparedPlanCache().Delete(cacheKey)
+		if !vars.IgnoreCloseStmtCmd { // keep the plan in cache if this flag is true
+			e.ctx.PreparedPlanCache().Delete(cacheKey)
+		}
 	}
 	vars.RemovePreparedStmt(id)
 	return nil

executor/set_test.go

@@ -570,6 +570,15 @@ func TestSetVar(t *testing.T) {
 	require.Error(t, tk.ExecToErr("select @@session.tidb_enable_column_tracking"))
 	require.Error(t, tk.ExecToErr("set tidb_enable_column_tracking = 0"))
 	require.Error(t, tk.ExecToErr("set global tidb_enable_column_tracking = -1"))
+
+	// test for tidb_ignore_close_stmt_cmd
+	tk.MustQuery("select @@global.tidb_ignore_close_stmt_cmd").Check(testkit.Rows("0")) // default value is 0
+	tk.MustExec("set global tidb_ignore_close_stmt_cmd=1")
+	tk.MustQuery("select @@global.tidb_ignore_close_stmt_cmd").Check(testkit.Rows("1"))
+	tk.MustQuery("show global variables like 'tidb_ignore_close_stmt_cmd'").Check(testkit.Rows("tidb_ignore_close_stmt_cmd ON"))
+	tk.MustExec("set global tidb_ignore_close_stmt_cmd=0")
+	tk.MustQuery("select @@global.tidb_ignore_close_stmt_cmd").Check(testkit.Rows("0"))
+	tk.MustQuery("show global variables like 'tidb_ignore_close_stmt_cmd'").Check(testkit.Rows("tidb_ignore_close_stmt_cmd OFF"))
 }
 
 func TestTruncateIncorrectIntSessionVar(t *testing.T) {

planner/core/prepare_test.go

@@ -42,6 +42,80 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestPreparePointGetWithDML(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+	orgEnable := core.PreparedPlanCacheEnabled()
+	defer core.SetPreparedPlanCache(orgEnable)
+	core.SetPreparedPlanCache(true)
+	se, err := session.CreateSession4TestWithOpt(store, &session.Opt{
+		PreparedPlanCache: kvcache.NewSimpleLRUCache(100, 0.1, math.MaxUint64),
+	})
+	require.NoError(t, err)
+	tk := testkit.NewTestKitWithSession(t, store, se)
+
+	tk.MustExec(`use test`)
+	tk.MustExec(`drop table if exists t`)
+	tk.MustExec(`create table t(a int, unique key(a))`)
+	tk.MustExec(`insert into t values(1), (2)`)
+
+	// txn1 left a cached plan
+	tk.MustExec(`begin`)
+	tk.MustExec(`prepare stmt from 'update t set a = ? where  a = ?'`)
+	tk.MustExec(`set @a=1`)
+	tk.MustExec(`execute stmt using @a, @a`)
+	tk.MustExec(`commit`)
+
+	// txn2 can reuse the cached plan generated by txn1 directly
+	tk.MustExec(`begin`)
+	tk.MustExec(`prepare stmt from 'update t set a = ? where  a = ?'`)
+	tk.MustExec(`set @a=2`)
+	tk.MustExec(`execute stmt using @a, @a`) // can reuse the cached plan directly
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+	tk.MustExec(`rollback`)
+}
+
+func TestPrepareIgnoreCloseStmtCmd(t *testing.T) {
+	store, clean := testkit.CreateMockStore(t)
+	defer clean()
+	orgEnable := core.PreparedPlanCacheEnabled()
+	defer core.SetPreparedPlanCache(orgEnable)
+	core.SetPreparedPlanCache(true)
+	se, err := session.CreateSession4TestWithOpt(store, &session.Opt{
+		PreparedPlanCache: kvcache.NewSimpleLRUCache(100, 0.1, math.MaxUint64),
+	})
+	require.NoError(t, err)
+	tk := testkit.NewTestKitWithSession(t, store, se)
+
+	tk.MustExec(`use test`)
+	tk.MustExec(`drop table if exists t`)
+	tk.MustExec(`create table t(a int, unique key(a))`)
+
+	// disable the ignore-stmt-cmd
+	tk.MustExec(`set @@tidb_ignore_close_stmt_cmd=0`)
+	tk.MustQuery(`select @@tidb_ignore_close_stmt_cmd`).Check(testkit.Rows("0"))
+	tk.MustExec(`prepare stmt from 'select * from t'`)
+	tk.MustQuery(`execute stmt`).Check(testkit.Rows())
+	tk.MustQuery(`execute stmt`).Check(testkit.Rows())
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+	tk.MustExec(`deallocate prepare stmt`) // close/deallocate this stmt
+	tk.MustExec(`prepare stmt from 'select * from t'`)
+	tk.MustQuery(`execute stmt`).Check(testkit.Rows())
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0")) // cannot reuse last plan since it is closed
+
+	// enable the ignore-stmt-cmd
+	tk.MustExec(`set tidb_ignore_close_stmt_cmd=1`)
+	tk.MustQuery(`select @@tidb_ignore_close_stmt_cmd`).Check(testkit.Rows("1"))
+	tk.MustExec(`prepare stmt from 'select * from t'`)
+	tk.MustQuery(`execute stmt`).Check(testkit.Rows())
+	tk.MustQuery(`execute stmt`).Check(testkit.Rows())
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1"))
+	tk.MustExec(`deallocate prepare stmt`) // close/deallocate this stmt
+	tk.MustExec(`prepare stmt from 'select * from t'`)
+	tk.MustQuery(`execute stmt`).Check(testkit.Rows())
+	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("1")) // can reuse last plan since last close-cmd was ignored
+}
+
 func TestRandomFlushPlanCache(t *testing.T) {
 	store, clean := testkit.CreateMockStore(t)
 	defer clean()

server/driver_tidb.go

@@ -169,7 +169,9 @@ func (ts *TiDBStatement) Close() error {
 			if err != nil {
 				return err
 			}
-			ts.ctx.PreparedPlanCache().Delete(cacheKey)
+			if !ts.ctx.GetSessionVars().IgnoreCloseStmtCmd { // keep the plan in cache if this flag is true
+				ts.ctx.PreparedPlanCache().Delete(cacheKey)
+			}
 		}
 		ts.ctx.GetSessionVars().RemovePreparedStmt(ts.id)
 	}

session/session.go

@@ -341,7 +341,9 @@ func (s *session) cleanRetryInfo() {
 			if i > 0 && preparedAst != nil {
 				plannercore.SetPstmtIDSchemaVersion(cacheKey, stmtText, preparedAst.SchemaVersion, s.sessionVars.IsolationReadEngines)
 			}
-			s.PreparedPlanCache().Delete(cacheKey)
+			if !s.sessionVars.IgnoreCloseStmtCmd { // keep the plan in cache if this flag is true
+				s.PreparedPlanCache().Delete(cacheKey)
+			}
 		}
 		s.sessionVars.RemovePreparedStmt(stmtID)
 	}

sessionctx/variable/session.go

@@ -1011,6 +1011,8 @@ type SessionVars struct {
 	EnableMutationChecker bool
 	// AssertionLevel controls how strict the assertions on data mutations should be.
 	AssertionLevel AssertionLevel
+	// IgnoreCloseStmtCmd controls if ignore the close-prepared-stmt command.
+	IgnoreCloseStmtCmd bool
 }
 
 // InitStatementContext initializes a StatementContext, the object is reused to reduce allocation.

sessionctx/variable/sysvar.go

@@ -1363,6 +1363,12 @@ var defaultSysVars = []*SysVar{
 		}
 		return string(info), nil
 	}},
+	{Scope: ScopeGlobal | ScopeSession, Name: TiDBIgnoreCloseStmtCmd, Value: BoolToOnOff(DefTiDBIgnoreCloseStmtCmd), Type: TypeBool,
+		SetSession: func(vars *SessionVars, s string) error {
+			vars.IgnoreCloseStmtCmd = TiDBOptOn(s)
+			return nil
+		},
+	},
 }
 
 // FeedbackProbability points to the FeedbackProbability in statistics package.

sessionctx/variable/tidb_vars.go

@@ -607,6 +607,9 @@ const (
 	// TiDBTxnAssertionLevel indicates how strict the assertion will be, which helps to detect and preventing data &
 	// index inconsistency problems.
 	TiDBTxnAssertionLevel = "tidb_txn_assertion_level"
+
+	// TiDBIgnoreClosePreparedStmt indicates whether to ignore close-stmt commands for prepared statements.
+	TiDBIgnoreCloseStmtCmd = "tidb_ignore_close_stmt_cmd"
 )
 
 // TiDB vars that have only global scope
@@ -801,6 +804,7 @@ const (
 	DefSysdateIsNow                       = false
 	DefTiDBEnableMutationChecker          = false
 	DefTiDBTxnAssertionLevel              = AssertionOffStr
+	DefTiDBIgnoreCloseStmtCmd             = false
 )
 
 // Process global variables.