planner: prepared plan cache support cached plan with placeholder in limit clause

executor/prepared.go

@@ -204,8 +204,9 @@ func (e *DeallocateExec) Next(ctx context.Context, req *chunk.Chunk) error {
 	delete(vars.PreparedStmtNameToID, e.Name)
 	if e.ctx.GetSessionVars().EnablePreparedPlanCache {
 		bindSQL, _ := plannercore.GetBindSQL4PlanCache(e.ctx, preparedObj)
+		limitOffsetAndCount, _ := plannercore.ExtractLimitFromAst(prepared.Stmt)
 		cacheKey, err := plannercore.NewPlanCacheKey(vars, preparedObj.StmtText, preparedObj.StmtDB, prepared.SchemaVersion,
-			0, bindSQL)
+			0, bindSQL, limitOffsetAndCount)
 		if err != nil {
 			return err
 		}

planner/core/plan_cache.go

@@ -131,6 +131,7 @@ func GetPlanFromSessionPlanCache(ctx context.Context, sctx sessionctx.Context,
 
 	var bindSQL string
 	var ignorePlanCache = false
+	var canBeCached = true
 
 	// In rc or for update read, we need the latest schema version to decide whether we need to
 	// rebuild the plan. So we set this value in rc or for update read. In other cases, let it be 0.
@@ -144,8 +145,12 @@ func GetPlanFromSessionPlanCache(ctx context.Context, sctx sessionctx.Context,
 			// up-to-date schema version which can lead plan cache miss and thus, the plan will be rebuilt.
 			latestSchemaVersion = domain.GetDomain(sctx).InfoSchema().SchemaMetaVersion()
 		}
+		var limitOffsetAndCount []int64
+		if limitOffsetAndCount, canBeCached = ExtractLimitFromAst(stmt.PreparedAst.Stmt); !canBeCached {
+			sctx.GetSessionVars().StmtCtx.AppendWarning(errors.New("plan with limit count more than 10000 can't be cached"))
+		}
 		if cacheKey, err = NewPlanCacheKey(sctx.GetSessionVars(), stmt.StmtText,
-			stmt.StmtDB, stmtAst.SchemaVersion, latestSchemaVersion, bindSQL); err != nil {
+			stmt.StmtDB, stmtAst.SchemaVersion, latestSchemaVersion, bindSQL, limitOffsetAndCount); err != nil {
 			return nil, nil, err
 		}
 	}
@@ -165,7 +170,7 @@ func GetPlanFromSessionPlanCache(ctx context.Context, sctx sessionctx.Context,
 		}
 	}
 
-	return generateNewPlan(ctx, sctx, isNonPrepared, is, stmt, ignorePlanCache, cacheKey,
+	return generateNewPlan(ctx, sctx, isNonPrepared, is, stmt, ignorePlanCache && canBeCached, cacheKey,
 		latestSchemaVersion, paramNum, paramTypes, bindSQL)
 }
 
@@ -284,8 +289,9 @@ func generateNewPlan(ctx context.Context, sctx sessionctx.Context, isNonPrepared
 		// rebuild key to exclude kv.TiFlash when stmt is not read only
 		if _, isolationReadContainTiFlash := sessVars.IsolationReadEngines[kv.TiFlash]; isolationReadContainTiFlash && !IsReadOnly(stmtAst.Stmt, sessVars) {
 			delete(sessVars.IsolationReadEngines, kv.TiFlash)
+			limitOffsetAndCount, _ := ExtractLimitFromAst(stmt.PreparedAst.Stmt)
 			if cacheKey, err = NewPlanCacheKey(sessVars, stmt.StmtText, stmt.StmtDB,
-				stmtAst.SchemaVersion, latestSchemaVersion, bindSQL); err != nil {
+				stmtAst.SchemaVersion, latestSchemaVersion, bindSQL, limitOffsetAndCount); err != nil {
 				return nil, nil, err
 			}
 			sessVars.IsolationReadEngines[kv.TiFlash] = struct{}{}

planner/core/plan_cache_test.go

@@ -283,12 +283,6 @@ func TestPlanCacheDiagInfo(t *testing.T) {
 	tk.MustExec("prepare stmt from 'select /*+ ignore_plan_cache() */ * from t'")
 	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 skip plan-cache: ignore plan cache by hint"))
 
-	tk.MustExec("prepare stmt from 'select * from t limit ?'")
-	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 skip plan-cache: query has 'limit ?' is un-cacheable"))
-
-	tk.MustExec("prepare stmt from 'select * from t limit ?, 1'")
-	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 skip plan-cache: query has 'limit ?, 10' is un-cacheable"))
-
 	tk.MustExec("prepare stmt from 'select * from t order by ?'")
 	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 skip plan-cache: query has 'order by ?' is un-cacheable"))
 
@@ -302,3 +296,48 @@ func TestPlanCacheDiagInfo(t *testing.T) {
 	tk.MustExec("execute stmt using @a, @b") // a=1 and a=1 -> a=1
 	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 skip plan-cache: some parameters may be overwritten"))
 }
+
+func TestPlanCacheWithLimit(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int, b int, key(a))")
+
+	testCases := []struct {
+		sql    string
+		params []int
+	}{
+		{"prepare stmt from 'select * from t limit ?'", []int{1}},
+		{"prepare stmt from 'select * from t limit ?, ?'", []int{1, 2}},
+		{"prepare stmt from 'delete from t order by a limit ?'", []int{1}},
+		{"prepare stmt from 'insert into t select * from t order by a desc limit ?'", []int{1}},
+		{"prepare stmt from 'insert into t select * from t order by a desc limit ?, ?'", []int{1, 2}},
+		{"prepare stmt from 'update t set a = 1 limit ?'", []int{1}},
+		{" prepare stmt from '(select * from t order by a limit ?) union (select * from t order by a desc limit ?)';", []int{1, 2}},
+	}
+
+	for _, testCase := range testCases {
+		tk.MustExec(testCase.sql)
+		tk.MustExec("set @a = 1")
+		var using []string
+		for i, p := range testCase.params {
+			tk.MustExec(fmt.Sprintf("set @a%d = %d", i, p))
+			using = append(using, fmt.Sprintf("@a%d", i))
+		}
+
+		tk.MustExec("execute stmt using " + strings.Join(using, ", "))
+		tk.MustExec("execute stmt using " + strings.Join(using, ", "))
+		tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1"))
+
+		tk.MustExec("set @a0 = 10086")
+		tk.MustExec("execute stmt using " + strings.Join(using, ", "))
+		tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0"))
+	}
+
+	// error case
+	tk.MustExec("prepare stmt from 'select * from t limit ?'")
+	tk.MustExec("set @a = 10001")
+	tk.MustExec("execute stmt using @a")
+	tk.MustQuery("show warnings").Check(testkit.Rows("Warning 1105 plan with limit count more than 10000 can't be cached"))
+}

planner/core/plan_cache_utils.go

@@ -190,6 +190,7 @@ type planCacheKey struct {
 	inRestrictedSQL          bool
 	restrictedReadOnly       bool
 	TiDBSuperReadOnly        bool
+	limitOffsetAndCount      []int64
 
 	memoryUsage int64 // Do not include in hash
 	hash        []byte
@@ -226,6 +227,9 @@ func (key *planCacheKey) Hash() []byte {
 		key.hash = append(key.hash, hack.Slice(strconv.FormatBool(key.inRestrictedSQL))...)
 		key.hash = append(key.hash, hack.Slice(strconv.FormatBool(key.restrictedReadOnly))...)
 		key.hash = append(key.hash, hack.Slice(strconv.FormatBool(key.TiDBSuperReadOnly))...)
+		for _, l := range key.limitOffsetAndCount {
+			key.hash = codec.EncodeInt(key.hash, l)
+		}
 	}
 	return key.hash
 }
@@ -267,7 +271,7 @@ func SetPstmtIDSchemaVersion(key kvcache.Key, stmtText string, schemaVersion int
 // Note: lastUpdatedSchemaVersion will only be set in the case of rc or for update read in order to
 // differentiate the cache key. In other cases, it will be 0.
 func NewPlanCacheKey(sessionVars *variable.SessionVars, stmtText, stmtDB string, schemaVersion int64,
-	lastUpdatedSchemaVersion int64, bindSQL string) (kvcache.Key, error) {
+	lastUpdatedSchemaVersion int64, bindSQL string, offsetAndCount []int64) (kvcache.Key, error) {
 	if stmtText == "" {
 		return nil, errors.New("no statement text")
 	}
@@ -295,6 +299,7 @@ func NewPlanCacheKey(sessionVars *variable.SessionVars, stmtText, stmtDB string,
 		inRestrictedSQL:          sessionVars.InRestrictedSQL,
 		restrictedReadOnly:       variable.RestrictedReadOnly.Load(),
 		TiDBSuperReadOnly:        variable.VarTiDBSuperReadOnly.Load(),
+		limitOffsetAndCount:      offsetAndCount,
 	}
 	for k, v := range sessionVars.IsolationReadEngines {
 		key.isolationReadEngines[k] = v
@@ -445,3 +450,48 @@ func GetPreparedStmt(stmt *ast.ExecuteStmt, vars *variable.SessionVars) (*PlanCa
 	}
 	return nil, ErrStmtNotFound
 }
+
+type limitExtractor struct {
+	cacheable      bool // For safety considerations, check if limit count less than 10000
+	offsetAndCount []int64
+}
+
+// Enter implements Visitor interface.
+func (checker *limitExtractor) Enter(in ast.Node) (out ast.Node, skipChildren bool) {
+	switch node := in.(type) {
+	case *ast.Limit:
+		if node.Count != nil {
+			if count, isParamMarker := node.Count.(*driver.ParamMarkerExpr); isParamMarker {
+				countNum := count.GetInt64()
+				checker.offsetAndCount = append(checker.offsetAndCount, countNum)
+				if countNum > 10000 {
+					checker.cacheable = false
+				}
+			}
+		}
+		if node.Offset != nil {
+			if offset, isParamMarker := node.Offset.(*driver.ParamMarkerExpr); isParamMarker {
+				checker.offsetAndCount = append(checker.offsetAndCount, offset.GetInt64())
+			}
+		}
+	}
+	return in, false
+}
+
+// Leave implements Visitor interface.
+func (checker *limitExtractor) Leave(in ast.Node) (out ast.Node, ok bool) {
+	return in, checker.cacheable
+}
+
+// ExtractLimitFromAst extract limit offset and count from ast for plan cache key encode
+func ExtractLimitFromAst(node ast.Node) ([]int64, bool) {
+	if node == nil {
+		return []int64{}, true
+	}
+	checker := limitExtractor{
+		cacheable:      true,
+		offsetAndCount: []int64{},
+	}
+	node.Accept(&checker)
+	return checker.offsetAndCount, checker.cacheable
+}

planner/core/plan_cache_utils_test.go

@@ -32,20 +32,20 @@ func TestCacheKey(t *testing.T) {
 	ctx.GetSessionVars().InRestrictedSQL = false
 	variable.RestrictedReadOnly.Store(false)
 	variable.VarTiDBSuperReadOnly.Store(false)
-	key, err := NewPlanCacheKey(ctx.GetSessionVars(), "", "test", 1, 1, "")
+	key, err := NewPlanCacheKey(ctx.GetSessionVars(), "", "test", 1, 1, "", []int64{})
 	if err.Error() != "no statement text" {
 		t.Fail() // no statement text
 	}
-	key, err = NewPlanCacheKey(ctx.GetSessionVars(), "select 1", "", 1, 1, "")
+	key, err = NewPlanCacheKey(ctx.GetSessionVars(), "select 1", "", 1, 1, "", []int64{})
 	if err != nil {
 		t.Fail() // schema can be nil
 	}
 	key, err = NewPlanCacheKey(ctx.GetSessionVars(), "select 1", "test", 1, 1,
-		"select /*+ ignore_plan_cache() */ * from t")
+		"select /*+ ignore_plan_cache() */ * from t", []int64{})
 	if err != nil {
 		t.Fail()
 	}
-	key, err = NewPlanCacheKey(ctx.GetSessionVars(), "select 1", "test", 1, 1, "")
+	key, err = NewPlanCacheKey(ctx.GetSessionVars(), "select 1", "test", 1, 1, "", []int64{})
 	if err != nil {
 		t.Fail()
 	}

planner/core/plan_cacheable_checker.go

@@ -135,21 +135,6 @@ func (checker *cacheableChecker) Enter(in ast.Node) (out ast.Node, skipChildren
 				return in, true
 			}
 		}
-	case *ast.Limit:
-		if node.Count != nil {
-			if _, isParamMarker := node.Count.(*driver.ParamMarkerExpr); isParamMarker {
-				checker.cacheable = false
-				checker.reason = "query has 'limit ?' is un-cacheable"
-				return in, true
-			}
-		}
-		if node.Offset != nil {
-			if _, isParamMarker := node.Offset.(*driver.ParamMarkerExpr); isParamMarker {
-				checker.cacheable = false
-				checker.reason = "query has 'limit ?, 10' is un-cacheable"
-				return in, true
-			}
-		}
 	case *ast.FrameBound:
 		if _, ok := node.Expr.(*driver.ParamMarkerExpr); ok {
 			checker.cacheable = false

planner/core/plan_cacheable_checker_test.go

@@ -87,7 +87,7 @@ func TestCacheable(t *testing.T) {
 		TableRefs: tableRefsClause,
 		Limit:     limitStmt,
 	}
-	require.False(t, core.Cacheable(stmt, is))
+	require.True(t, core.Cacheable(stmt, is))
 
 	limitStmt = &ast.Limit{
 		Offset: &driver.ParamMarkerExpr{},
@@ -96,7 +96,7 @@ func TestCacheable(t *testing.T) {
 		TableRefs: tableRefsClause,
 		Limit:     limitStmt,
 	}
-	require.False(t, core.Cacheable(stmt, is))
+	require.True(t, core.Cacheable(stmt, is))
 
 	limitStmt = &ast.Limit{}
 	stmt = &ast.DeleteStmt{
@@ -139,7 +139,7 @@ func TestCacheable(t *testing.T) {
 		TableRefs: tableRefsClause,
 		Limit:     limitStmt,
 	}
-	require.False(t, core.Cacheable(stmt, is))
+	require.True(t, core.Cacheable(stmt, is))
 
 	limitStmt = &ast.Limit{
 		Offset: &driver.ParamMarkerExpr{},
@@ -148,7 +148,7 @@ func TestCacheable(t *testing.T) {
 		TableRefs: tableRefsClause,
 		Limit:     limitStmt,
 	}
-	require.False(t, core.Cacheable(stmt, is))
+	require.True(t, core.Cacheable(stmt, is))
 
 	limitStmt = &ast.Limit{}
 	stmt = &ast.UpdateStmt{
@@ -188,15 +188,15 @@ func TestCacheable(t *testing.T) {
 	stmt = &ast.SelectStmt{
 		Limit: limitStmt,
 	}
-	require.False(t, core.Cacheable(stmt, is))
+	require.True(t, core.Cacheable(stmt, is))
 
 	limitStmt = &ast.Limit{
 		Offset: &driver.ParamMarkerExpr{},
 	}
 	stmt = &ast.SelectStmt{
 		Limit: limitStmt,
 	}
-	require.False(t, core.Cacheable(stmt, is))
+	require.True(t, core.Cacheable(stmt, is))
 
 	limitStmt = &ast.Limit{}
 	stmt = &ast.SelectStmt{

server/driver_tidb.go

@@ -171,8 +171,9 @@ func (ts *TiDBStatement) Close() error {
 				return errors.Errorf("invalid PlanCacheStmt type")
 			}
 			bindSQL, _ := core.GetBindSQL4PlanCache(ts.ctx, preparedObj)
+			limitOffsetAndCount, _ := core.ExtractLimitFromAst(preparedObj.PreparedAst.Stmt)
 			cacheKey, err := core.NewPlanCacheKey(ts.ctx.GetSessionVars(), preparedObj.StmtText, preparedObj.StmtDB,
-				preparedObj.PreparedAst.SchemaVersion, 0, bindSQL)
+				preparedObj.PreparedAst.SchemaVersion, 0, bindSQL, limitOffsetAndCount)
 			if err != nil {
 				return err
 			}

session/session.go

@@ -382,8 +382,9 @@ func (s *session) cleanRetryInfo() {
 				preparedAst = preparedObj.PreparedAst
 				stmtText, stmtDB = preparedObj.StmtText, preparedObj.StmtDB
 				bindSQL, _ := plannercore.GetBindSQL4PlanCache(s, preparedObj)
+				limitOffsetAndCount, _ := plannercore.ExtractLimitFromAst(preparedAst.Stmt)
 				cacheKey, err = plannercore.NewPlanCacheKey(s.sessionVars, stmtText, stmtDB, preparedAst.SchemaVersion,
-					0, bindSQL)
+					0, bindSQL, limitOffsetAndCount)
 				if err != nil {
 					logutil.Logger(s.currentCtx).Warn("clean cached plan failed", zap.Error(err))
 					return