planner,txn: fix the bug that RESOURCE_GROUP() hint can not take effe…

…ct for write … (#44513) (#44704)

close #44512

planner/optimize.go

@@ -138,7 +138,7 @@ func getPlanFromNonPreparedPlanCache(ctx context.Context, sctx sessionctx.Contex
 
 // Optimize does optimization and creates a Plan.
 // The node must be prepared first.
-func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is infoschema.InfoSchema) (core.Plan, types.NameSlice, error) {
+func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is infoschema.InfoSchema) (plan core.Plan, slice types.NameSlice, retErr error) {
 	sessVars := sctx.GetSessionVars()
 	if sessVars.StmtCtx.EnableOptimizerDebugTrace {
 		debugtrace.EnterContextCommon(sctx)
@@ -176,6 +176,25 @@ func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 	for _, warn := range warns {
 		sessVars.StmtCtx.AppendWarning(warn)
 	}
+
+	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
+		if retErr == nil && sessVars.StmtCtx.StmtHints.HasResourceGroup {
+			if variable.EnableResourceControl.Load() {
+				sessVars.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() {
+					txn.SetOption(kv.ResourceGroupName, sessVars.ResourceGroupName)
+				}
+			} else {
+				err := infoschema.ErrResourceGroupSupportDisabled
+				sessVars.StmtCtx.AppendWarning(err)
+			}
+		}
+	}()
+
 	warns = warns[:0]
 	for name, val := range originStmtHints.SetVars {
 		err := sessVars.SetStmtVar(name, val)
@@ -309,17 +328,6 @@ func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 		hint.BindHint(stmtNode, originHints)
 	}
 
-	// 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
-	if sessVars.StmtCtx.StmtHints.HasResourceGroup {
-		if variable.EnableResourceControl.Load() {
-			sessVars.ResourceGroupName = sessVars.StmtCtx.StmtHints.ResourceGroup
-		} else {
-			err := infoschema.ErrResourceGroupSupportDisabled
-			sessVars.StmtCtx.AppendWarning(err)
-		}
-	}
-
 	if sessVars.StmtCtx.EnableOptimizerDebugTrace && bestPlanFromBind != nil {
 		core.DebugTraceBestBinding(sctx, chosenBinding.Hint)
 	}

session/session.go

@@ -640,7 +640,6 @@ func (s *session) doCommit(ctx context.Context) error {
 	s.txn.SetOption(kv.EnableAsyncCommit, sessVars.EnableAsyncCommit)
 	s.txn.SetOption(kv.Enable1PC, sessVars.Enable1PC)
 	s.txn.SetOption(kv.ResourceGroupTagger, sessVars.StmtCtx.GetResourceGroupTagger())
-	s.txn.SetOption(kv.ResourceGroupName, sessVars.ResourceGroupName)
 	if sessVars.StmtCtx.KvExecCounter != nil {
 		// Bind an interceptor for client-go to count the number of SQL executions of each TiKV.
 		s.txn.SetOption(kv.RPCInterceptor, sessVars.StmtCtx.KvExecCounter.RPCInterceptor())

session/txn.go

@@ -272,7 +272,7 @@ func (txn *LazyTxn) changeToPending(future *txnFuture) {
 	txn.txnFuture = future
 }
 
-func (txn *LazyTxn) changePendingToValid(ctx context.Context) error {
+func (txn *LazyTxn) changePendingToValid(ctx context.Context, sctx sessionctx.Context) error {
 	if txn.txnFuture == nil {
 		return errors.New("transaction future is not set")
 	}
@@ -307,6 +307,9 @@ func (txn *LazyTxn) changePendingToValid(ctx context.Context) error {
 		txn.mu.TxnInfo.CurrentSQLDigest,
 		txn.mu.TxnInfo.AllSQLDigests)
 
+	// set resource group name for kv request such as lock pessimistic keys.
+	txn.SetOption(kv.ResourceGroupName, sctx.GetSessionVars().ResourceGroupName)
+
 	return nil
 }
 
@@ -594,16 +597,14 @@ func (txn *LazyTxn) Wait(ctx context.Context, sctx sessionctx.Context) (kv.Trans
 		// Transaction is lazy initialized.
 		// PrepareTxnCtx is called to get a tso future, makes s.txn a pending txn,
 		// If Txn() is called later, wait for the future to get a valid txn.
-		if err := txn.changePendingToValid(ctx); err != nil {
+		if err := txn.changePendingToValid(ctx, sctx); err != nil {
 			logutil.BgLogger().Error("active transaction fail",
 				zap.Error(err))
 			txn.cleanup()
 			sctx.GetSessionVars().TxnCtx.StartTS = 0
 			return txn, err
 		}
 		txn.lazyUniquenessCheckEnabled = !sctx.GetSessionVars().ConstraintCheckInPlacePessimistic
-		// set resource group name for kv request such as lock pessimistic keys.
-		txn.SetOption(kv.ResourceGroupName, sctx.GetSessionVars().ResourceGroupName)
 	}
 	return txn, nil
 }