From 5a7de4b2dd9f146bf63ce6432d84aa94470e0792 Mon Sep 17 00:00:00 2001 From: qw4990 Date: Tue, 20 Dec 2022 15:30:10 +0800 Subject: [PATCH] fixup --- planner/core/plan_cache.go | 4 +++ planner/core/plan_cache_param.go | 11 ++++++-- planner/core/plan_cache_param_test.go | 5 ++-- planner/core/plan_cache_test.go | 40 +++++++++++++++++++++++++++ planner/core/plan_cache_utils.go | 4 +++ planner/optimize.go | 11 ++++++-- 6 files changed, 68 insertions(+), 7 deletions(-) diff --git a/planner/core/plan_cache.go b/planner/core/plan_cache.go index d9df658aaa46e..ab4eb4e4912ab 100644 --- a/planner/core/plan_cache.go +++ b/planner/core/plan_cache.go @@ -115,6 +115,10 @@ func planCachePreprocess(ctx context.Context, sctx sessionctx.Context, isNonPrep func GetPlanFromSessionPlanCache(ctx context.Context, sctx sessionctx.Context, isNonPrepared bool, is infoschema.InfoSchema, stmt *PlanCacheStmt, params []expression.Expression) (plan Plan, names []*types.FieldName, err error) { + if v := ctx.Value("____GetPlanFromSessionPlanCacheErr"); v != nil { // for testing + return nil, nil, errors.New("____GetPlanFromSessionPlanCacheErr") + } + if err := planCachePreprocess(ctx, sctx, isNonPrepared, is, stmt, params); err != nil { return nil, nil, err } diff --git a/planner/core/plan_cache_param.go b/planner/core/plan_cache_param.go index 7c79b2a6416a0..9094edec621c0 100644 --- a/planner/core/plan_cache_param.go +++ b/planner/core/plan_cache_param.go @@ -15,6 +15,7 @@ package core import ( + "context" "errors" "strings" "sync" @@ -70,7 +71,7 @@ func (pr *paramReplacer) Reset() { pr.params = nil } // ParameterizeAST parameterizes this StmtNode. // e.g. `select * from t where a<10 and b<23` --> `select * from t where a `select * from t where a<10 and b<23`. -func RestoreASTWithParams(_ sessionctx.Context, stmt ast.StmtNode, params []*driver.ValueExpr) error { +func RestoreASTWithParams(ctx context.Context, _ sessionctx.Context, stmt ast.StmtNode, params []*driver.ValueExpr) error { + if v := ctx.Value("____RestoreASTWithParamsErr"); v != nil { + return errors.New("____RestoreASTWithParamsErr") + } + pr := paramRestorerPool.Get().(*paramRestorer) defer func() { pr.Reset() diff --git a/planner/core/plan_cache_param_test.go b/planner/core/plan_cache_param_test.go index 5c65b89767a60..ee4a8e9ae65c5 100644 --- a/planner/core/plan_cache_param_test.go +++ b/planner/core/plan_cache_param_test.go @@ -15,6 +15,7 @@ package core import ( + "context" "strings" "testing" @@ -61,7 +62,7 @@ func TestParameterize(t *testing.T) { for _, c := range cases { stmt, err := parser.New().ParseOneStmt(c.sql, "", "") require.Nil(t, err) - paramSQL, params, err := ParameterizeAST(sctx, stmt) + paramSQL, params, err := ParameterizeAST(context.Background(), sctx, stmt) require.Nil(t, err) require.Equal(t, c.paramSQL, paramSQL) require.Equal(t, len(c.params), len(params)) @@ -69,7 +70,7 @@ func TestParameterize(t *testing.T) { require.Equal(t, c.params[i], params[i].Datum.GetValue()) } - err = RestoreASTWithParams(sctx, stmt, params) + err = RestoreASTWithParams(context.Background(), sctx, stmt, params) require.Nil(t, err) var buf strings.Builder rCtx := format.NewRestoreCtx(format.DefaultRestoreFlags, &buf) diff --git a/planner/core/plan_cache_test.go b/planner/core/plan_cache_test.go index 5c3c7c45b9702..7a4ac860d8593 100644 --- a/planner/core/plan_cache_test.go +++ b/planner/core/plan_cache_test.go @@ -15,6 +15,7 @@ package core_test import ( + "context" "errors" "fmt" "math/rand" @@ -112,6 +113,45 @@ func TestNonPreparedPlanCacheWithExplain(t *testing.T) { tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1")) } +func TestNonPreparedPlanCacheFallback(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec(`use test`) + tk.MustExec(`create table t (a int)`) + for i := 0; i < 5; i++ { + tk.MustExec(fmt.Sprintf("insert into t values (%v)", i)) + } + tk.MustExec("set tidb_enable_non_prepared_plan_cache=1") + + // inject a fault to GeneratePlanCacheStmtWithAST + ctx := context.WithValue(context.Background(), "____GeneratePlanCacheStmtWithASTErr", struct{}{}) + tk.MustQueryWithContext(ctx, "select * from t where a in (1, 2)").Sort().Check(testkit.Rows("1", "2")) + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0")) // cannot generate PlanCacheStmt + tk.MustQueryWithContext(ctx, "select * from t where a in (1, 3)").Sort().Check(testkit.Rows("1", "3")) + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0")) // cannot generate PlanCacheStmt + tk.MustQuery("select * from t where a in (1, 2)").Sort().Check(testkit.Rows("1", "2")) + tk.MustQuery("select * from t where a in (1, 3)").Sort().Check(testkit.Rows("1", "3")) + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1")) // no error + + // inject a fault to GetPlanFromSessionPlanCache + tk.MustQuery("select * from t where a=1").Check(testkit.Rows("1")) // cache this plan + tk.MustQuery("select * from t where a=2").Check(testkit.Rows("2")) // plan from cache + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1")) + ctx = context.WithValue(context.Background(), "____GetPlanFromSessionPlanCacheErr", struct{}{}) + tk.MustQueryWithContext(ctx, "select * from t where a=3").Check(testkit.Rows("3")) + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0")) // fallback to the normal opt-path + tk.MustQueryWithContext(ctx, "select * from t where a=4").Check(testkit.Rows("4")) + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("0")) // fallback to the normal opt-path + tk.MustQueryWithContext(context.Background(), "select * from t where a=0").Check(testkit.Rows("0")) + tk.MustQuery("select @@last_plan_from_cache").Check(testkit.Rows("1")) // use the cached plan if no error + + // inject a fault to RestoreASTWithParams + ctx = context.WithValue(context.Background(), "____GetPlanFromSessionPlanCacheErr", struct{}{}) + ctx = context.WithValue(ctx, "____RestoreASTWithParamsErr", struct{}{}) + _, err := tk.ExecWithContext(ctx, "select * from t where a=1") + require.NotNil(t, err) +} + func TestNonPreparedPlanCacheBasically(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) diff --git a/planner/core/plan_cache_utils.go b/planner/core/plan_cache_utils.go index 1cb66d3cdeb6a..8dc867316207d 100644 --- a/planner/core/plan_cache_utils.go +++ b/planner/core/plan_cache_utils.go @@ -67,6 +67,10 @@ func (e *paramMarkerExtractor) Leave(in ast.Node) (ast.Node, bool) { // paramSQL is the corresponding parameterized sql like 'select * from t where a?'. // paramStmt is the Node of paramSQL. func GeneratePlanCacheStmtWithAST(ctx context.Context, sctx sessionctx.Context, paramSQL string, paramStmt ast.StmtNode) (*PlanCacheStmt, Plan, int, error) { + if v := ctx.Value("____GeneratePlanCacheStmtWithASTErr"); v != nil { // for testing + return nil, nil, 0, errors.New("____GeneratePlanCacheStmtWithASTErr") + } + vars := sctx.GetSessionVars() var extractor paramMarkerExtractor paramStmt.Accept(&extractor) diff --git a/planner/optimize.go b/planner/optimize.go index 0f37a59788a5c..b627307fb4fdc 100644 --- a/planner/optimize.go +++ b/planner/optimize.go @@ -74,15 +74,22 @@ func matchSQLBinding(sctx sessionctx.Context, stmtNode ast.StmtNode) (bindRecord } // getPlanFromNonPreparedPlanCache tries to get an available cached plan from the NonPrepared Plan Cache for this stmt. -func getPlanFromNonPreparedPlanCache(ctx context.Context, sctx sessionctx.Context, stmt ast.StmtNode, is infoschema.InfoSchema) (core.Plan, types.NameSlice, bool, error) { +func getPlanFromNonPreparedPlanCache(ctx context.Context, sctx sessionctx.Context, stmt ast.StmtNode, is infoschema.InfoSchema) (p core.Plan, ns types.NameSlice, ok bool, err error) { if sctx.GetSessionVars().StmtCtx.InPreparedPlanBuilding || // already in cached plan rebuilding phase !core.NonPreparedPlanCacheableWithCtx(sctx, stmt, is) { return nil, nil, false, nil } - paramSQL, params, err := core.ParameterizeAST(sctx, stmt) + paramSQL, params, err := core.ParameterizeAST(ctx, sctx, stmt) if err != nil { return nil, nil, false, err } + defer func() { + if err != nil { + // keep the stmt unchanged if err so that it can fallback to the normal optimization path. + // TODO: add metrics + err = core.RestoreASTWithParams(ctx, sctx, stmt, params) + } + }() val := sctx.GetSessionVars().GetNonPreparedPlanCacheStmt(paramSQL) if val == nil { cachedStmt, _, _, err := core.GeneratePlanCacheStmtWithAST(ctx, sctx, paramSQL, stmt)