mm: memcontrol: don't lose track of exiting OOM victims

When oom_evaluate_task() comes across a victim we've already killed, it backs off until the victim's mm has MMF_OOM_SKIP. This gives the victim a chance to free some memory before we kill anything else. However, since commit a72f73c4dd9b ("cgroup: Don't expose dead tasks in cgroup"), css_task_iter skips exiting tasks once signal->live reaches zero. That happens before exit_mm(), so a single-threaded victim can disappear from the memcg scan while it still has all its memory.

In tests I used a 64 MiB memcg with swap disabled, and found that the OOM killer could kill both processes if the first was still waiting for AIO. The first process faults in 44 MiB of anon and submits a direct AIO read to null_blk with a 10 second completion delay. The second then tries to fault in 40 MiB. The first OOM kills the larger process, which gets stuck in exit_aio() waiting for the read. About 5 ms later we kill the second process too, even though the reaper gets to the first one about two seconds later. memory.events showed oom_kill 2 in all three runs.

The cgroup iterator hides these tasks so that cgroup.procs doesn't show processes that waitpid() may already have reaped. That's fine for cgroup.procs, but the OOM scan needs to see a victim even after its last thread has started exiting. We already have CSS_TASK_ITER_WITH_DEAD for sched_ext, so use it for mem_cgroup_scan_tasks() as well.

With this applied, the second process finishes faulting in its memory and memory.events shows oom_kill 1 in all three runs.

---

Zhan's patch for migrating a process whose leader has exited also changes this iterator call. This patch deals with the case where the last thread is already exiting, which came up in review of that patch.

https://lore.kernel.org/linux-mm/[email protected]/

 mm/memcontrol.c | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

index 84f4a8fd4d16..cb7cfae9ce5a 100644
@@ -1451,7 +1451,13 @@ void mem_cgroup_scan_tasks(struct mem_cgroup *memcg,
 		struct css_task_iter it;
 		struct task_struct *task;
 
-		css_task_iter_start(&iter->css, CSS_TASK_ITER_PROCS, &it);
+		/*
+		 * An OOM victim can start exiting well before its mm is
+		 * torn down. Keep it visible so oom_evaluate_task() can
+		 * back off while there's still memory to free.
+		 */
+		css_task_iter_start(&iter->css, CSS_TASK_ITER_PROCS |
+				    CSS_TASK_ITER_WITH_DEAD, &it);
 		while (!ret && (task = css_task_iter_next(&it))) {
 			ret = fn(task, arg);
 			/* Avoid potential softlockup warning */

base-commit: fa4ad798f6f681031775266044ba1b0cf3c670bd