[PATCH] sched/cache: honor migrate_llc_task semantics in active load balance

[Lu Wang <wanglu.priv@gmail.com>]

A passive load-balance pass marks group_llc_balance as migrate_llc_task
and queues active balance when it cannot move a task. The CPU stopper
callback constructs a fresh lb_env, so preserve the migration type on
the runqueue across the asynchronous boundary.

For CAS-directed active balance, reject a candidate whose preferred LLC
does not match the destination LLC. This keeps the fallback from moving
a task away from its preferred LLC.

Fixes: e4c9a4cb244a ("sched/cache: Add migrate_llc_task migration type for cache-aware balancing")
Signed-off-by: Lu Wang <wanglu.priv@gmail.com>
---
 kernel/sched/fair.c  | 25 ++++++++++++++++++++++---
 kernel/sched/sched.h |  1 +
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index d78467ec6..cf036a8a0 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -10645,6 +10645,18 @@ alb_break_llc(struct lb_env *env)
 	return false;
 }
 
+/*
+ * Returns true if p's preferred LLC does not match the destination CPU,
+ * meaning this task should not be migrated under migrate_llc_task semantics.
+ */
+static inline bool
+migrate_llc_task_wrong_dst(struct task_struct *p, struct lb_env *env)
+{
+	return sched_cache_enabled() &&
+	       env->migration_type == migrate_llc_task &&
+	       READ_ONCE(p->preferred_llc) != llc_id(env->dst_cpu);
+}
+
 /*
  * Check if migrating task p from env->src_cpu to
  * env->dst_cpu breaks LLC localiy.
@@ -10673,8 +10685,7 @@ static bool migrate_degrades_llc(struct task_struct *p, struct lb_env *env)
 	 * run on env->dst_cpu, skip the tasks do not prefer
 	 * env->dst_cpu, and find the one that prefers.
 	 */
-	if (env->migration_type == migrate_llc_task &&
-	    READ_ONCE(p->preferred_llc) != llc_id(env->dst_cpu))
+	if (migrate_llc_task_wrong_dst(p, env))
 		return true;
 
 	if (can_migrate_llc_task(env->src_cpu,
@@ -10697,6 +10708,12 @@ alb_break_llc(struct lb_env *env)
 	return false;
 }
 
+static inline bool
+migrate_llc_task_wrong_dst(struct task_struct *p, struct lb_env *env)
+{
+	return false;
+}
+
 static inline bool
 migrate_degrades_llc(struct task_struct *p, struct lb_env *env)
 {
@@ -10796,7 +10813,7 @@ int can_migrate_task(struct task_struct *p, struct lb_env *env)
 	 * 4) too many balance attempts have failed.
 	 */
 	if (env->flags & LBF_ACTIVE_LB)
-		return 1;
+		return !migrate_llc_task_wrong_dst(p, env);
 
 	degrades = migrate_degrades_locality(p, env);
 	if (!degrades) {
@@ -13485,6 +13502,7 @@ static int sched_balance_rq(int this_cpu, struct rq *this_rq,
 			if (!busiest->active_balance) {
 				busiest->active_balance = 1;
 				busiest->push_cpu = this_cpu;
+				busiest->active_balance_type = env.migration_type;
 				active_balance = 1;
 			}
 
@@ -13654,6 +13672,7 @@ static int active_load_balance_cpu_stop(void *data)
 			.src_rq		= busiest_rq,
 			.idle		= CPU_IDLE,
 			.flags		= LBF_ACTIVE_LB,
+			.migration_type	= (enum migration_type)busiest_rq->active_balance_type,
 		};
 
 		schedstat_inc(sd->alb_count);
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 56acf502b..82084d405 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1266,6 +1266,7 @@ struct rq {
 	/* For active balancing */
 	int			active_balance;
 	int			push_cpu;
+	int			active_balance_type; /* enum migration_type */
 	struct cpu_stop_work	active_balance_work;
 
 	/* CPU of this runqueue: */
-- 
2.43.0