Commit 7900aa69 authored by Tejun Heo's avatar Tejun Heo
Browse files

sched_ext: Fix cgroup exit ordering by moving sched_ext_free() to finish_task_switch() 



sched_ext_free() was called from __put_task_struct() when the last reference
to the task is dropped, which could be long after the task has finished
running. This causes cgroup-related problems:

- ops.init_task() can be called on a cgroup which didn't get ops.cgroup_init()'d
  during scheduler load, because the cgroup might be destroyed/unlinked
  while the zombie or dead task is still lingering on the scx_tasks list.

- ops.cgroup_exit() could be called before ops.exit_task() is called on all
  member tasks, leading to incorrect exit ordering.

Fix by moving it to finish_task_switch() to be called right after the final
context switch away from the dying task, matching when sched_class->task_dead()
is called. Rename it to sched_ext_dead() to match the new calling context.

By calling sched_ext_dead() before cgroup_task_dead(), we ensure that:

- Tasks visible on scx_tasks list have valid cgroups during scheduler load,
  as cgroup_mutex prevents cgroup destruction while the task is still linked.

- All member tasks have ops.exit_task() called and are removed from scx_tasks
  before the cgroup can be destroyed and trigger ops.cgroup_exit().

This fix is made possible by the cgroup_task_dead() split in the previous patch.

This also makes more sense resource-wise as there's no point in keeping
scheduler side resources around for dead tasks.

Reported-by: default avatarDan Schatzberg <dschatzberg@meta.com>
Cc: Peter Zijlstra <peterz@infradead.org>
Reviewed-by: default avatarAndrea Righi <arighi@nvidia.com>
Signed-off-by: default avatarTejun Heo <tj@kernel.org>
parent 587eb08a

include/linux/sched/ext.h

+2 −2
Original line number Diff line number Diff line
@@ -208,14 +208,14 @@ struct sched_ext_entity { 
	struct list_head	tasks_node;

};



void sched_ext_free(struct task_struct *p);

void sched_ext_dead(struct task_struct *p);

void print_scx_info(const char *log_lvl, struct task_struct *p);

void scx_softlockup(u32 dur_s);

bool scx_rcu_cpu_stall(void);



#else	/* !CONFIG_SCHED_CLASS_EXT */



static inline void sched_ext_free(struct task_struct *p) {}

static inline void sched_ext_dead(struct task_struct *p) {}

static inline void print_scx_info(const char *log_lvl, struct task_struct *p) {}

static inline void scx_softlockup(u32 dur_s) {}

static inline bool scx_rcu_cpu_stall(void) { return false; }

kernel/fork.c

+0 −1
Original line number Diff line number Diff line
@@ -736,7 +736,6 @@ void __put_task_struct(struct task_struct *tsk) 
	WARN_ON(tsk == current);



	unwind_task_free(tsk);

	sched_ext_free(tsk);

	io_uring_free(tsk);

	cgroup_task_free(tsk);

	task_numa_free(tsk, true);

kernel/sched/core.c

+6 −0
Original line number Diff line number Diff line
@@ -5151,6 +5151,12 @@ static struct rq *finish_task_switch(struct task_struct *prev) 
		if (prev->sched_class->task_dead)

			prev->sched_class->task_dead(prev);



		/*

		 * sched_ext_dead() must come before cgroup_task_dead() to

		 * prevent cgroups from being removed while its member tasks are

		 * visible to SCX schedulers.

		 */

		sched_ext_dead(prev);

		cgroup_task_dead(prev);



		/* Task is done with its stack. */

kernel/sched/ext.c

+1 −1
Original line number Diff line number Diff line
@@ -2966,7 +2966,7 @@ void scx_cancel_fork(struct task_struct *p) 
	percpu_up_read(&scx_fork_rwsem);

}



void sched_ext_free(struct task_struct *p)

void sched_ext_dead(struct task_struct *p)

{

	unsigned long flags;