From 538741fdedcf26bc2018a389254b9c560fe537c0 Mon Sep 17 00:00:00 2001 From: you06 Date: Fri, 4 Aug 2023 14:17:12 +0800 Subject: [PATCH 1/2] store the hints of session variable Signed-off-by: you06 lint Signed-off-by: you06 clone stmt hints & fix fast path Signed-off-by: you06 add test for clone stmt hints Signed-off-by: you06 lint & bazel Signed-off-by: you06 address comment Signed-off-by: you06 fix lint Signed-off-by: you06 --- planner/core/plan_cache.go | 11 ++++- planner/core/plan_cache_utils.go | 6 ++- planner/core/point_get_plan.go | 2 + session/test/vars/BUILD.bazel | 3 +- session/test/vars/vars_test.go | 68 ++++++++++++++++++++++++++++++ sessionctx/stmtctx/BUILD.bazel | 2 +- sessionctx/stmtctx/stmtctx.go | 38 ++++++++++++++++- sessionctx/stmtctx/stmtctx_test.go | 23 ++++++++++ 8 files changed, 147 insertions(+), 6 deletions(-) diff --git a/planner/core/plan_cache.go b/planner/core/plan_cache.go index 8e7477a8c34bd..8623f59173c12 100644 --- a/planner/core/plan_cache.go +++ b/planner/core/plan_cache.go @@ -245,6 +245,9 @@ func getCachedPointPlan(stmt *ast.Prepared, sessVars *variable.SessionVars, stmt } sessVars.FoundInPlanCache = true stmtCtx.PointExec = true + if pointGetPlan, ok := plan.(*PointGetPlan); ok && pointGetPlan != nil && pointGetPlan.stmtHints != nil { + sessVars.StmtCtx.StmtHints = *pointGetPlan.stmtHints + } return plan, names, true, nil } @@ -285,6 +288,7 @@ func getCachedPlan(sctx sessionctx.Context, isNonPrepared bool, cacheKey kvcache core_metrics.GetPlanCacheHitCounter(isNonPrepared).Inc() } stmtCtx.SetPlanDigest(stmt.NormalizedPlan, stmt.PlanDigest) + stmtCtx.StmtHints = *cachedVal.stmtHints return cachedVal.Plan, cachedVal.OutPutNames, true, nil } @@ -327,7 +331,7 @@ func generateNewPlan(ctx context.Context, sctx sessionctx.Context, isNonPrepared } sessVars.IsolationReadEngines[kv.TiFlash] = struct{}{} } - cached := NewPlanCacheValue(p, names, stmtCtx.TblInfo2UnionScan, matchOpts) + cached := NewPlanCacheValue(p, names, stmtCtx.TblInfo2UnionScan, matchOpts, &stmtCtx.StmtHints) stmt.NormalizedPlan, stmt.PlanDigest = NormalizePlan(p) stmtCtx.SetPlan(p) stmtCtx.SetPlanDigest(stmt.NormalizedPlan, stmt.PlanDigest) @@ -757,12 +761,15 @@ func tryCachePointPlan(_ context.Context, sctx sessionctx.Context, names types.NameSlice ) - if _, _ok := p.(*PointGetPlan); _ok { + if plan, _ok := p.(*PointGetPlan); _ok { ok, err = IsPointGetWithPKOrUniqueKeyByAutoCommit(sctx, p) names = p.OutputNames() if err != nil { return err } + if ok { + plan.stmtHints = sctx.GetSessionVars().StmtCtx.StmtHints.Clone() + } } if ok { diff --git a/planner/core/plan_cache_utils.go b/planner/core/plan_cache_utils.go index 2bb3ac62a4316..ca5683b317603 100644 --- a/planner/core/plan_cache_utils.go +++ b/planner/core/plan_cache_utils.go @@ -32,6 +32,7 @@ import ( "github.com/pingcap/tidb/planner/util" "github.com/pingcap/tidb/planner/util/fixcontrol" "github.com/pingcap/tidb/sessionctx" + "github.com/pingcap/tidb/sessionctx/stmtctx" "github.com/pingcap/tidb/sessionctx/variable" "github.com/pingcap/tidb/statistics" "github.com/pingcap/tidb/types" @@ -337,6 +338,8 @@ type PlanCacheValue struct { // matchOpts stores some fields help to choose a suitable plan matchOpts *utilpc.PlanCacheMatchOpts + // stmtHints stores the hints which set session variables, because the hints won't be processed using cached plan. + stmtHints *stmtctx.StmtHints } // unKnownMemoryUsage represent the memory usage of uncounted structure, maybe need implement later @@ -383,7 +386,7 @@ func (v *PlanCacheValue) MemoryUsage() (sum int64) { // NewPlanCacheValue creates a SQLCacheValue. func NewPlanCacheValue(plan Plan, names []*types.FieldName, srcMap map[*model.TableInfo]bool, - matchOpts *utilpc.PlanCacheMatchOpts) *PlanCacheValue { + matchOpts *utilpc.PlanCacheMatchOpts, stmtHints *stmtctx.StmtHints) *PlanCacheValue { dstMap := make(map[*model.TableInfo]bool) for k, v := range srcMap { dstMap[k] = v @@ -397,6 +400,7 @@ func NewPlanCacheValue(plan Plan, names []*types.FieldName, srcMap map[*model.Ta OutPutNames: names, TblInfo2UnionScan: dstMap, matchOpts: matchOpts, + stmtHints: stmtHints.Clone(), } } diff --git a/planner/core/point_get_plan.go b/planner/core/point_get_plan.go index bf6190575858f..0e388d3cbea7b 100644 --- a/planner/core/point_get_plan.go +++ b/planner/core/point_get_plan.go @@ -97,6 +97,8 @@ type PointGetPlan struct { // probeParents records the IndexJoins and Applys with this operator in their inner children. // Please see comments in PhysicalPlan for details. probeParents []PhysicalPlan + // stmtHints should restore in executing context. + stmtHints *stmtctx.StmtHints } func (p *PointGetPlan) getEstRowCountForDisplay() float64 { diff --git a/session/test/vars/BUILD.bazel b/session/test/vars/BUILD.bazel index f93be7371b697..a53c88a43ed62 100644 --- a/session/test/vars/BUILD.bazel +++ b/session/test/vars/BUILD.bazel @@ -8,7 +8,7 @@ go_test( "vars_test.go", ], flaky = True, - shard_count = 12, + shard_count = 13, deps = [ "//config", "//domain", @@ -16,6 +16,7 @@ go_test( "//kv", "//parser/mysql", "//parser/terror", + "//sessionctx/stmtctx", "//sessionctx/variable", "//testkit", "//testkit/testmain", diff --git a/session/test/vars/vars_test.go b/session/test/vars/vars_test.go index 29e51ec787577..d533c175a880f 100644 --- a/session/test/vars/vars_test.go +++ b/session/test/vars/vars_test.go @@ -27,6 +27,7 @@ import ( tikv "github.com/pingcap/tidb/kv" "github.com/pingcap/tidb/parser/mysql" "github.com/pingcap/tidb/parser/terror" + "github.com/pingcap/tidb/sessionctx/stmtctx" "github.com/pingcap/tidb/sessionctx/variable" "github.com/pingcap/tidb/testkit" "github.com/stretchr/testify/require" @@ -615,3 +616,70 @@ func TestGetSysVariables(t *testing.T) { tk.MustExec("select @@local.performance_schema_max_mutex_classes") tk.MustGetErrMsg("select @@global.last_insert_id", "[variable:1238]Variable 'last_insert_id' is a SESSION variable") } + +func TestPrepareExecuteWithSQLHints(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + se := tk.Session() + se.SetConnectionID(1) + tk.MustExec("use test") + tk.MustExec("create table t(a int primary key)") + + type hintCheck struct { + hint string + check func(*stmtctx.StmtHints) + } + + hintChecks := []hintCheck{ + { + hint: "MEMORY_QUOTA(1024 MB)", + check: func(stmtHint *stmtctx.StmtHints) { + require.True(t, stmtHint.HasMemQuotaHint) + require.Equal(t, int64(1024*1024*1024), stmtHint.MemQuotaQuery) + }, + }, + { + hint: "READ_CONSISTENT_REPLICA()", + check: func(stmtHint *stmtctx.StmtHints) { + require.True(t, stmtHint.HasReplicaReadHint) + require.Equal(t, byte(tikv.ReplicaReadFollower), stmtHint.ReplicaRead) + }, + }, + { + hint: "MAX_EXECUTION_TIME(1000)", + check: func(stmtHint *stmtctx.StmtHints) { + require.True(t, stmtHint.HasMaxExecutionTime) + require.Equal(t, uint64(1000), stmtHint.MaxExecutionTime) + }, + }, + { + hint: "USE_TOJA(TRUE)", + check: func(stmtHint *stmtctx.StmtHints) { + require.True(t, stmtHint.HasAllowInSubqToJoinAndAggHint) + require.True(t, stmtHint.AllowInSubqToJoinAndAgg) + }, + }, + { + hint: "RESOURCE_GROUP(rg1)", + check: func(stmtHint *stmtctx.StmtHints) { + require.True(t, stmtHint.HasResourceGroup) + require.Equal(t, "rg1", stmtHint.ResourceGroup) + }, + }, + } + + for i, check := range hintChecks { + // common path + tk.MustExec(fmt.Sprintf("prepare stmt%d from 'select /*+ %s */ * from t'", i, check.hint)) + for j := 0; j < 10; j++ { + tk.MustQuery(fmt.Sprintf("execute stmt%d", i)) + check.check(&tk.Session().GetSessionVars().StmtCtx.StmtHints) + } + // fast path + tk.MustExec(fmt.Sprintf("prepare fast%d from 'select /*+ %s */ * from t where a = 1'", i, check.hint)) + for j := 0; j < 10; j++ { + tk.MustQuery(fmt.Sprintf("execute fast%d", i)) + check.check(&tk.Session().GetSessionVars().StmtCtx.StmtHints) + } + } +} diff --git a/sessionctx/stmtctx/BUILD.bazel b/sessionctx/stmtctx/BUILD.bazel index b0b4239eab7ea..e16d9a44da20d 100644 --- a/sessionctx/stmtctx/BUILD.bazel +++ b/sessionctx/stmtctx/BUILD.bazel @@ -38,7 +38,7 @@ go_test( ], embed = [":stmtctx"], flaky = True, - shard_count = 5, + shard_count = 6, deps = [ "//kv", "//sessionctx/variable", diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go index 66e72d4cfccc1..a4872c2882b77 100644 --- a/sessionctx/stmtctx/stmtctx.go +++ b/sessionctx/stmtctx/stmtctx.go @@ -421,7 +421,6 @@ type StatementContext struct { type StmtHints struct { // Hint Information MemQuotaQuery int64 - ApplyCacheCapacity int64 MaxExecutionTime uint64 TidbKvReadTimeout uint64 ReplicaRead byte @@ -454,6 +453,43 @@ func (sh *StmtHints) TaskMapNeedBackUp() bool { return sh.ForceNthPlan != -1 } +// Clone the StmtHints struct and returns the pointer of the new one. +func (sh *StmtHints) Clone() *StmtHints { + var ( + vars map[string]string + tableHints []*ast.TableOptimizerHint + ) + if len(sh.SetVars) > 0 { + vars = make(map[string]string, len(sh.SetVars)) + for k, v := range sh.SetVars { + vars[k] = v + } + } + if len(sh.OriginalTableHints) > 0 { + tableHints = make([]*ast.TableOptimizerHint, len(sh.OriginalTableHints)) + copy(tableHints, sh.OriginalTableHints) + } + return &StmtHints{ + MemQuotaQuery: sh.MemQuotaQuery, + MaxExecutionTime: sh.MaxExecutionTime, + ReplicaRead: sh.ReplicaRead, + AllowInSubqToJoinAndAgg: sh.AllowInSubqToJoinAndAgg, + NoIndexMergeHint: sh.NoIndexMergeHint, + StraightJoinOrder: sh.StraightJoinOrder, + EnableCascadesPlanner: sh.EnableCascadesPlanner, + ForceNthPlan: sh.ForceNthPlan, + ResourceGroup: sh.ResourceGroup, + HasAllowInSubqToJoinAndAggHint: sh.HasAllowInSubqToJoinAndAggHint, + HasMemQuotaHint: sh.HasMemQuotaHint, + HasReplicaReadHint: sh.HasReplicaReadHint, + HasMaxExecutionTime: sh.HasMaxExecutionTime, + HasEnableCascadesPlannerHint: sh.HasEnableCascadesPlannerHint, + HasResourceGroup: sh.HasResourceGroup, + SetVars: vars, + OriginalTableHints: tableHints, + } +} + // StmtCacheKey represents the key type in the StmtCache. type StmtCacheKey int diff --git a/sessionctx/stmtctx/stmtctx_test.go b/sessionctx/stmtctx/stmtctx_test.go index 461168ee4b607..9a3951278befc 100644 --- a/sessionctx/stmtctx/stmtctx_test.go +++ b/sessionctx/stmtctx/stmtctx_test.go @@ -19,6 +19,7 @@ import ( "encoding/json" "fmt" "math/rand" + "reflect" "sort" "testing" "time" @@ -273,3 +274,25 @@ func TestApproxRuntimeInfo(t *testing.T) { require.Equal(t, d.TotBackoffTime[backoff], timeSum) } } + +func TestStmtHintsClone(t *testing.T) { + hints := stmtctx.StmtHints{} + value := reflect.ValueOf(&hints).Elem() + for i := 0; i < value.NumField(); i++ { + field := value.Field(i) + switch field.Kind() { + case reflect.Int, reflect.Int32, reflect.Int64: + field.SetInt(1) + case reflect.Uint, reflect.Uint32, reflect.Uint64: + field.SetUint(1) + case reflect.Uint8: // byte + field.SetUint(1) + case reflect.Bool: + field.SetBool(true) + case reflect.String: + field.SetString("test") + default: + } + } + require.Equal(t, hints, *hints.Clone()) +} From b2a10a46d6d65c55a257328838037dc12e8e2792 Mon Sep 17 00:00:00 2001 From: you06 Date: Mon, 14 Aug 2023 11:56:07 +0800 Subject: [PATCH 2/2] clone kv timeout Signed-off-by: you06 --- sessionctx/stmtctx/stmtctx.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sessionctx/stmtctx/stmtctx.go b/sessionctx/stmtctx/stmtctx.go index a4872c2882b77..5e3ead54d0640 100644 --- a/sessionctx/stmtctx/stmtctx.go +++ b/sessionctx/stmtctx/stmtctx.go @@ -472,6 +472,7 @@ func (sh *StmtHints) Clone() *StmtHints { return &StmtHints{ MemQuotaQuery: sh.MemQuotaQuery, MaxExecutionTime: sh.MaxExecutionTime, + TidbKvReadTimeout: sh.TidbKvReadTimeout, ReplicaRead: sh.ReplicaRead, AllowInSubqToJoinAndAgg: sh.AllowInSubqToJoinAndAgg, NoIndexMergeHint: sh.NoIndexMergeHint, @@ -483,6 +484,7 @@ func (sh *StmtHints) Clone() *StmtHints { HasMemQuotaHint: sh.HasMemQuotaHint, HasReplicaReadHint: sh.HasReplicaReadHint, HasMaxExecutionTime: sh.HasMaxExecutionTime, + HasTidbKvReadTimeout: sh.HasTidbKvReadTimeout, HasEnableCascadesPlannerHint: sh.HasEnableCascadesPlannerHint, HasResourceGroup: sh.HasResourceGroup, SetVars: vars,