re-order the initialization of resource group runaway checker

Signed-off-by: Yang Keao <[email protected]>
pingcap · Mar 14, 2024 · 48c8112 · 48c8112
1 parent ab4b5cd
commit 48c8112
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 16 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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()
 	})