[PATCH v4] sched/cache: Remove the old cache group footprint on exec

[Jemmy Wong]

Exec switches cache groups without subtracting the task's old footprint.
CLONE_VM tasks without CLONE_THREAD can retain the old mm and group
despite de_thread(), leaving an undecayed contribution that can suppress
cache-aware aggregation. This mainly affects longer-lived CLONE_VM tasks:
vfork() children rarely accumulate NUMA faults before exec because scanning
starts late. Found by inspection, without a workload reproducer.

Rename sched_cache_replace_grp() to sched_cache_switch_grp() and share
exec/exit accounting: subtract the old contribution, publish the new
pointer, and return the old reference for release. Exec subtraction must
precede task_numa_free() clearing total_numa_faults.

Convert footprint to atomic_long_t so concurrent placement updates
cannot overwrite exec/exit subtractions. Use non-returning atomic
add/sub, requiring no ordering against unrelated accesses, and apply
the net total_numa_faults delta once per placement pass to reduce
cache-line traffic. Assuming signed long accounting does not overflow,
updates balance without a write-side clamp; retain the defensive
read-side zero clamp and negative-value warning.

Keep NUMA reset ordering unchanged. If exec fails after switching groups
but before resetting statistics, exit can subtract the old total from
the new, private group. The task is past the point of no return, so that
group dies on exit; the surviving old shared group remains correctly
accounted.

Fixes: b636fef85bda ("sched/cache: Introduce task_struct->sched_cache_grp to fix UAF")
Suggested-by: Chen Yu <yu.c.chen@xxxxxxxxx>
Reviewed-by: Chen Yu <yu.c.chen@xxxxxxxxx>
Tested-by: Chen Yu <yu.c.chen@xxxxxxxxx>
Signed-off-by: Jemmy Wong <jemmywong512@xxxxxxxxx>
Assisted-by: LLM

---

v3: https://lore.kernel.org/all/20261015053604.29927-1-jemmywong512@xxxxxxxxx/
v2: https://lore.kernel.org/all/20261010111612.31345-1-jemmywong512@xxxxxxxxx/
v1: https://lore.kernel.org/all/20261009151215.62878-1-jemmywong512@xxxxxxxxx/

Changes in v4:
- Restore the footprint convergence comment next to the batched update in
  task_numa_placement(), referring to atomic_long_t as Chen Yu requested.
- Add Chen Yu's Reviewed-by, Tested-by and Assisted-by tags.
- Shorten the commit description while retaining the accounting rationale.

Changes in v3:
- Use atomic_long_t and atomic add/sub to prevent lost accounting updates.
- Batch placement updates using the net total_numa_faults delta.
- Rename sched_cache_replace_grp() to sched_cache_switch_grp(); share
  exec/exit subtraction, guarded by old && old != new.
- Drop the write-side clamp; use an atomic read, warning and zero clamp.
- Document failed-exec over-subtraction in the dying task's private group;
  keep NUMA reset ordering unchanged.
- Add Chen Yu's Suggested-by; drop Tim Chen's Reviewed-by pending renewed
  review of the accounting changes.

Changes in v2:
- Mention vfork(), explain that longer-lived CLONE_VM tasks are the main
  concern, and state that the issue was found by code inspection.
- Add the Fixes tag and document that exec footprint subtraction must
  precede task_numa_free() resetting p->total_numa_faults.
- Add Tim Chen's Reviewed-by tag (not carried forward to v3).

Validation:
- v4: Strict checkpatch and whitespace checks passed.
- v3: Built kernel/sched/{fair,core,build_utility}.o on x86-64 with:
  - CONFIG_SCHED_CACHE=y, CONFIG_NUMA_BALANCING=y
  - CONFIG_SCHED_CACHE=y, CONFIG_NUMA_BALANCING=n
  - CONFIG_SCHED_CACHE=n, CONFIG_NUMA_BALANCING=y
- Chen Yu tested v3 with hackbench on a Xeon with multiple LLCs per node
  and reported no obvious performance difference. Some apparent
  regressions had large run-to-run variance; selected reruns were also
  unstable. His report is at:
  https://lore.kernel.org/all/atiP8TwaDDtbImRn@chenyu-dev/
- Build and runtime tests have not been rerun for this comment-only v4
  revision.


 include/linux/sched.h |  3 +-
 kernel/sched/fair.c   | 90 +++++++++++++++++++------------------------
 2 files changed, 41 insertions(+), 52 deletions(-)

diff --git a/include/linux/sched.h b/include/linux/sched.h
index 12a7e094a3e6..bb0a424b6ad8 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -2391,7 +2391,8 @@ struct sched_cache_group {
 	unsigned long epoch;
 	u64 nr_running_avg;
 	unsigned long next_scan;
-	unsigned long footprint;
+	/* Sum of the member tasks' accounted footprints, in pages. */
+	atomic_long_t footprint;
 	int cpu;
 	refcount_t refcnt;
 	struct rcu_head rcu;
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index b5546cd27dba..ab5c117d32a2 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1427,6 +1427,7 @@ static bool exceed_llc_capacity(struct sched_cache_group *grp, int cpu)
 #ifdef CONFIG_NUMA_BALANCING
 	unsigned long llc, footprint;
 	struct sched_domain *sd;
+	long fp;
 	int scale;

 	guard(rcu)();
@@ -1441,7 +1442,9 @@ static bool exceed_llc_capacity(struct sched_cache_group *grp, int cpu)
 		 * excluded.
 		 */
 		llc = sd->llc_bytes;
-		footprint = READ_ONCE(grp->footprint);
+		fp = atomic_long_read(&grp->footprint);
+		WARN_ON_ONCE(fp < 0);
+		footprint = max(fp, 0L);

 		/*
 		 * Scale the LLC size by 256*llc_aggr_tolerance
@@ -1621,7 +1624,7 @@ int mm_init_sched(struct mm_struct *mm,
 	grp->cpu = -1;
 	grp->next_scan = jiffies;
 	grp->nr_running_avg = 0;
-	grp->footprint = 0;
+	atomic_long_set(&grp->footprint, 0);
 	refcount_set(&grp->refcnt, 1);
 	/*
 	 * The update to grp->pcpu_sched should not be reordered
@@ -1663,12 +1666,26 @@ DEFINE_FREE(sched_cache_group_put, struct sched_cache_group *,
 #define rcu_deref_sched_cache_grp(tsk) \
 	rcu_dereference_check((tsk)->sched_cache_grp, (tsk) == current)

-static struct sched_cache_group *sched_cache_replace_grp(struct task_struct *p,
-							 struct sched_cache_group *new)
+static inline void sched_cache_footprint_update(struct task_struct *p, long delta)
+{
+	struct sched_cache_group *grp = rcu_deref_sched_cache_grp(p);
+
+	if (grp && delta)
+		atomic_long_add(delta, &grp->footprint);
+}
+
+/* Remove the old contribution and return the old group for the caller to put. */
+static struct sched_cache_group *sched_cache_switch_grp(struct task_struct *p,
+							struct sched_cache_group *new)
 {
 	struct sched_cache_group *old;

 	old = rcu_deref_sched_cache_grp(p);
+#ifdef CONFIG_NUMA_BALANCING
+	/* Other tasks sharing the old mm can keep its group alive after exec. */
+	if (old && old != new)
+		atomic_long_sub(p->total_numa_faults, &old->footprint);
+#endif
 	rcu_assign_pointer(p->sched_cache_grp, new);

 	return old;
@@ -1722,28 +1739,18 @@ void sched_cache_exec_mmap(struct task_struct *p, struct mm_struct *mm)
 	/*
 	 * Acquire the new reference before publishing the pointer, then drop
 	 * the old one. @p is current and the only writer of its own pointer.
+	 *
+	 * Footprint subtraction must precede task_numa_free() in bprm_execve(),
+	 * which clears p->total_numa_faults. Nothing here enforces that order.
 	 */
-	old = sched_cache_replace_grp(p, sched_cache_group_get(mm->sched_cache_grp));
+	old = sched_cache_switch_grp(p, sched_cache_group_get(mm->sched_cache_grp));
 	sched_cache_group_put(old);
 }

 void sched_cache_exit_mm(struct task_struct *p)
 {
-	struct sched_cache_group *grp = sched_cache_replace_grp(p, NULL);
+	struct sched_cache_group *grp = sched_cache_switch_grp(p, NULL);

-#ifdef CONFIG_NUMA_BALANCING
-	/*
-	 * Subtract this task's footprint from the group before dropping the
-	 * reference, so the group footprint converges as its threads exit.
-	 * Unlocked for performance; clamp to avoid underflow.
-	 */
-	if (grp && p->total_numa_faults) {
-		unsigned long fp = READ_ONCE(grp->footprint);
-		unsigned long sub = min(fp, p->total_numa_faults);
-
-		WRITE_ONCE(grp->footprint, fp - sub);
-	}
-#endif
 	sched_cache_group_put(grp);
 }

@@ -2133,6 +2140,9 @@ void init_sched_mm(struct task_struct *p)
 static inline void account_mm_sched(struct rq *rq, struct task_struct *p,
 				    s64 delta_exec) { }

+static inline void sched_cache_footprint_update(struct task_struct *p,
+						long delta) { }
+
 void init_sched_mm(struct task_struct *p) { }

 static void task_tick_cache(struct rq *rq, struct task_struct *p) { }
@@ -3830,14 +3840,13 @@ static int preferred_group_nid(struct task_struct *p, int nid)
 static void task_numa_placement(struct task_struct *p)
 	__context_unsafe(/* conditional locking */)
 {
-	struct sched_cache_group __maybe_unused *grp;
 	int seq, nid, max_nid = NUMA_NO_NODE;
 	unsigned long max_faults = 0;
 	unsigned long fault_types[2] = { 0, 0 };
+	unsigned long old_faults = p->total_numa_faults;
 	unsigned long total_faults;
 	u64 runtime, period;
 	spinlock_t *group_lock = NULL;
-	long __maybe_unused new_fp;
 	struct numa_group *ng;

 	/*
@@ -3912,36 +3921,6 @@ static void task_numa_placement(struct task_struct *p)
 				ng->total_faults += diff;
 				group_faults += ng->faults[mem_idx];
 			}
-#ifdef CONFIG_SCHED_CACHE
-			/*
-			 * Per task p->numa_faults[mem_idx] converges,
-			 * so the accumulation of each task's faults
-			 * converges too - Given the number of threads,
-			 * it cannot overflow an unsigned long.
-			 * Racy with concurrent updates from other threads
-			 * sharing this mm. Acceptable since footprint is a
-			 * heuristic and occasional lost updates are tolerable.
-			 *
-			 * If a task exits, its corresponding footprint must
-			 * be subtracted from p->sched_cache_grp->footprint,
-			 * otherwise the footprint will not converge: the
-			 * exiting thread's footprint remains unchanged/undecayed.
-			 * See exit_mm().
-			 *
-			 * Lost updates and unsynchronized subtraction
-			 * in exit_mm() can cause footprint + diff to
-			 * go negative. Clamp to zero to prevent the
-			 * unsigned footprint from wrapping.
-			 */
-			scoped_guard(rcu) {
-				grp = rcu_dereference(p->sched_cache_grp);
-
-				if (grp) {
-					new_fp = (long)READ_ONCE(grp->footprint) + diff;
-					WRITE_ONCE(grp->footprint, max(new_fp, 0L));
-				}
-			}
-#endif
 		}

 		if (!ng) {
@@ -3964,6 +3943,15 @@ static void task_numa_placement(struct task_struct *p)
 		max_nid = preferred_group_nid(p, max_nid);
 	}

+	/*
+	 * Per task p->numa_faults[mem_idx] converges, so the accumulation of each
+	 * task's faults converges too. Given the number of threads, it cannot
+	 * overflow an atomic_long_t.
+	 *
+	 * Apply the net change once per placement pass to limit cache-line traffic.
+	 */
+	sched_cache_footprint_update(p, p->total_numa_faults - old_faults);
+
 	if (max_faults) {
 		/* Set the new preferred node */
 		if (max_nid != p->numa_preferred_nid)
--
2.54.0 (Apple Git-157)