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.IgnorePreparedCacheCloseStmt { // keep the plan in cache
+			e.ctx.PreparedPlanCache().Delete(cacheKey)
+		}
 	}
 	vars.RemovePreparedStmt(id)
 	return nil

executor/set_test.go

@@ -576,6 +576,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_prepared_cache_close_stmt
+	tk.MustQuery("select @@global.tidb_ignore_prepared_cache_close_stmt").Check(testkit.Rows("0")) // default value is 0
+	tk.MustExec("set global tidb_ignore_prepared_cache_close_stmt=1")
+	tk.MustQuery("select @@global.tidb_ignore_prepared_cache_close_stmt").Check(testkit.Rows("1"))
+	tk.MustQuery("show global variables like 'tidb_ignore_prepared_cache_close_stmt'").Check(testkit.Rows("tidb_ignore_prepared_cache_close_stmt ON"))
+	tk.MustExec("set global tidb_ignore_prepared_cache_close_stmt=0")
+	tk.MustQuery("select @@global.tidb_ignore_prepared_cache_close_stmt").Check(testkit.Rows("0"))
+	tk.MustQuery("show global variables like 'tidb_ignore_prepared_cache_close_stmt'").Check(testkit.Rows("tidb_ignore_prepared_cache_close_stmt 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_prepared_cache_close_stmt=0`)
+	tk.MustQuery(`select @@tidb_ignore_prepared_cache_close_stmt`).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_prepared_cache_close_stmt=1`)
+	tk.MustQuery(`select @@tidb_ignore_prepared_cache_close_stmt`).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().IgnorePreparedCacheCloseStmt { // keep the plan in cache
+				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.IgnorePreparedCacheCloseStmt { // keep the plan in cache
+				s.PreparedPlanCache().Delete(cacheKey)
+			}
 		}
 		s.sessionVars.RemovePreparedStmt(stmtID)
 	}

sessionctx/variable/session.go

@@ -1004,6 +1004,8 @@ type SessionVars struct {
 	EnableMutationChecker bool
 	// AssertionLevel controls how strict the assertions on data mutations should be.
 	AssertionLevel AssertionLevel
+	// IgnorePreparedCacheCloseStmt controls if ignore the close-stmt command for prepared statement.
+	IgnorePreparedCacheCloseStmt bool
 	// BatchPendingTiFlashCount shows the threshold of pending TiFlash tables when batch adding.
 	BatchPendingTiFlashCount int
 }

sessionctx/variable/sysvar.go

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

sessionctx/variable/tidb_vars.go

@@ -605,6 +605,9 @@ const (
 	// index inconsistency problems.
 	TiDBTxnAssertionLevel = "tidb_txn_assertion_level"
 
+	// TiDBIgnorePreparedCacheCloseStmt indicates whether to ignore close-stmt commands for prepared statements.
+	TiDBIgnorePreparedCacheCloseStmt = "tidb_ignore_prepared_cache_close_stmt"
+
 	// TiDBBatchPendingTiFlashCount indicates the maximum count of non-available TiFlash tables.
 	TiDBBatchPendingTiFlashCount = "tidb_batch_pending_tiflash_count"
 )
@@ -800,6 +803,7 @@ const (
 	DefSysdateIsNow                       = false
 	DefTiDBEnableMutationChecker          = false
 	DefTiDBTxnAssertionLevel              = AssertionOffStr
+	DefTiDBIgnorePreparedCacheCloseStmt   = false
 	DefTiDBBatchPendingTiFlashCount       = 4000
 )