runaway: fix the issue where COOLDOWN/SWITCH_GROUP can't be triggered (pingcap#60457)

close pingcap#60404

Author: JmPotato

pkg/executor/adapter.go

@@ -562,12 +562,7 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) {
 			return nil, err
 		}
 		if len(switchGroupName) > 0 {
-			group, err := rm.ResourceGroupCtl.GetResourceGroup(switchGroupName)
-			if err != nil || group == nil {
-				logutil.BgLogger().Debug("invalid switch resource group", zap.String("switch-group-name", switchGroupName), zap.Error(err))
-			} else {
-				stmtCtx.ResourceGroupName = switchGroupName
-			}
+			stmtCtx.ResourceGroupName = switchGroupName
 		}
 	}
 	ctx = a.observeStmtBeginForTopSQL(ctx)

pkg/resourcegroup/runaway/checker.go

@@ -112,6 +112,13 @@ func (rm *Manager) DeriveChecker(resourceGroupName, originalSQL, sqlDigest, plan
 	return NewChecker(rm, resourceGroupName, group.RunawaySettings, originalSQL, sqlDigest, planDigest, startTime)
 }
 
+func (r *Checker) isMarkedByIdentifyInRunawaySettings() bool {
+	if r == nil {
+		return false
+	}
+	return r.markedByIdentifyInRunawaySettings.Load()
+}
+
 // BeforeExecutor checks whether query is in watch list before executing and after compiling.
 func (r *Checker) BeforeExecutor() (string, error) {
 	if r == nil {
@@ -149,34 +156,43 @@ func (r *Checker) BeforeExecutor() (string, error) {
 			return "", nil
 		case rmpb.RunawayAction_SwitchGroup:
 			// Return the switch group name to switch the resource group before executing.
-			return switchGroupName, nil
+			return r.checkSwitchGroupName(switchGroupName), nil
 		default:
 			// Continue to examine other convicts.
 		}
 	}
 	return "", nil
 }
 
+func (r *Checker) checkSwitchGroupName(groupName string) string {
+	if len(groupName) == 0 {
+		return ""
+	}
+	group, err := r.manager.ResourceGroupCtl.GetResourceGroup(groupName)
+	if err != nil || group == nil {
+		logutil.BgLogger().Debug("invalid switch resource group", zap.String("switch-group-name", groupName), zap.Error(err))
+		return ""
+	}
+	return groupName
+}
+
 // BeforeCopRequest checks runaway and modifies the request if necessary before sending coprocessor request.
 func (r *Checker) BeforeCopRequest(req *tikvrpc.Request) error {
 	if r == nil {
 		return nil
 	}
-	// If the group settings are not available, and it's not marked by watch, skip this part.
-	if r.settings == nil && !r.markedByQueryWatchRule {
-		return nil
-	}
-	// If it's marked by watch and the action is cooldown, override the priority,
+	// If it's marked by watch and the action is cooldown, override the priority.
+	// Other watch actions should already been handled in `BeforeExecutor` before.
 	if r.markedByQueryWatchRule && r.watchAction == rmpb.RunawayAction_CoolDown {
 		req.ResourceControlContext.OverridePriority = 1 // set priority to lowest
 	}
-	// If group settings are available and the query is not marked by a rule,
-	// verify if it matches any rules in the settings.
-	if r.settings != nil && !r.markedByIdentifyInRunawaySettings.Load() {
+	// If group settings are available, verify if it matches any rules in the settings.
+	if r.settings != nil {
 		now := time.Now()
-		// only check time and need to ensure deadline existed.
+		// Check and ensure the deadline exists.
 		exceedCause := r.exceedsThresholds(now, nil, 0)
-		if exceedCause == "" { // only set timeout when the query is not runaway.
+		// Only set timeout when the query has not been marked as runaway yet.
+		if !r.isMarkedByIdentifyInRunawaySettings() && len(exceedCause) == 0 {
 			if r.settings.Action == rmpb.RunawayAction_Kill {
 				until := r.deadline.Sub(now)
 				// if the execution time is close to the threshold, set a timeout
@@ -186,7 +202,8 @@ func (r *Checker) BeforeCopRequest(req *tikvrpc.Request) error {
 			}
 			return nil
 		}
-		// execution time exceeds the threshold, mark the query as runaway
+		// Try to mark the query as runaway, it's safe to call this method concurrently.
+		// So it's possible that the query has already been marked as runaway in `CheckThresholds`.
 		r.markRunawayByIdentifyInRunawaySettings(&now, exceedCause)
 		// Take action if needed.
 		switch r.settings.Action {
@@ -196,7 +213,9 @@ func (r *Checker) BeforeCopRequest(req *tikvrpc.Request) error {
 			req.ResourceControlContext.OverridePriority = 1 // set priority to lowest
 			return nil
 		case rmpb.RunawayAction_SwitchGroup:
-			req.ResourceControlContext.ResourceGroupName = r.settings.SwitchGroupName
+			if switchGroupName := r.checkSwitchGroupName(r.settings.SwitchGroupName); len(switchGroupName) != 0 {
+				req.ResourceControlContext.ResourceGroupName = switchGroupName
+			}
 			return nil
 		default:
 			return nil
@@ -214,7 +233,7 @@ func (r *Checker) CheckAction() rmpb.RunawayAction {
 	if r.markedByQueryWatchRule {
 		return r.watchAction
 	}
-	if r.markedByIdentifyInRunawaySettings.Load() {
+	if r.isMarkedByIdentifyInRunawaySettings() {
 		return r.settings.Action
 	}
 	return rmpb.RunawayAction_NoneAction
@@ -227,7 +246,7 @@ func (r *Checker) CheckRuleKillAction() (string, bool) {
 		return "", false
 	}
 	// If the group settings are available, and it's not marked by rule, check the execution time.
-	if r.settings != nil && !r.markedByIdentifyInRunawaySettings.Load() {
+	if r.settings != nil && !r.isMarkedByIdentifyInRunawaySettings() {
 		now := time.Now()
 		exceedCause := r.exceedsThresholds(now, nil, 0)
 		if exceedCause == "" {
@@ -249,15 +268,14 @@ func (r *Checker) markQuarantine(now *time.Time, exceedCause string) {
 		r.settings.Action, r.settings.SwitchGroupName, ttl, now, exceedCause)
 }
 
-func (r *Checker) markRunawayByIdentifyInRunawaySettings(now *time.Time, exceedCause string) bool {
+func (r *Checker) markRunawayByIdentifyInRunawaySettings(now *time.Time, exceedCause string) {
 	swapped := r.markedByIdentifyInRunawaySettings.CompareAndSwap(false, true)
 	if swapped {
 		r.markRunaway("identify", r.settings.Action, r.settings.SwitchGroupName, now, exceedCause)
 		if !r.markedByQueryWatchRule {
 			r.markQuarantine(now, exceedCause)
 		}
 	}
-	return swapped
 }
 
 func (r *Checker) markRunawayByQueryWatchRule(action rmpb.RunawayAction, switchGroupName, exceedCause string) {
@@ -320,7 +338,7 @@ func (r *Checker) CheckThresholds(ruDetail *util.RUDetails, processKeys int64, e
 			processKeys = int64(100)
 		}
 	})
-	if r.settings == nil || r.settings.Action != rmpb.RunawayAction_Kill {
+	if r.settings == nil {
 		return err
 	}
 
@@ -333,18 +351,15 @@ func (r *Checker) CheckThresholds(ruDetail *util.RUDetails, processKeys int64, e
 	atomic.AddInt64(&r.totalProcessedKeys, processKeys)
 	totalProcessedKeys := atomic.LoadInt64(&r.totalProcessedKeys)
 	exceedCause := r.exceedsThresholds(checkTime, ruDetail, totalProcessedKeys)
-	if !r.markedByIdentifyInRunawaySettings.Load() {
-		if exceedCause != "" && r.markRunawayByIdentifyInRunawaySettings(&now, exceedCause) {
-			if r.markRunawayByIdentifyInRunawaySettings(&now, exceedCause) {
-				return exeerrors.ErrResourceGroupQueryRunawayInterrupted.FastGenByArgs(exceedCause)
-			}
-		}
+	// No need to mark as runaway if the query is not exceeded any threshold.
+	if len(exceedCause) == 0 {
+		return err
 	}
-	// Due to concurrency, check again.
-	if r.markedByIdentifyInRunawaySettings.Load() {
+	r.markRunawayByIdentifyInRunawaySettings(&now, exceedCause)
+	// Other actions will be handled in `BeforeCopRequest` since they need to modify the request.
+	if r.settings.Action == rmpb.RunawayAction_Kill {
 		return exeerrors.ErrResourceGroupQueryRunawayInterrupted.FastGenByArgs(exceedCause)
 	}
-
 	return err
 }