session: add resource group name in stmt context

pkg/ddl/backfilling_scheduler.go

@@ -174,8 +174,7 @@ func initSessCtx(
 		WithIgnoreZeroInDate(!sqlMode.HasStrictMode() || sqlMode.HasAllowInvalidDatesMode()).
 		WithCastTimeToYearThroughConcat(true)
 	sessCtx.GetSessionVars().StmtCtx.SetTypeFlags(typeFlags)
-
-	sessCtx.GetSessionVars().ResourceGroupName = resGroupName
+	sessCtx.GetSessionVars().StmtCtx.ResourceGroupName = resGroupName
 
 	// Prevent initializing the mock context in the workers concurrently.
 	// For details, see https://github.com/pingcap/tidb/issues/40879.
@@ -199,7 +198,7 @@ func restoreSessCtx(sessCtx sessionctx.Context) func(sessCtx sessionctx.Context)
 	overflowAsWarn := sv.StmtCtx.OverflowAsWarning
 	dividedZeroAsWarn := sv.StmtCtx.DividedByZeroAsWarning
 	typeFlags := sv.StmtCtx.TypeFlags()
-	resGroupName := sv.ResourceGroupName
+	resGroupName := sv.StmtCtx.ResourceGroupName
 	return func(usedSessCtx sessionctx.Context) {
 		uv := usedSessCtx.GetSessionVars()
 		uv.RowEncoder.Enable = rowEncoder
@@ -209,7 +208,7 @@ func restoreSessCtx(sessCtx sessionctx.Context) func(sessCtx sessionctx.Context)
 		uv.StmtCtx.OverflowAsWarning = overflowAsWarn
 		uv.StmtCtx.DividedByZeroAsWarning = dividedZeroAsWarn
 		uv.StmtCtx.SetTypeFlags(typeFlags)
-		uv.ResourceGroupName = resGroupName
+		uv.StmtCtx.ResourceGroupName = resGroupName
 	}
 }
 

pkg/ddl/ddl_api.go

@@ -9031,7 +9031,7 @@ func NewDDLReorgMeta(ctx sessionctx.Context) *model.DDLReorgMeta {
 		Warnings:          make(map[errors.ErrorID]*terror.Error),
 		WarningsCount:     make(map[errors.ErrorID]int64),
 		Location:          &model.TimeZoneLocation{Name: tzName, Offset: tzOffset},
-		ResourceGroupName: ctx.GetSessionVars().ResourceGroupName,
+		ResourceGroupName: ctx.GetSessionVars().StmtCtx.ResourceGroupName,
 		Version:           model.CurrentReorgMetaVersion,
 	}
 }

pkg/ddl/tests/resourcegroup/resource_group_test.go

@@ -492,8 +492,10 @@ func TestBindHints(t *testing.T) {
 	tk.MustExec("create global binding for select * from t using select /*+ resource_group(rg1) */ * from t")
 	tk.MustQuery("select * from t")
 	re.Equal("rg1", tk.Session().GetSessionVars().StmtCtx.ResourceGroup)
+	re.Equal("rg1", tk.Session().GetSessionVars().StmtCtx.ResourceGroupName)
 	re.Equal("default", tk.Session().GetSessionVars().ResourceGroupName)
 	tk.MustQuery("select a, b from t")
 	re.Equal("", tk.Session().GetSessionVars().StmtCtx.ResourceGroup)
+	re.Equal("default", tk.Session().GetSessionVars().StmtCtx.ResourceGroupName)
 	re.Equal("default", tk.Session().GetSessionVars().ResourceGroupName)
 }

pkg/distsql/request_builder.go

@@ -309,7 +309,7 @@ func (builder *RequestBuilder) SetFromSessionVars(sv *variable.SessionVars) *Req
 	builder.RequestSource.RequestSourceType = sv.RequestSourceType
 	builder.RequestSource.ExplicitRequestSourceType = sv.ExplicitRequestSourceType
 	builder.StoreBatchSize = sv.StoreBatchSize
-	builder.Request.ResourceGroupName = sv.ResourceGroupName
+	builder.Request.ResourceGroupName = sv.StmtCtx.ResourceGroupName
 	builder.Request.StoreBusyThreshold = sv.LoadBasedReplicaReadThreshold
 	builder.Request.RunawayChecker = sv.StmtCtx.RunawayChecker
 	builder.Request.TiKVClientReadTimeout = sv.GetTiKVClientReadTimeout()

pkg/distsql/request_builder_test.go

@@ -579,7 +579,7 @@ func TestRequestBuilder7(t *testing.T) {
 
 func TestRequestBuilder8(t *testing.T) {
 	sv := variable.NewSessionVars(nil)
-	sv.ResourceGroupName = "test"
+	sv.StmtCtx.ResourceGroupName = "test"
 	actual, err := (&RequestBuilder{}).
 		SetFromSessionVars(sv).
 		Build()

pkg/executor/adapter.go

@@ -557,7 +557,7 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) {
 		stmtCtx := sctx.GetSessionVars().StmtCtx
 		_, planDigest := GetPlanDigest(stmtCtx)
 		_, digest := stmtCtx.SQLDigest()
-		stmtCtx.RunawayChecker = domain.GetDomain(sctx).RunawayManager().DeriveChecker(sctx.GetSessionVars().ResourceGroupName, stmtCtx.OriginalSQL, digest.String(), planDigest.String())
+		stmtCtx.RunawayChecker = domain.GetDomain(sctx).RunawayManager().DeriveChecker(sctx.GetSessionVars().StmtCtx.ResourceGroupName, stmtCtx.OriginalSQL, digest.String(), planDigest.String())
 		if err := stmtCtx.RunawayChecker.BeforeExecutor(); err != nil {
 			return nil, err
 		}
@@ -1620,7 +1620,7 @@ func (a *ExecStmt) LogSlowQuery(txnTS uint64, succ bool, hasMoreResults bool) {
 		UsedStats:         stmtCtx.GetUsedStatsInfo(false),
 		IsSyncStatsFailed: stmtCtx.IsSyncStatsFailed,
 		Warnings:          collectWarningsForSlowLog(stmtCtx),
-		ResourceGroupName: sessVars.ResourceGroupName,
+		ResourceGroupName: sessVars.StmtCtx.ResourceGroupName,
 		RRU:               ruDetails.RRU(),
 		WRU:               ruDetails.WRU(),
 		WaitRUDuration:    ruDetails.RUWaitDuration(),

pkg/executor/analyze_col.go

@@ -122,7 +122,7 @@ func (e *AnalyzeColumnsExec) buildResp(ranges []*ranger.Range) (distsql.SelectRe
 		SetKeepOrder(true).
 		SetConcurrency(e.concurrency).
 		SetMemTracker(e.memTracker).
-		SetResourceGroupName(e.ctx.GetSessionVars().ResourceGroupName).
+		SetResourceGroupName(e.ctx.GetSessionVars().StmtCtx.ResourceGroupName).
 		SetExplicitRequestSourceType(e.ctx.GetSessionVars().ExplicitRequestSourceType).
 		Build()
 	if err != nil {

pkg/executor/analyze_idx.go

@@ -159,7 +159,7 @@ func (e *AnalyzeIndexExec) fetchAnalyzeResult(ranges []*ranger.Range, isNullRang
 		SetStartTS(startTS).
 		SetKeepOrder(true).
 		SetConcurrency(e.concurrency).
-		SetResourceGroupName(e.ctx.GetSessionVars().ResourceGroupName).
+		SetResourceGroupName(e.ctx.GetSessionVars().StmtCtx.ResourceGroupName).
 		SetExplicitRequestSourceType(e.ctx.GetSessionVars().ExplicitRequestSourceType).
 		Build()
 	if err != nil {

pkg/executor/builder.go

@@ -1926,7 +1926,7 @@ func (b *executorBuilder) getSnapshot() (kv.Snapshot, error) {
 	snapshot.SetOption(kv.ReadReplicaScope, b.readReplicaScope)
 	snapshot.SetOption(kv.TaskID, sessVars.StmtCtx.TaskID)
 	snapshot.SetOption(kv.TiKVClientReadTimeout, sessVars.GetTiKVClientReadTimeout())
-	snapshot.SetOption(kv.ResourceGroupName, sessVars.ResourceGroupName)
+	snapshot.SetOption(kv.ResourceGroupName, sessVars.StmtCtx.ResourceGroupName)
 	snapshot.SetOption(kv.ExplicitRequestSourceType, sessVars.ExplicitRequestSourceType)
 
 	if replicaReadType.IsClosestRead() && b.readReplicaScope != kv.GlobalTxnScope {

pkg/executor/checksum.go

@@ -247,7 +247,7 @@ func (c *checksumContext) buildTableRequest(ctx sessionctx.Context, tableID int6
 		SetChecksumRequest(checksum).
 		SetStartTS(c.StartTs).
 		SetConcurrency(ctx.GetSessionVars().DistSQLScanConcurrency()).
-		SetResourceGroupName(ctx.GetSessionVars().ResourceGroupName).
+		SetResourceGroupName(ctx.GetSessionVars().StmtCtx.ResourceGroupName).
 		SetExplicitRequestSourceType(ctx.GetSessionVars().ExplicitRequestSourceType).
 		Build()
 }
@@ -266,7 +266,7 @@ func (c *checksumContext) buildIndexRequest(ctx sessionctx.Context, tableID int6
 		SetChecksumRequest(checksum).
 		SetStartTS(c.StartTs).
 		SetConcurrency(ctx.GetSessionVars().DistSQLScanConcurrency()).
-		SetResourceGroupName(ctx.GetSessionVars().ResourceGroupName).
+		SetResourceGroupName(ctx.GetSessionVars().StmtCtx.ResourceGroupName).
 		SetExplicitRequestSourceType(ctx.GetSessionVars().ExplicitRequestSourceType).
 		Build()
 }

pkg/executor/executor.go

@@ -1989,6 +1989,7 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 	sc.OptimizerCETrace = nil
 	sc.IsSyncStatsFailed = false
 	sc.IsExplainAnalyzeDML = false
+	sc.ResourceGroupName = vars.ResourceGroupName
 	// Firstly we assume that UseDynamicPruneMode can be enabled according session variable, then we will check other conditions
 	// in PlanBuilder.buildDataSource
 	if ctx.GetSessionVars().IsDynamicPartitionPruneEnabled() {

pkg/executor/internal/mpp/local_mpp_coordinator.go

@@ -209,7 +209,7 @@ func (c *localMppCoordinator) appendMPPDispatchReq(pf *plannercore.Fragment) err
 			return errors.Trace(err)
 		}
 
-		rgName := c.sessionCtx.GetSessionVars().ResourceGroupName
+		rgName := c.sessionCtx.GetSessionVars().StmtCtx.ResourceGroupName
 		if !variable.EnableResourceControl.Load() {
 			rgName = ""
 		}

pkg/expression/builtin_info.go

@@ -32,6 +32,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/privilege"
 	"github.com/pingcap/tidb/pkg/sessionctx"
+	"github.com/pingcap/tidb/pkg/sessionctx/variable"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/chunk"
@@ -280,7 +281,20 @@ func (b *builtinCurrentResourceGroupSig) evalString(ctx sessionctx.Context, row
 	if data == nil {
 		return "", true, errors.Errorf("Missing session variable when eval builtin")
 	}
-	return data.ResourceGroupName, false, nil
+
+	return getHintResourceGroupName(data), false, nil
+}
+
+// get statement resource group name with hint in consideration
+// NOTE: because function `CURRENT_RESOURCE_GROUP()` maybe evaluated in optimizer
+// before we assign the hint value to StmtCtx.ResourceGroupName, so we have to
+// explicitly check the hint here.
+func getHintResourceGroupName(vars *variable.SessionVars) string {
+	groupName := vars.StmtCtx.ResourceGroupName
+	if vars.StmtCtx.HasResourceGroup {
+		groupName = vars.StmtCtx.StmtHints.ResourceGroup
+	}
+	return groupName
 }
 
 type userFunctionClass struct {

pkg/expression/builtin_info_test.go

@@ -99,7 +99,7 @@ func TestCurrentUser(t *testing.T) {
 func TestCurrentResourceGroup(t *testing.T) {
 	ctx := mock.NewContext()
 	sessionVars := ctx.GetSessionVars()
-	sessionVars.ResourceGroupName = "rg1"
+	sessionVars.StmtCtx.ResourceGroupName = "rg1"
 
 	fc := funcs[ast.CurrentResourceGroup]
 	f, err := fc.getFunction(ctx, nil)

pkg/expression/builtin_info_vec.go

@@ -130,8 +130,9 @@ func (b *builtinCurrentResourceGroupSig) vecEvalString(ctx sessionctx.Context, i
 	}
 	n := input.NumRows()
 	result.ReserveString(n)
+	groupName := getHintResourceGroupName(data)
 	for i := 0; i < n; i++ {
-		result.AppendString(data.ResourceGroupName)
+		result.AppendString(groupName)
 	}
 	return nil
 }

pkg/planner/optimize.go

@@ -188,11 +188,10 @@ func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 	}
 
 	defer func() {
-		// Override the resource group if necessary
-		// TODO: we didn't check the existence of the hinted resource group now to save the cost per query
+		// Override the resource group if the hint is set.
 		if retErr == nil && sessVars.StmtCtx.StmtHints.HasResourceGroup {
 			if variable.EnableResourceControl.Load() {
-				sessVars.ResourceGroupName = sessVars.StmtCtx.StmtHints.ResourceGroup
+				sessVars.StmtCtx.ResourceGroupName = sessVars.StmtCtx.StmtHints.ResourceGroup
 				// if we are in a txn, should update the txn resource name to let the txn
 				// commit with the hint resource group.
 				if txn, err := sctx.Txn(false); err == nil && txn != nil && txn.Valid() {

pkg/server/conn.go

@@ -1211,7 +1211,7 @@ func (cc *clientConn) addMetrics(cmd byte, startTime time.Time, err error) {
 
 	vars := cc.getCtx().GetSessionVars()
 	for _, dbName := range session.GetDBNames(vars) {
-		metrics.QueryDurationHistogram.WithLabelValues(sqlType, dbName, vars.ResourceGroupName).Observe(cost.Seconds())
+		metrics.QueryDurationHistogram.WithLabelValues(sqlType, dbName, vars.StmtCtx.ResourceGroupName).Observe(cost.Seconds())
 	}
 }
 

pkg/session/session.go

@@ -1512,7 +1512,7 @@ func (s *session) SetProcessInfo(sql string, t time.Time, command byte, maxExecu
 		IndexNames:            s.sessionVars.StmtCtx.IndexNames,
 		MaxExecutionTime:      maxExecutionTime,
 		RedactSQL:             s.sessionVars.EnableRedactLog,
-		ResourceGroupName:     s.sessionVars.ResourceGroupName,
+		ResourceGroupName:     s.sessionVars.StmtCtx.ResourceGroupName,
 		SessionAlias:          s.sessionVars.SessionAlias,
 	}
 	oldPi := s.ShowProcess()
@@ -1883,7 +1883,7 @@ func (s *session) ExecRestrictedStmt(ctx context.Context, stmtNode ast.StmtNode,
 
 	vars := se.GetSessionVars()
 	for _, dbName := range GetDBNames(vars) {
-		metrics.QueryDurationHistogram.WithLabelValues(metrics.LblInternal, dbName, vars.ResourceGroupName).Observe(time.Since(startTime).Seconds())
+		metrics.QueryDurationHistogram.WithLabelValues(metrics.LblInternal, dbName, vars.StmtCtx.ResourceGroupName).Observe(time.Since(startTime).Seconds())
 	}
 	return rows, rs.Fields(), err
 }
@@ -2059,7 +2059,7 @@ func (s *session) ExecRestrictedSQL(ctx context.Context, opts []sqlexec.OptionFu
 
 		vars := se.GetSessionVars()
 		for _, dbName := range GetDBNames(vars) {
-			metrics.QueryDurationHistogram.WithLabelValues(metrics.LblInternal, dbName, vars.ResourceGroupName).Observe(time.Since(startTime).Seconds())
+			metrics.QueryDurationHistogram.WithLabelValues(metrics.LblInternal, dbName, vars.StmtCtx.ResourceGroupName).Observe(time.Since(startTime).Seconds())
 		}
 		return rows, rs.Fields(), err
 	})
@@ -2168,28 +2168,19 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 	}
 	s.setRequestSource(ctx, stmtLabel, stmtNode)
 
-	// Backup the original resource group name since sql hint might change it during optimization
-	originalResourceGroup := s.GetSessionVars().ResourceGroupName
-
 	// Transform abstract syntax tree to a physical plan(stored in executor.ExecStmt).
 	compiler := executor.Compiler{Ctx: s}
 	stmt, err := compiler.Compile(ctx, stmtNode)
-	// session resource-group might be changed by query hint, ensure restore it back when
-	// the execution finished.
-	if sessVars.ResourceGroupName != originalResourceGroup {
+	// check if resource group hint is valid, can't do this in planner.Optimize because we can access
+	// infoschema there.
+	if sessVars.StmtCtx.ResourceGroupName != sessVars.ResourceGroupName {
 		// if target resource group doesn't exist, fallback to the origin resource group.
-		if _, ok := domain.GetDomain(s).InfoSchema().ResourceGroupByName(model.NewCIStr(sessVars.ResourceGroupName)); !ok {
-			logutil.Logger(ctx).Warn("Unknown resource group from hint", zap.String("name", sessVars.ResourceGroupName))
-			sessVars.ResourceGroupName = originalResourceGroup
-			// if we are in a txn, should also reset the txn resource group.
+		if _, ok := domain.GetDomain(s).InfoSchema().ResourceGroupByName(model.NewCIStr(sessVars.StmtCtx.ResourceGroupName)); !ok {
+			logutil.Logger(ctx).Warn("Unknown resource group from hint", zap.String("name", sessVars.StmtCtx.ResourceGroupName))
+			sessVars.StmtCtx.ResourceGroupName = sessVars.ResourceGroupName
 			if txn, err := s.Txn(false); err == nil && txn != nil && txn.Valid() {
-				kv.SetTxnResourceGroup(txn, originalResourceGroup)
+				kv.SetTxnResourceGroup(txn, sessVars.ResourceGroupName)
 			}
-		} else {
-			defer func() {
-				// Restore the resource group for the session
-				sessVars.ResourceGroupName = originalResourceGroup
-			}()
 		}
 	}
 	if err != nil {
@@ -2238,7 +2229,7 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 
 	// Observe the resource group query total counter if the resource control is enabled and the
 	// current session is attached with a resource group.
-	resourceGroupName := s.GetSessionVars().ResourceGroupName
+	resourceGroupName := s.GetSessionVars().StmtCtx.ResourceGroupName
 	if len(resourceGroupName) > 0 {
 		metrics.ResourceGroupQueryTotalCounter.WithLabelValues(resourceGroupName).Inc()
 	}

pkg/session/txn.go

@@ -307,7 +307,7 @@ func (txn *LazyTxn) changePendingToValid(ctx context.Context, sctx sessionctx.Co
 		txn.mu.TxnInfo.AllSQLDigests)
 
 	// set resource group name for kv request such as lock pessimistic keys.
-	kv.SetTxnResourceGroup(txn, sctx.GetSessionVars().ResourceGroupName)
+	kv.SetTxnResourceGroup(txn, sctx.GetSessionVars().StmtCtx.ResourceGroupName)
 
 	return nil
 }

pkg/sessionctx/sessionstates/session_states_test.go

@@ -548,12 +548,14 @@ func TestSessionCtx(t *testing.T) {
 			setFunc: func(tk *testkit.TestKit) any {
 				tk.MustExec("SET GLOBAL tidb_enable_resource_control='on'")
 				tk.MustExec("CREATE RESOURCE GROUP rg1 ru_per_sec = 100")
+				tk.MustExec("CREATE RESOURCE GROUP rg2 ru_per_sec = 100")
 				tk.MustExec("SET RESOURCE GROUP `rg1`")
 				require.Equal(t, "rg1", tk.Session().GetSessionVars().ResourceGroupName)
 				return nil
 			},
 			checkFunc: func(tk *testkit.TestKit, param any) {
 				tk.MustQuery("SELECT CURRENT_RESOURCE_GROUP()").Check(testkit.Rows("rg1"))
+				tk.MustQuery("SELECT /*+ RESOURCE_GROUP(rg2) */ CURRENT_RESOURCE_GROUP()").Check(testkit.Rows("rg2"))
 			},
 		},
 		{

pkg/sessionctx/stmtctx/stmtctx.go

@@ -250,18 +250,21 @@ type StatementContext struct {
 	MaxRowID  int64
 
 	// Copied from SessionVars.TimeZone.
-	Priority         mysql.PriorityEnum
-	NotFillCache     bool
-	MemTracker       *memory.Tracker
-	DiskTracker      *disk.Tracker
-	RunawayChecker   *resourcegroup.RunawayChecker
-	IsTiFlash        atomic2.Bool
-	RuntimeStatsColl *execdetails.RuntimeStatsColl
-	TableIDs         []int64
-	IndexNames       []string
-	StmtType         string
-	OriginalSQL      string
-	digestMemo       struct {
+	Priority     mysql.PriorityEnum
+	NotFillCache bool
+	MemTracker   *memory.Tracker
+	DiskTracker  *disk.Tracker
+	// per statement resource group name
+	// hint /* +ResourceGroup(name) */ can change the statement group name
+	ResourceGroupName string
+	RunawayChecker    *resourcegroup.RunawayChecker
+	IsTiFlash         atomic2.Bool
+	RuntimeStatsColl  *execdetails.RuntimeStatsColl
+	TableIDs          []int64
+	IndexNames        []string
+	StmtType          string
+	OriginalSQL       string
+	digestMemo        struct {
 		sync.Once
 		normalized string
 		digest     *parser.Digest

pkg/sessionctx/variable/session.go

@@ -1485,6 +1485,7 @@ type SessionVars struct {
 	shardRand *rand.Rand
 
 	// Resource group name
+	// NOTE: all statement relate opeartion should use StmtCtx.ResourceGroupName instead.
 	ResourceGroupName string
 
 	// PessimisticTransactionFairLocking controls whether fair locking for pessimistic transaction
@@ -2029,6 +2030,7 @@ func NewSessionVars(hctx HookContext) *SessionVars {
 		ResourceGroupName:             resourcegroup.DefaultResourceGroupName,
 		DefaultCollationForUTF8MB4:    mysql.DefaultCollationName,
 	}
+	vars.StmtCtx.ResourceGroupName = resourcegroup.DefaultResourceGroupName
 	vars.KVVars = tikvstore.NewVariables(&vars.SQLKiller.Signal)
 	vars.Concurrency = Concurrency{
 		indexLookupConcurrency:            DefIndexLookupConcurrency,

tests/integrationtest/r/executor/simple.result

@@ -338,6 +338,9 @@ SET RESOURCE GROUP default;
 SELECT CURRENT_RESOURCE_GROUP();
 CURRENT_RESOURCE_GROUP()
 default
+SELECT /*+ RESOURCE_GROUP(rg1)*/ CURRENT_RESOURCE_GROUP();
+CURRENT_RESOURCE_GROUP()
+rg1
 SELECT CURRENT_RESOURCE_GROUP();
 CURRENT_RESOURCE_GROUP()
 rg1

tests/integrationtest/t/executor/simple.test

@@ -391,6 +391,7 @@ SET RESOURCE GROUP ``;
 SELECT CURRENT_RESOURCE_GROUP();
 SET RESOURCE GROUP default;
 SELECT CURRENT_RESOURCE_GROUP();
+SELECT /*+ RESOURCE_GROUP(rg1)*/ CURRENT_RESOURCE_GROUP();
 
 connection conn1;
 SELECT CURRENT_RESOURCE_GROUP();