planner: use prepare text to build the plan cache key for reuse

executor/explainfor_test.go

@@ -1428,6 +1428,7 @@ func (s *testPrepareSerialSuite) TestValidity4PlanCache(c *C) {
 	tk.MustExec("create table t(a int);")
 	tk.MustExec("insert into t values(1);")
 	tk.MustQuery("execute stmt;").Check(testkit.Rows())
+	// The schema version has changed.
 	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("0"))
 	tk.MustQuery("execute stmt;").Check(testkit.Rows())
 	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("1"))
@@ -1439,7 +1440,7 @@ func (s *testPrepareSerialSuite) TestValidity4PlanCache(c *C) {
 
 	tk.MustExec("use test")
 	tk.MustQuery("execute stmt;").Check(testkit.Rows("1"))
-	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("0"))
+	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("1"))
 }
 
 func (s *testPrepareSerialSuite) TestListPartition4PlanCache(c *C) {

executor/prepared.go

@@ -227,9 +227,10 @@ func (e *PrepareExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	if e.name != "" {
 		vars.PreparedStmtNameToID[e.name] = e.ID
 	}
-
 	preparedObj := &plannercore.CachedPrepareStmt{
+		PreparedDB:          e.ctx.GetSessionVars().CurrentDB,
 		PreparedAst:         prepared,
+		PreparedStmtText:    prepared.Stmt.Text(),
 		VisitInfos:          destBuilder.GetVisitInfo(),
 		NormalizedSQL:       normalizedSQL,
 		SQLDigest:           digest,
@@ -315,7 +316,7 @@ func (e *DeallocateExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	delete(vars.PreparedStmtNameToID, e.Name)
 	if plannercore.PreparedPlanCacheEnabled() {
 		e.ctx.PreparedPlanCache().Delete(plannercore.NewPSTMTPlanCacheKey(
-			vars, id, prepared.SchemaVersion,
+			vars, preparedObj.PreparedDB, preparedObj.PreparedStmtText, prepared.SchemaVersion,
 		))
 	}
 	vars.RemovePreparedStmt(id)

planner/core/cache.go

@@ -68,7 +68,7 @@ func PreparedPlanCacheEnabled() bool {
 type pstmtPlanCacheKey struct {
 	database             string
 	connID               uint64
-	pstmtID              uint32
+	preparedStmtText     string
 	schemaVersion        int64
 	sqlMode              mysql.SQLMode
 	timezoneOffset       int
@@ -90,7 +90,7 @@ func (key *pstmtPlanCacheKey) Hash() []byte {
 		}
 		key.hash = append(key.hash, dbBytes...)
 		key.hash = codec.EncodeInt(key.hash, int64(key.connID))
-		key.hash = codec.EncodeInt(key.hash, int64(key.pstmtID))
+		key.hash = append(key.hash, hack.Slice(key.preparedStmtText)...)
 		key.hash = codec.EncodeInt(key.hash, key.schemaVersion)
 		key.hash = codec.EncodeInt(key.hash, int64(key.sqlMode))
 		key.hash = codec.EncodeInt(key.hash, int64(key.timezoneOffset))
@@ -108,32 +108,16 @@ func (key *pstmtPlanCacheKey) Hash() []byte {
 	return key.hash
 }
 
-// SetPstmtIDSchemaVersion implements PstmtCacheKeyMutator interface to change pstmtID and schemaVersion of cacheKey.
-// so we can reuse Key instead of new every time.
-func SetPstmtIDSchemaVersion(key kvcache.Key, pstmtID uint32, schemaVersion int64, isolationReadEngines map[kv.StoreType]struct{}) {
-	psStmtKey, isPsStmtKey := key.(*pstmtPlanCacheKey)
-	if !isPsStmtKey {
-		return
-	}
-	psStmtKey.pstmtID = pstmtID
-	psStmtKey.schemaVersion = schemaVersion
-	psStmtKey.isolationReadEngines = make(map[kv.StoreType]struct{})
-	for k, v := range isolationReadEngines {
-		psStmtKey.isolationReadEngines[k] = v
-	}
-	psStmtKey.hash = psStmtKey.hash[:0]
-}
-
 // NewPSTMTPlanCacheKey creates a new pstmtPlanCacheKey object.
-func NewPSTMTPlanCacheKey(sessionVars *variable.SessionVars, pstmtID uint32, schemaVersion int64) kvcache.Key {
+func NewPSTMTPlanCacheKey(sessionVars *variable.SessionVars, db string, preparedStmtText string, schemaVersion int64) kvcache.Key {
 	timezoneOffset := 0
 	if sessionVars.TimeZone != nil {
 		_, timezoneOffset = time.Now().In(sessionVars.TimeZone).Zone()
 	}
 	key := &pstmtPlanCacheKey{
-		database:             sessionVars.CurrentDB,
+		database:             db,
 		connID:               sessionVars.ConnectionID,
-		pstmtID:              pstmtID,
+		preparedStmtText:     preparedStmtText,
 		schemaVersion:        schemaVersion,
 		sqlMode:              sessionVars.SQLMode,
 		timezoneOffset:       timezoneOffset,
@@ -199,7 +183,9 @@ func NewPSTMTPlanCacheValue(plan Plan, names []*types.FieldName, srcMap map[*mod
 
 // CachedPrepareStmt store prepared ast from PrepareExec and other related fields
 type CachedPrepareStmt struct {
+	PreparedDB          string
 	PreparedAst         *ast.Prepared
+	PreparedStmtText    string
 	VisitInfos          []visitInfo
 	ColumnInfos         interface{}
 	Executor            interface{}

planner/core/cache_test.go

@@ -29,6 +29,6 @@ func TestCacheKey(t *testing.T) {
 	ctx.GetSessionVars().SQLMode = mysql.ModeNone
 	ctx.GetSessionVars().TimeZone = time.UTC
 	ctx.GetSessionVars().ConnectionID = 0
-	key := NewPSTMTPlanCacheKey(ctx.GetSessionVars(), 1, 1)
-	require.Equal(t, []byte{0x74, 0x65, 0x73, 0x74, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x74, 0x69, 0x64, 0x62, 0x74, 0x69, 0x6b, 0x76, 0x74, 0x69, 0x66, 0x6c, 0x61, 0x73, 0x68, 0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}, key.Hash())
+	key := NewPSTMTPlanCacheKey(ctx.GetSessionVars(), "", "", 1)
+	require.Equal(t, []byte{0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x80, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x74, 0x69, 0x64, 0x62, 0x74, 0x69, 0x6b, 0x76, 0x74, 0x69, 0x66, 0x6c, 0x61, 0x73, 0x68, 0x7f, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff}, key.Hash())
 }

planner/core/common_plans.go

@@ -387,7 +387,7 @@ func (e *Execute) getPhysicalPlan(ctx context.Context, sctx sessionctx.Context,
 	stmtCtx.UseCache = prepared.UseCache
 	var cacheKey kvcache.Key
 	if prepared.UseCache {
-		cacheKey = NewPSTMTPlanCacheKey(sctx.GetSessionVars(), e.ExecID, prepared.SchemaVersion)
+		cacheKey = NewPSTMTPlanCacheKey(sctx.GetSessionVars(), preparedStmt.PreparedDB, preparedStmt.PreparedStmtText, prepared.SchemaVersion)
 	}
 	tps := make([]*types.FieldType, len(e.UsingVars))
 	for i, param := range e.UsingVars {
@@ -485,7 +485,7 @@ REBUILD:
 		// rebuild key to exclude kv.TiFlash when stmt is not read only
 		if _, isolationReadContainTiFlash := sessVars.IsolationReadEngines[kv.TiFlash]; isolationReadContainTiFlash && !IsReadOnly(stmt, sessVars) {
 			delete(sessVars.IsolationReadEngines, kv.TiFlash)
-			cacheKey = NewPSTMTPlanCacheKey(sctx.GetSessionVars(), e.ExecID, prepared.SchemaVersion)
+			cacheKey = NewPSTMTPlanCacheKey(sctx.GetSessionVars(), preparedStmt.PreparedDB, preparedStmt.PreparedStmtText, prepared.SchemaVersion)
 			sessVars.IsolationReadEngines[kv.TiFlash] = struct{}{}
 		}
 		cached := NewPSTMTPlanCacheValue(p, names, stmtCtx.TblInfo2UnionScan, tps)

planner/core/prepare_test.go

@@ -1135,6 +1135,40 @@ func (s *testPlanSerialSuite) TestIssue29303(c *C) {
 	tk.MustQuery(`select @@last_plan_from_cache`).Check(testkit.Rows("0"))
 }
 
+func (s *testPlanSerialSuite) TestIssue31056(c *C) {
+	defer testleak.AfterTest(c)()
+	store, dom, err := newStoreWithBootstrap()
+	c.Assert(err, IsNil)
+	tk := testkit.NewTestKit(c, store)
+	orgEnable := core.PreparedPlanCacheEnabled()
+	defer func() {
+		dom.Close()
+		err = store.Close()
+		c.Assert(err, IsNil)
+		core.SetPreparedPlanCache(orgEnable)
+	}()
+	core.SetPreparedPlanCache(true)
+
+	tk.Se, err = session.CreateSession4TestWithOpt(store, &session.Opt{
+		PreparedPlanCache: kvcache.NewSimpleLRUCache(100, 0.1, math.MaxUint64),
+	})
+
+	tk.MustExec(`use test`)
+	tk.MustExec(`drop table if exists t`)
+	tk.MustExec(`create table t(a int, b int)`)
+	tk.MustExec(`prepare stmt from 'select * from t where a > ? and b = 2';`)
+	tk.MustExec("set @a = 1;")
+	tk.MustExec("execute stmt using @a;")
+	tk.MustExec(`prepare stmt1 from 'select * from t where a > 1 and b = ?';`)
+	tk.MustExec("set @b = 2;")
+	tk.MustExec("execute stmt1 using @b;")
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0"))
+	tk.MustExec(`prepare stmt2 from 'select * from t where a > ? and b = 2';`)
+	tk.MustExec("set @a = 1;")
+	tk.MustExec("execute stmt2 using @a;")
+	tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
+}
+
 func (s *testPlanSerialSuite) TestIssue28942(c *C) {
 	defer testleak.AfterTest(c)()
 	store, dom, err := newStoreWithBootstrap()

server/driver_tidb.go

@@ -19,7 +19,6 @@ import (
 	"crypto/tls"
 	"sync/atomic"
 
-	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/charset"
@@ -158,15 +157,6 @@ func (ts *TiDBStatement) Close() error {
 			return err
 		}
 	} else {
-		if core.PreparedPlanCacheEnabled() {
-			preparedPointer := ts.ctx.GetSessionVars().PreparedStmts[ts.id]
-			preparedObj, ok := preparedPointer.(*core.CachedPrepareStmt)
-			if !ok {
-				return errors.Errorf("invalid CachedPrepareStmt type")
-			}
-			ts.ctx.PreparedPlanCache().Delete(core.NewPSTMTPlanCacheKey(
-				ts.ctx.GetSessionVars(), ts.id, preparedObj.PreparedAst.SchemaVersion))
-		}
 		ts.ctx.GetSessionVars().RemovePreparedStmt(ts.id)
 	}
 	delete(ts.ctx.stmts, int(ts.id))

session/session.go

@@ -305,22 +305,16 @@ func (s *session) cleanRetryInfo() {
 	planCacheEnabled := plannercore.PreparedPlanCacheEnabled()
 	var cacheKey kvcache.Key
 	var preparedAst *ast.Prepared
-	if planCacheEnabled {
-		firstStmtID := retryInfo.DroppedPreparedStmtIDs[0]
-		if preparedPointer, ok := s.sessionVars.PreparedStmts[firstStmtID]; ok {
-			preparedObj, ok := preparedPointer.(*plannercore.CachedPrepareStmt)
-			if ok {
-				preparedAst = preparedObj.PreparedAst
-				cacheKey = plannercore.NewPSTMTPlanCacheKey(s.sessionVars, firstStmtID, preparedAst.SchemaVersion)
-			}
-		}
-	}
-	for i, stmtID := range retryInfo.DroppedPreparedStmtIDs {
+	for _, stmtID := range retryInfo.DroppedPreparedStmtIDs {
 		if planCacheEnabled {
-			if i > 0 && preparedAst != nil {
-				plannercore.SetPstmtIDSchemaVersion(cacheKey, stmtID, preparedAst.SchemaVersion, s.sessionVars.IsolationReadEngines)
+			if preparedPointer, ok := s.sessionVars.PreparedStmts[stmtID]; ok {
+				preparedObj, ok := preparedPointer.(*plannercore.CachedPrepareStmt)
+				if ok {
+					preparedAst = preparedObj.PreparedAst
+					cacheKey = plannercore.NewPSTMTPlanCacheKey(s.sessionVars, preparedObj.PreparedDB, preparedObj.PreparedStmtText, preparedAst.SchemaVersion)
+					s.PreparedPlanCache().Delete(cacheKey)
+				}
 			}
-			s.PreparedPlanCache().Delete(cacheKey)
 		}
 		s.sessionVars.RemovePreparedStmt(stmtID)
 	}