[PATCH v3] sched/fair: Distribute burst-woken tasks when rq depth reaches 3x average

[Li kunyu]

When a task burst-wakes a large number of tasks in a short period, tasks
can pile up on a single CPU regardless of whether the system has idle
CPUs or NO_HZ enabled.

During the periodic load-balance pass at the top-level domain, compute
and cache the 3x average runqueue depth threshold (floored at 3 tasks)
in root_domain, updating it only when the value changes.

In select_task_rq_fair(), before scanning domains or siblings, check if
the waking CPU's runqueue depth has reached this threshold. If so, fall
back to cpumask_any_and_distribute(p->cpus_ptr, cpu_active_mask) to
rotate waking tasks across allowed CPUs (provided the target has a
smaller queue depth) and pull back rq->next_balance to expedite the next
periodic load-balance pass.

Suggested-by: Peter Zijlstra <peterz@xxxxxxxxxxxxx>
Link: https://lore.kernel.org/all/20261021125419.GK2823843@xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx/
Signed-off-by: Li kunyu <likunyu@xxxxxxxxxxxxxx>
---
Changes in v3:
- Remove READ_ONCE and WRITE_ONCE from hotspot functions to avoid
  excessive overhead.
- Move the overload check to the end of the select_task_rq_fair function
  and apply it only to the final task dispersal effect.

Changes in v2:
- Rebase onto latest mainline (Kayra Cizmeci).
- Cache pre-calculated 3x threshold in root_domain and only WRITE_ONCE()
  when the threshold changes to avoid hotpath math and cacheline dirtying
  (K Prateek Nayak).
- Expedite periodic balance via rq->next_balance gated by time_after()
  instead of walking domains on every wakeup (K Prateek Nayak, Kayra Cizmeci).
---
 kernel/sched/fair.c     | 68 ++++++++++++++++++++++++++++++++++++++---
 kernel/sched/sched.h    |  3 ++
 kernel/sched/topology.c |  2 ++
 3 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index cfc982205b98..a4201fcff132 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -9900,6 +9900,36 @@ static int find_energy_efficient_cpu(struct task_struct *p, int prev_cpu)
 	return target;
 }
 
+/*
+ * Check if a CPU's rq depth reaches the 3x system-wide average threshold
+ * cached in root_domain (floored at 3 tasks).
+ */
+static inline bool rq_burst_overloaded(int cpu)
+{
+	struct rq *rq = cpu_rq(cpu);
+
+	if (unlikely(!rq->rd))
+		return false;
+
+	/* Directly compare the values to avoid excessive costs. */
+	return rq->nr_running >= rq->rd->overload_nr_threshold;
+}
+
+static inline void reset_periodic_lb_interval(int cpu)
+{
+	struct sched_domain *sd;
+	unsigned long next = jiffies + 1;
+	struct rq *rq = cpu_rq(cpu);
+
+	if (time_after(rq->next_balance, next))
+		rq->next_balance = next;
+
+	for_each_domain(cpu, sd) {
+		if (sd->balance_interval > sd->min_interval)
+			sd->balance_interval = sd->min_interval;
+	}
+}
+
 /*
  * select_task_rq_fair: Select target runqueue for the waking task in domains
  * that have the relevant SD flag set. In practice, this is SD_BALANCE_WAKE,
@@ -9970,12 +10000,21 @@ select_task_rq_fair(struct task_struct *p, int prev_cpu, int wake_flags)
 	/* Slow path */
 	if (unlikely(sd)) {
 		new_cpu = sched_balance_find_dst_cpu(sd, p, cpu, prev_cpu, sd_flag);
-		return select_idle_smt_cpu(p, new_cpu);
+		new_cpu = select_idle_smt_cpu(p, new_cpu);
+	/* Fast path */
+	} else if (wake_flags & WF_TTWU) {
+		new_cpu = select_idle_sibling(p, prev_cpu, new_cpu);
 	}
 
-	/* Fast path */
-	if (wake_flags & WF_TTWU)
-		return select_idle_sibling(p, prev_cpu, new_cpu);
+	if (unlikely(rq_burst_overloaded(new_cpu))) {
+		int dist_cpu = cpumask_any_and_distribute(p->cpus_ptr, cpu_active_mask);
+
+		if (likely(dist_cpu < nr_cpumask_bits && cpu_rq(dist_cpu)->nr_running <
+			   cpu_rq(new_cpu)->nr_running)) {
+			reset_periodic_lb_interval(dist_cpu);
+			return dist_cpu;
+		}
+	}
 
 	return new_cpu;
 }
@@ -11682,6 +11721,7 @@ struct sg_lb_stats {
 struct sd_lb_stats {
 	struct sched_group *busiest;		/* Busiest group in this sd */
 	struct sched_group *local;		/* Local group in this sd */
+	unsigned long total_nr_running;		/* Sum of nr_running in this sd */
 	unsigned long total_load;		/* Total load of all groups in sd */
 	unsigned long total_capacity;		/* Total capacity of all groups in sd */
 	unsigned long avg_load;			/* Average load across all groups in sd */
@@ -11705,6 +11745,7 @@ static inline void init_sd_lb_stats(struct sd_lb_stats *sds)
 		.local = NULL,
 		.total_load = 0UL,
 		.total_capacity = 0UL,
+		.total_nr_running = 0UL,
 		.busiest_stat = {
 			.idle_cpus = UINT_MAX,
 			.group_type = group_has_spare,
@@ -12977,6 +13018,7 @@ static inline void update_sd_lb_stats(struct lb_env *env, struct sd_lb_stats *sd
 		/* Now, start updating sd_lb_stats */
 		sds->total_load += sgs->group_load;
 		sds->total_capacity += sgs->group_capacity;
+		sds->total_nr_running += sgs->sum_nr_running;
 
 		sum_util += sgs->group_util;
 		sg = sg->next;
@@ -12995,11 +13037,29 @@ static inline void update_sd_lb_stats(struct lb_env *env, struct sd_lb_stats *sd
 		env->fbq_type = fbq_classify_group(&sds->busiest_stat);
 
 	if (!env->sd->parent) {
+		unsigned int weight = cpumask_weight(env->cpus);
+
 		/* update overload indicator if we are at root domain */
 		set_rd_overloaded(env->dst_rq->rd, sg_overloaded);
 
 		/* Update over-utilization (tipping point, U >= 0) indicator */
 		set_rd_overutilized(env->dst_rq->rd, sg_overutilized);
+
+		/*
+		 * We update the average task count at the top-level (root)
+		 * domain, as this incurs the least overhead and comes
+		 * closest to the ideal scenario. The trade-off is that we
+		 * must accept some temporary data inaccuracy.
+		 */
+		if (weight) {
+			unsigned long thresh;
+
+			thresh = max_t(unsigned long,
+				       (sds->total_nr_running * 3) / weight, 3);
+
+			if (READ_ONCE(env->dst_rq->rd->overload_nr_threshold) != thresh)
+				WRITE_ONCE(env->dst_rq->rd->overload_nr_threshold, thresh);
+		}
 	} else if (sg_overutilized) {
 		set_rd_overutilized(env->dst_rq->rd, sg_overutilized);
 	}
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 6c7c5633524a..71d946da9600 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1096,6 +1096,9 @@ struct root_domain {
 	cpumask_var_t		rto_mask;
 	struct cpupri		cpupri;
 
+	/* 3x average per-CPU rq depth threshold for burst wakeup distribution */
+	unsigned long		overload_nr_threshold;
+
 	/*
 	 * NULL-terminated list of performance domains intersecting with the
 	 * CPUs of the rd. Protected by RCU.
diff --git a/kernel/sched/topology.c b/kernel/sched/topology.c
index 919d0fb00bd9..fe154dfdf23e 100644
--- a/kernel/sched/topology.c
+++ b/kernel/sched/topology.c
@@ -572,6 +572,8 @@ static int init_rootdomain(struct root_domain *rd)
 
 	if (cpupri_init(&rd->cpupri) != 0)
 		goto free_cpudl;
+
+	rd->overload_nr_threshold = 3;
 	return 0;
 
 free_cpudl:
-- 
2.47.3