[PATCHSET cgroup/for-6.19] cgroup: Fix task exit ordering

[Tejun Heo <tj@kernel.org>]

Hello,

This series fixes a cgroup task exit ordering issue that is generally
suboptimal for all cgroup controllers and has caused real breakage for
sched_ext schedulers.

Currently, when a task exits, cgroup_task_exit() in do_exit() immediately
unlinks the task from its cgroup via css_set_move_task(). From the cgroup's
perspective, the task is now gone. If this makes the cgroup empty, it can be
destroyed, triggering ->css_offline() callbacks that notify controllers the
cgroup is going offline resource-wise.

However, the exiting task continues to run, perform memory operations, and
schedule until the final context switch in finish_task_switch(). This creates
a problematic window where controllers are told a cgroup is offline while
resource activities are still occurring in it. While this hasn't broken
existing controllers, it's clearly suboptimal and has caused real breakage
for sched_ext schedulers.

The sched_ext breakage manifests in two ways:

1. When a sched_ext scheduler is loaded, it walks all tasks and calls
   ops.init_task() on each. For tasks in a cgroup, it first ensures the
   cgroup has been initialized via ops.cgroup_init(). However, if the task
   is in the dying state (still running but already unlinked from its
   cgroup), the cgroup may already be offline. This results in
   ops.init_task() being called on a cgroup that never received
   ops.cgroup_init(), breaking the initialization invariant.

   This broke the scx_mitosis scheduler with errors like "cgrp_ctx lookup
   failed for cgid 3869" where the BPF program couldn't find cgroup context
   that was never created. See: https://github.com/sched-ext/scx/issues/2846

2. Because sched_ext_free() was called from __put_task_struct() (which can
   happen long after the task stops running), ops.cgroup_exit() could be
   called before ops.exit_task() was called on all member tasks, violating
   the expected ordering where all tasks exit before their cgroup does.

The fix defers the cgroup unlinking from do_exit() to finish_task_switch(),
ensuring the task remains linked to its cgroup until it's truly done running.
For sched_ext specifically, we also move the cleanup earlier to
finish_task_switch() to ensure proper ordering with cgroup operations.

This adds two new calls to finish_task_switch() that operate on the dead task
after the final switch: cgroup_task_dead() and sched_ext_dead(). It may make
sense to factor these into a helper function located in kernel/exit.c if this
pattern continues to grow.

Patch 0004 changes sched_ext and can be applied to sched_ext/for-6.19 after
pulling cgroup/for-6.19. Alternatively, would it be easier for some or all of
this series to go through the tip tree?

Based on cgroup/for-6.19 (d5cf4d34a333).

 0001 cgroup: Rename cgroup lifecycle hooks to cgroup_task_*()
 0002 cgroup: Move dying_tasks cleanup from cgroup_task_release() to cgroup_task_free()
 0003 cgroup: Defer task cgroup unlink until after the task is done switching out
 0004 sched_ext: Fix cgroup exit ordering by moving sched_ext_free() to finish_task_switch()

Git tree: git://git.kernel.org/pub/scm/linux/kernel/git/tj/cgroup.git cgroup-fix-exit-ordering

 include/linux/cgroup.h    | 14 ++++++++------
 include/linux/sched/ext.h |  4 ++--
 kernel/cgroup/cgroup.c    | 39 +++++++++++++++++++++++----------------
 kernel/exit.c             |  4 ++--
 kernel/fork.c             |  3 +--
 kernel/sched/autogroup.c  |  4 ++--
 kernel/sched/core.c       |  8 ++++++++
 kernel/sched/ext.c        |  2 +-
 8 files changed, 47 insertions(+), 31 deletions(-)

--
tejun