From 48c8112a636b0a49965fc94c76dc10825bb39c52 Mon Sep 17 00:00:00 2001 From: Yang Keao Date: Thu, 14 Mar 2024 16:06:50 +0800 Subject: [PATCH] re-order the initialization of resource group runaway checker Signed-off-by: Yang Keao --- pkg/distsql/context/context.go | 5 ++--- pkg/distsql/request_builder.go | 4 +--- pkg/executor/adapter.go | 13 +++++++------ pkg/session/session.go | 4 ++-- pkg/sessionctx/stmtctx/stmtctx.go | 7 +++++-- 5 files changed, 17 insertions(+), 16 deletions(-) diff --git a/pkg/distsql/context/context.go b/pkg/distsql/context/context.go index 9547a45d9fddb..beb66cf968af0 100644 --- a/pkg/distsql/context/context.go +++ b/pkg/distsql/context/context.go @@ -79,9 +79,8 @@ type DistSQLContext struct { StoreBatchSize int ResourceGroupName string LoadBasedReplicaReadThreshold time.Duration - // TODO: remove this double pointer. This is a double-pointer because the `RunawayChecker` in stmtctx will be modified after building executors - RunawayChecker **resourcegroup.RunawayChecker - TiKVClientReadTimeout uint64 + RunawayChecker *resourcegroup.RunawayChecker + TiKVClientReadTimeout uint64 ExecDetails *execdetails.SyncExecDetails } diff --git a/pkg/distsql/request_builder.go b/pkg/distsql/request_builder.go index 7d3ca4d40221e..a3e6264e9e417 100644 --- a/pkg/distsql/request_builder.go +++ b/pkg/distsql/request_builder.go @@ -310,9 +310,7 @@ func (builder *RequestBuilder) SetFromSessionVars(dctx *distsqlctx.DistSQLContex builder.StoreBatchSize = dctx.StoreBatchSize builder.Request.ResourceGroupName = dctx.ResourceGroupName builder.Request.StoreBusyThreshold = dctx.LoadBasedReplicaReadThreshold - if dctx.RunawayChecker != nil { - builder.Request.RunawayChecker = *dctx.RunawayChecker - } + builder.Request.RunawayChecker = dctx.RunawayChecker builder.Request.TiKVClientReadTimeout = dctx.TiKVClientReadTimeout return builder } diff --git a/pkg/executor/adapter.go b/pkg/executor/adapter.go index cd45db602ed7b..b90501e6f3db0 100644 --- a/pkg/executor/adapter.go +++ b/pkg/executor/adapter.go @@ -504,12 +504,6 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) { sctx.GetSessionVars().MemTracker.SetBytesLimit(sctx.GetSessionVars().StmtCtx.MemQuotaQuery) } - e, err := a.buildExecutor() - if err != nil { - return nil, err - } - // ExecuteExec will rewrite `a.Plan`, so set plan label should be executed after `a.buildExecutor`. - ctx = a.observeStmtBeginForTopSQL(ctx) if variable.EnableResourceControl.Load() && domain.GetDomain(sctx).RunawayManager() != nil { stmtCtx := sctx.GetSessionVars().StmtCtx _, planDigest := GetPlanDigest(stmtCtx) @@ -520,6 +514,13 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) { } } + e, err := a.buildExecutor() + if err != nil { + return nil, err + } + // ExecuteExec will rewrite `a.Plan`, so set plan label should be executed after `a.buildExecutor`. + ctx = a.observeStmtBeginForTopSQL(ctx) + cmd32 := atomic.LoadUint32(&sctx.GetSessionVars().CommandValue) cmd := byte(cmd32) var pi processinfoSetter diff --git a/pkg/session/session.go b/pkg/session/session.go index 326c12583c3b5..f01234d461152 100644 --- a/pkg/session/session.go +++ b/pkg/session/session.go @@ -2626,7 +2626,7 @@ func (s *session) GetDistSQLCtx() *distsqlctx.DistSQLContext { vars := s.GetSessionVars() sc := vars.StmtCtx - dctx := sc.GetOrInitDistSQLFromCache(func() any { + dctx := sc.GetOrInitDistSQLFromCache(func() *distsqlctx.DistSQLContext { return &distsqlctx.DistSQLContext{ AppendWarning: sc.AppendWarning, InRestrictedSQL: sc.InRestrictedSQL, @@ -2668,7 +2668,7 @@ func (s *session) GetDistSQLCtx() *distsqlctx.DistSQLContext { StoreBatchSize: vars.StoreBatchSize, ResourceGroupName: sc.ResourceGroupName, LoadBasedReplicaReadThreshold: vars.LoadBasedReplicaReadThreshold, - RunawayChecker: &sc.RunawayChecker, + RunawayChecker: sc.RunawayChecker, TiKVClientReadTimeout: vars.GetTiKVClientReadTimeout(), ExecDetails: &sc.SyncExecDetails, diff --git a/pkg/sessionctx/stmtctx/stmtctx.go b/pkg/sessionctx/stmtctx/stmtctx.go index 5a035898cbca1..00eb339e14fcc 100644 --- a/pkg/sessionctx/stmtctx/stmtctx.go +++ b/pkg/sessionctx/stmtctx/stmtctx.go @@ -28,6 +28,7 @@ import ( "time" "github.com/pingcap/errors" + distsqlctx "github.com/pingcap/tidb/pkg/distsql/context" "github.com/pingcap/tidb/pkg/domain/resourcegroup" "github.com/pingcap/tidb/pkg/errctx" "github.com/pingcap/tidb/pkg/parser" @@ -167,11 +168,13 @@ type StatementContext struct { // errCtx is used to indicate how to handle the errors errCtx errctx.Context + DistSQLContext distsqlctx.DistSQLContext + // distSQLCtxCache is used to persist all variables and tools needed by the `distsql` // this cache is set on `StatementContext` because it has to be updated after each statement. distSQLCtxCache struct { init sync.Once - dctx any + dctx *distsqlctx.DistSQLContext } // Set the following variables before execution @@ -1246,7 +1249,7 @@ func (sc *StatementContext) TypeCtxOrDefault() types.Context { // GetOrInitDistSQLFromCache returns the `DistSQLContext` inside cache. If it didn't exist, return a new one created by // the `create` function. It uses the `any` to avoid cycle dependency. -func (sc *StatementContext) GetOrInitDistSQLFromCache(create func() any) any { +func (sc *StatementContext) GetOrInitDistSQLFromCache(create func() *distsqlctx.DistSQLContext) any { sc.distSQLCtxCache.init.Do(func() { sc.distSQLCtxCache.dctx = create() })