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

[Jemmy Wong <jemmywong512@gmail.com>]

Exec replaces the task's cache group before resetting its NUMA fault
statistics, but only the exit path subtracts the task's contribution
from the old group's footprint.

Although de_thread() removes other members of the executing task's
thread group, tasks sharing the mm through CLONE_VM without CLONE_THREAD
can retain the old mm and cache group. The executing task's contribution
then remains without further updates or decay, potentially suppressing
cache-aware aggregation through the LLC capacity check. A vfork() child
rarely accumulates NUMA faults before exec because scanning starts late;
the issue mainly concerns longer-lived CLONE_VM tasks. This was found by
code inspection, not a workload reproducer.

Simply subtracting on exec is insufficient: the unlocked read/modify/write
updates can race with NUMA placement by another task sharing the group,
which can overwrite the subtraction. Convert footprint to atomic_long_t
and use non-returning atomic add/sub operations, which require no ordering
against unrelated accesses. Apply the net change in total_numa_faults once
per placement pass to limit shared cache-line traffic.

Rename sched_cache_replace_grp() to sched_cache_switch_grp() to reflect
its expanded responsibility: subtract the task's contribution when leaving
the old group, publish the new pointer, and return the old reference for
the caller to release. Use it for both exec and exit, and document that
exec subtraction must precede task_numa_free() resetting the NUMA total.

With atomic updates, placement deltas balance the departing task's total,
so no write-side clamp is needed, provided the accounting fits in signed
long. Retain a defensive read-side clamp and warn on a negative footprint
observed by the LLC capacity check.

Keep the existing NUMA reset ordering. If exec fails after switching groups
but before resetting the statistics, exit can subtract the old total from
the new group. This mismatch is confined to the execing task's private
group: the task is past the point of no return and the group disappears 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@intel.com>
Signed-off-by: Jemmy Wong <jemmywong512@gmail.com>
---
v1: https://lore.kernel.org/all/20261009151215.62878-1-jemmywong512@gmail.com/
v2: https://lore.kernel.org/all/20261010111612.31345-1-jemmywong512@gmail.com/

Changes in v3:
- Convert grp->footprint to atomic_long_t and use atomic add/sub to avoid
  losing placement updates or exec/exit subtractions, as Chen Yu suggested.
- Update the group footprint once per NUMA placement pass, using the net
  change in p->total_numa_faults rather than updating each fault bucket.
- Rename sched_cache_replace_grp() to sched_cache_switch_grp() and move
  leave-group accounting into it, as agreed with Chen Yu. Exec and exit
  share the subtraction, guarded by old && old != new.
- Remove the write-side clamp now that updates are atomic. Use an atomic
  read, WARN_ON_ONCE(), and a defensive zero clamp in the capacity check.
- Explain the failed-exec over-subtraction case and why it is confined to
  the dying task's private group; leave NUMA reset ordering unchanged.
- Add Suggested-by: Chen Yu. Drop Tim Chen's Reviewed-by tag because the
  accounting changes are substantial and need renewed review.

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:
- Strict checkpatch and whitespace checks passed.
- 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
- Runtime and performance testing have not been performed.

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

diff --git a/include/linux/sched.h b/include/linux/sched.h
index d35ae49a991f..4a02a47c988e 100644
--- a/include/linux/sched.h
+++ b/include/linux/sched.h
@@ -2412,7 +2412,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 57360f5cdde4..47faa0bf12c7 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1483,6 +1483,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)();
@@ -1497,7 +1498,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
@@ -1677,7 +1680,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
@@ -1719,12 +1722,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;
@@ -1778,28 +1795,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);
-
-#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);
+	struct sched_cache_group *grp = sched_cache_switch_grp(p, NULL);

-		WRITE_ONCE(grp->footprint, fp - sub);
-	}
-#endif
 	sched_cache_group_put(grp);
 }

@@ -2189,6 +2196,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) { }
@@ -3883,14 +3893,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;

 	/*
@@ -3965,36 +3974,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) {
@@ -4017,6 +3996,9 @@ static void task_numa_placement(struct task_struct *p)
 		max_nid = preferred_group_nid(p, max_nid);
 	}

+	/* 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.25.1