Re: [RFC PATCH v3 12/17] sched/cache: Apply multi-LLC aggregation policy to task migration

[sashiko-bot@kernel.org]

Thank you for your contribution! Sashiko AI review found 3 potential issue(s) to consider:
- [Critical] sched/fair: NULL deref and mm_struct UAF in alb_break_llc()
- [High] sched/fair: alb_break_llc() blocks active balance when pref_llc < 0
- [Medium] sched/fair: active balance bypasses migrate_degrades_llc() when degrades > 0
--

commit 7bd76e4bf626d210dfac7f3c5332cbe2a77e81f9
Author: Jianyong Wu <wujianyong@hygon.cn>

sched/cache: Apply multi-LLC aggregation policy to task migration

This commit updates cache-aware task migration and active load balancing to
use if_to_prefer() and can_migrate_llcs() across all LLC pairs, removes the
can_migrate_llc() helper, and moves the LBF_ACTIVE_LB check in
can_migrate_task() after the locality check.

> diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
> index 03bb5c85a8b7b..c3f9ef8798d6c 100644
> --- a/kernel/sched/fair.c
> +++ b/kernel/sched/fair.c

[ ... ]

> @@ -10827,55 +10779,52 @@ static enum llc_mig can_migrate_llc_task(int src_cpu, int dst_cpu,

[ ... ]

>  static inline bool
>  alb_break_llc(struct lb_env *env)
>  {
> +	int pref_llc = -1;
> +	bool to_pref = false;
> +
>  	if (!sched_cache_enabled())
>  		return false;
>  
>  	if (cpus_share_cache(env->src_cpu, env->dst_cpu))
>  		return false;
>  	/*
> -	 * All tasks prefer to stay on their current CPU.
> -	 * Do not pull a task from its preferred CPU if:
> -	 * 1. It is the only task running and does not exceed
> -	 *    imbalance allowance; OR
> -	 * 2. Migrating it away from its preferred LLC would violate
> -	 *    the cache-aware scheduling policy.
> +	 * We need the preferred LLC to decide whether we can perform migration.
> +	 * Therefore, we need to obtain task_struct, which is only meaningful
> +	 * in the case that only one task on the rq.
> +	 * For cases with more than one task on the rq, we need to check
> +	 * this in can_migrate_task().
>  	 */
> -	if (env->src_rq->nr_pref_llc_running &&
> -	    env->src_rq->nr_pref_llc_running == env->src_rq->cfs.h_nr_runnable) {
> -		unsigned long util = 0;
> +	if (env->src_rq->nr_running == 1) {
>  		struct task_struct *cur;
>  
> -		if (env->src_rq->nr_running <= 1)
> -			return true;
> -
>  		cur = rcu_dereference_all(env->src_rq->curr);
> -		if (cur && cur->sched_class == &fair_sched_class)
> -			util = task_util(cur);
> +		if (!cur || !cur->mm)
> +			return false;
> +
> +		if (cur->sched_class == &fair_sched_class)
> +			pref_llc = llc_id(cur->mm->sc_stat.cpu);

[Severity: Critical]
Can this dereference of cur->mm cause a NULL pointer dereference or a
use-after-free on struct mm_struct?

In sched_balance_rq(), need_active_balance() -> alb_break_llc() is called
before locking busiest (env->src_rq) and again after unlocking it, holding
only rcu_read_lock():

kernel/sched/fair.c:sched_balance_rq() {
	...
	if (need_active_balance(&env)) {
		unsigned long flags;

		raw_spin_rq_lock_irqsave(busiest, flags);
	...
}

Because cur is actively running on env->src_cpu and neither the runqueue lock
nor task_lock(cur) is held, cur can concurrently execute exit_mm() or
kthread_unuse_mm() and clear current->mm = NULL between the cur->mm check and
the reload in cur->mm->sc_stat.cpu:

kernel/exit.c:exit_mm() {
	...
	local_irq_disable();
	current->mm = NULL;
	membarrier_update_current_mm(NULL);
	enter_lazy_tlb(mm, current);
	local_irq_enable();
	task_unlock(current);
	mmap_read_unlock(mm);
	mm_update_next_owner(mm);
	mmput(mm);
}

Even if the mm pointer is loaded only once, rcu_read_lock() only protects
struct task_struct, whereas mm_cachep does not use SLAB_TYPESAFE_BY_RCU and
__mmdrop() frees struct mm_struct immediately via free_mm(). If cur replaces
its mm in exec_mmap() or exits and switches out on the unlocked remote
runqueue, can alb_break_llc() and can_migrate_llcs() end up reading from a
freed struct mm_struct?

>  
> -		if (can_migrate_llc(env->src_cpu, env->dst_cpu,
> -				    util, false) == mig_forbid)
> +		to_pref = if_to_prefer(env->src_cpu, env->dst_cpu, pref_llc);
> +
> +		if (can_migrate_llcs(env->src_cpu, env->dst_cpu,
> +				    cur, to_pref) == mig_forbid)
>  			return true;
>  	}

[Severity: High]
When pref_llc is negative (for instance, when cur->mm->sc_stat.cpu is -1 on
single-threaded tasks, or when cur is not in fair_sched_class), shouldn't
alb_break_llc() return false instead of calling can_migrate_llcs()?

When pref_llc < 0, if_to_prefer() returns false, and can_migrate_llcs() falls
back to using src_cpu as target_cpu:

kernel/sched/fair.c:can_migrate_llcs() {
	...
	if (p) {
		struct mm_struct *mm = p->mm;

		if (mm)
			target_cpu = READ_ONCE(mm->sc_stat.cpu);
		target_cpu = target_cpu >= 0 ? target_cpu : src_cpu;

		tsk_util = task_util(p);
	}
	...
}

With to_pref == false and target_cpu == src_cpu, can_migrate_llcs() treats
src_cpu as the task's preferred cache location and returns mig_forbid whenever
fits_llc_capacity() is true on the source LLC or node:

kernel/sched/fair.c:can_migrate_llcs() {
	...
			if (src_node == dst_node) {
				/* Leave the source LLC at or above the threshold. */
				if (fits_llc_capacity(src_util, src_cap, false))
					return mig_forbid;
	...
}

Unlike can_migrate_llc_task(), which returns mig_unrestricted when cpu < 0,
won't this regression cause alb_break_llc() to block cross-LLC active balance
for tasks that have no preferred LLC?

[ ... ]

> @@ -11027,14 +10976,11 @@ int can_migrate_task(struct task_struct *p, struct lb_env *env)
>  
>  	/*
>  	 * Aggressive migration if:
> -	 * 1) active balance
> -	 * 2) destination numa is preferred
> +	 * 1) destination numa is preferred
> +	 * 2) active balance
>  	 * 3) task is cache cold, or
>  	 * 4) too many balance attempts have failed.
>  	 */
> -	if (env->flags & LBF_ACTIVE_LB)
> -		return 1;
> -
>  	degrades = migrate_degrades_locality(p, env);
>  	if (!degrades) {
>  		/*
> @@ -11060,6 +11006,13 @@ int can_migrate_task(struct task_struct *p, struct lb_env *env)
>  		hot = degrades > 0;
>  	}
>  
> +	/*
> +	 * Keep this after the cache-aware check so active balance respects its
> +	 * migration veto when NUMA locality is neutral.
> +	 */
> +	if (env->flags & LBF_ACTIVE_LB)
> +		return 1;

[Severity: Medium]
What happens during cross-node active load balancing when
migrate_degrades_locality() returns a positive value (degrades > 0)?

In active_load_balance_cpu_stop(), the stopper task is queued on env->src_rq,
so env->src_rq->nr_running > env->src_rq->nr_preferred_running is true in
migrate_degrades_locality():

kernel/sched/fair.c:migrate_degrades_locality() {
	...
	/* Migrating away from the preferred node is always bad. */
	if (src_nid == p->numa_preferred_nid) {
		if (env->src_rq->nr_running > env->src_rq->nr_preferred_running)
			return 1;
		else
			return 0;
	}
	...
}

When a task is on its preferred NUMA node, migrate_degrades_locality() returns
1, which skips the if (!degrades) block containing migrate_degrades_llc() and
sets hot = 1. Because the LBF_ACTIVE_LB check right after that block returns 1
unconditionally, won't active balance bypass migrate_degrades_llc() whenever
migrating the task also degrades NUMA locality?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261029022036.4002331-1-wujianyong@hygon.cn?part=12