[PATCHSET sched_ext/for-7.3] sched_ext: Cgroup migration and op delivery for sub-schedulers

[PATCHSET sched_ext/for-7.3] sched_ext: Cgroup migration and op delivery for sub-schedulers

[Tejun Heo]

Hello,

A task's sched must match its cgroup's: each sub-scheduler serves the
cgroup2 subtree it is attached to and the root scheduler serves the rest.
Cgroup migration currently breaks this. A task moved across a
sub-scheduler boundary keeps its old sched, leading to wrong-sched
scheduling and, once the stale sched is freed, a use-after-free. Cgroup
ops are also always delivered to the root scheduler no matter which sched
serves the cgroup.

This patchset makes both follow the sub-scheduler topology:

- cgroup gains a task migration notifier (0001). sched_ext uses it to
  re-home tasks whose migration crosses a sched boundary: the destination
  sched runs the fallible ops.init_task() before the migration commits,
  so a rejection fails the cgroup.procs write (0002-0004).

- cgroup_init/exit and the knob ops are delivered to the sched each
  task_group is on. The move ops go to the task's sched, only for
  migrations that don't re-home it (0005).

- Sub-scheduler enable claims the subtree's cgroups from the parent and
  disable returns them. A parent that fails to re-init a returned cgroup
  fails in turn and the cgroups move up the chain, ending at the root
  scheduler (0006).

- scx_qmap consumes cgroup weights through ops.cgroup_set_weight() and
  gains fault injection modes to exercise the failure paths (0007-0008).

Verified with a trace-based suite on the demo scheduler covering op
delivery and routing, handover failures, multi-level punt chains, cgroup1
and stress.

Based on sched_ext/for-7.3 (7c2cd767705d).

Tejun Heo (8):
  cgroup: Add cgroup_task_notifier and task migration events
  sched_ext: Factor out scx_rehome_task() and scx_punt_task()
  sched_ext: Relocate scx_cgroup_enabled
  sched_ext: Re-home tasks on cgroup migration
  sched_ext: Deliver cgroup ops to each task_group's sched
  sched_ext: Hand over cgroups at sub-scheduler enable/disable
  tools/sched_ext: scx_qmap - Consume cgroup weights through set_weight
  tools/sched_ext: scx_qmap - Add init fault injection modes

Git tree: git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext.git scx-cgroup-migration

 include/linux/cgroup.h         |  26 +++
 include/linux/sched/ext.h      |  15 ++
 kernel/cgroup/cgroup.c         |  93 ++++++++-
 kernel/sched/ext/ext.c         | 247 +++++++++++++++++-------
 kernel/sched/ext/internal.h    |  44 ++++-
 kernel/sched/ext/sub.c         | 418 +++++++++++++++++++++++++++++++++++++----
 tools/sched_ext/scx_qmap.bpf.c |  95 ++++++++--
 tools/sched_ext/scx_qmap.c     | 110 +----------
 tools/sched_ext/scx_qmap.h     |   5 +-
 9 files changed, 819 insertions(+), 234 deletions(-)

--
tejun

[PATCH 1/8] cgroup: Add cgroup_task_notifier and task migration events

[Tejun Heo]

A subsystem can attach to the cgroup hierarchy itself, independent of which
controllers are enabled where - BPF hooks already behave this way and
sched_ext sub-schedulers do too. Controller callbacks can't track task
migrations for them: sched_ext must re-home a task whose migration crosses a
sub-scheduler boundary, but the cpu controller's attach callbacks fire only
when the task_group changes and miss moves whenever the controller topology
is coarser than the sub-scheduler topology.

Add cgroup_task_notifier with per-task migration events mirroring the
can_attach/attach/cancel_attach phases so that a consumer which prepares
per-task state can also veto a migration: CGROUP_TASK_MIGRATING fires
pre-commit, CGROUP_TASK_MIGRATED post-commit and
CGROUP_TASK_MIGRATE_CANCELED unwinds a failed migration. Only migrations
that change a task's dfl cgroup are reported.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/cgroup.h | 26 ++++++++++++
 kernel/cgroup/cgroup.c | 93 ++++++++++++++++++++++++++++++++++++++----
 2 files changed, 110 insertions(+), 9 deletions(-)

diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index f2aa46a4f871..aa92db5f05de 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -82,12 +82,38 @@ enum cgroup_lifetime_events {
 	CGROUP_LIFETIME_OFFLINE,
 };
 
+/*
+ * Events on cgroup_task_notifier, data is struct cgroup_task_migrate_ctx.
+ * MIGRATING fires per task before the migration commits and an error return
+ * from the chain fails the migration, in which case tasks that were already
+ * notified receive MIGRATE_CANCELED. MIGRATED fires per task after the
+ * migration is committed and can't fail. Only migrations that change a task's
+ * dfl cgroup are reported.
+ */
+enum cgroup_task_events {
+	CGROUP_TASK_MIGRATING,
+	CGROUP_TASK_MIGRATED,
+	CGROUP_TASK_MIGRATE_CANCELED,
+};
+
+/*
+ * @src_dcgrp and @dst_dcgrp are @task's dfl cgroups before and after the
+ * migration. @src_dcgrp is NULL for CGROUP_TASK_MIGRATED as per-task sources
+ * are not tracked past the commit point.
+ */
+struct cgroup_task_migrate_ctx {
+	struct task_struct	*task;
+	struct cgroup		*src_dcgrp;
+	struct cgroup		*dst_dcgrp;
+};
+
 extern struct file_system_type cgroup_fs_type;
 extern struct cgroup_root cgrp_dfl_root;
 extern struct css_set init_css_set;
 extern struct mutex cgroup_mutex;
 extern spinlock_t css_set_lock;
 extern struct blocking_notifier_head cgroup_lifetime_notifier;
+extern struct blocking_notifier_head cgroup_task_notifier;
 
 #define SUBSYS(_x) extern struct cgroup_subsys _x ## _cgrp_subsys;
 #include <linux/cgroup_subsys.h>
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index 38f8d9df8fbc..2f6b634c84ca 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -88,6 +88,8 @@ EXPORT_SYMBOL_GPL(css_set_lock);
 
 struct blocking_notifier_head cgroup_lifetime_notifier =
 	BLOCKING_NOTIFIER_INIT(cgroup_lifetime_notifier);
+struct blocking_notifier_head cgroup_task_notifier =
+	BLOCKING_NOTIFIER_INIT(cgroup_task_notifier);
 
 DEFINE_SPINLOCK(trace_cgroup_path_lock);
 char trace_cgroup_path[TRACE_CGROUP_PATH_LEN];
@@ -2676,14 +2678,27 @@ struct task_struct *cgroup_taskset_next(struct cgroup_taskset *tset,
 	return NULL;
 }
 
+static void cgroup_migrate_notify_canceled(struct css_set *src_cset,
+					   struct task_struct *task)
+{
+	struct cgroup_task_migrate_ctx ctx = {
+		.task = task,
+		.src_dcgrp = src_cset->dfl_cgrp,
+		.dst_dcgrp = src_cset->mg_dst_cset->dfl_cgrp,
+	};
+
+	blocking_notifier_call_chain(&cgroup_task_notifier,
+				     CGROUP_TASK_MIGRATE_CANCELED, &ctx);
+}
+
 /**
  * cgroup_migrate_execute - migrate a taskset
  * @mgctx: migration context
  *
- * Migrate tasks in @mgctx as setup by migration preparation functions.
- * This function fails iff one of the ->can_attach callbacks fails and
- * guarantees that either all or none of the tasks in @mgctx are migrated.
- * @mgctx is consumed regardless of success.
+ * Migrate tasks in @mgctx as setup by migration preparation functions. This
+ * function fails iff one of the ->can_attach callbacks or CGROUP_TASK_MIGRATING
+ * notifications fails and guarantees that either all or none of the tasks in
+ * @mgctx are migrated. @mgctx is consumed regardless of success.
  */
 static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 {
@@ -2691,6 +2706,7 @@ static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 	struct cgroup_subsys *ss;
 	struct task_struct *task, *tmp_task;
 	struct css_set *cset, *tmp_cset;
+	bool dfl_migration = false;
 	int ssid, failed_ssid, ret;
 
 	/* check that we can legitimately attach to the cgroup */
@@ -2707,6 +2723,33 @@ static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 		} while_each_subsys_mask();
 	}
 
+	/*
+	 * Notify each task about the impending migration. An error return fails
+	 * the migration. Only migrations on the default hierarchy are reported:
+	 * a migration modifies either every moved task's dfl cgroup or, on
+	 * cgroup1 or for subtree_control writes, none.
+	 */
+	list_for_each_entry(cset, &tset->src_csets, mg_node) {
+		if (cset->dfl_cgrp == cset->mg_dst_cset->dfl_cgrp)
+			continue;
+		dfl_migration = true;
+		list_for_each_entry(task, &cset->mg_tasks, cg_list) {
+			struct cgroup_task_migrate_ctx ctx = {
+				.task = task,
+				.src_dcgrp = cset->dfl_cgrp,
+				.dst_dcgrp = cset->mg_dst_cset->dfl_cgrp,
+			};
+
+			ret = blocking_notifier_call_chain_robust(&cgroup_task_notifier,
+								  CGROUP_TASK_MIGRATING,
+								  CGROUP_TASK_MIGRATE_CANCELED,
+								  &ctx);
+			ret = notifier_to_errno(ret);
+			if (ret)
+				goto out_cancel_migrating;
+		}
+	}
+
 	/*
 	 * Now that we're guaranteed success, proceed to move all tasks to
 	 * the new cgroup.  There are no failure cases after here, so this
@@ -2750,9 +2793,41 @@ static int cgroup_migrate_execute(struct cgroup_mgctx *mgctx)
 		} while_each_subsys_mask();
 	}
 
+	/*
+	 * Notify each task after successful migration. The operation can no
+	 * longer fail and the return value is ignored. The MIGRATING loop
+	 * above explains why only dfl migrations are reported. Per-task
+	 * sources are not tracked past the commit point, so src_dcgrp is
+	 * NULL.
+	 */
+	if (dfl_migration) {
+		list_for_each_entry(cset, &tset->dst_csets, mg_node) {
+			list_for_each_entry(task, &cset->mg_tasks, cg_list) {
+				struct cgroup_task_migrate_ctx ctx = {
+					.task = task,
+					.dst_dcgrp = cset->dfl_cgrp,
+				};
+
+				blocking_notifier_call_chain(
+					&cgroup_task_notifier,
+					CGROUP_TASK_MIGRATED, &ctx);
+			}
+		}
+	}
+
 	ret = 0;
 	goto out_release_tset;
 
+out_cancel_migrating:
+	list_for_each_entry_continue_reverse(task, &cset->mg_tasks, cg_list)
+		cgroup_migrate_notify_canceled(cset, task);
+	list_for_each_entry_continue_reverse(cset, &tset->src_csets, mg_node) {
+		if (cset->dfl_cgrp == cset->mg_dst_cset->dfl_cgrp)
+			continue;
+		list_for_each_entry_reverse(task, &cset->mg_tasks, cg_list)
+			cgroup_migrate_notify_canceled(cset, task);
+	}
+	failed_ssid = CGROUP_SUBSYS_COUNT;
 out_cancel_attach:
 	if (tset->nr_tasks) {
 		do_each_subsys_mask(ss, ssid, mgctx->ss_mask) {
@@ -2976,11 +3051,11 @@ int cgroup_migrate_prepare_dst(struct cgroup_mgctx *mgctx)
  * cgroup_migrate_prepare_dst() on the targets before invoking this
  * function and following up with cgroup_migrate_finish().
  *
- * As long as a controller's ->can_attach() doesn't fail, this function is
- * guaranteed to succeed.  This means that, excluding ->can_attach()
- * failure, when migrating multiple targets, the success or failure can be
- * decided for all targets by invoking group_migrate_prepare_dst() before
- * actually starting migrating.
+ * As long as a controller's ->can_attach() or a CGROUP_TASK_MIGRATING
+ * notification doesn't fail, this function is guaranteed to succeed.  This
+ * means that, excluding those failures, when migrating multiple targets,
+ * the success or failure can be decided for all targets by invoking
+ * group_migrate_prepare_dst() before actually starting migrating.
  */
 int cgroup_migrate(struct task_struct *leader, bool threadgroup,
 		   struct cgroup_mgctx *mgctx)
-- 
2.55.0

[PATCH 2/8] sched_ext: Factor out scx_rehome_task() and scx_punt_task()

[Tejun Heo]

Factor out scx_rehome_task() and scx_punt_task() from the sub-disable
re-home loop and scx_fail_parent(). The upcoming cgroup migration re-homing
also needs scx_rehome_task(). No functional changes.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/sub.c | 84 ++++++++++++++++++++++++++++--------------
 1 file changed, 56 insertions(+), 28 deletions(-)

diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index c7f70cf877d1..5f7ac6696d17 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -758,6 +758,57 @@ void drain_descendants(struct scx_sched *sch)
 	wait_event(scx_unlink_waitq, list_empty(&sch->children));
 }
 
+/**
+ * scx_rehome_task - Move a task to a sched it has been initialized for
+ * @to: sched taking over @p, @p's init on it already complete
+ * @p: task to re-home
+ *
+ * Exit @p from its current sched and switch it over to @to, overriding the
+ * state to %SCX_TASK_READY to account for the already completed init. A task
+ * on a non-ext class, possible under an %SCX_OPS_SWITCH_PARTIAL root, stays
+ * %READY and is enabled by switching_to_scx() if it switches over.
+ */
+static void scx_rehome_task(struct scx_sched *to, struct task_struct *p)
+{
+	lockdep_assert_held(&p->pi_lock);
+	lockdep_assert_rq_held(task_rq(p));
+
+	scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
+		scx_disable_and_exit_task(scx_task_sched(p), p);
+		scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
+		scx_set_task_state(p, SCX_TASK_INIT);
+		scx_set_task_sched(p, to);
+		scx_set_task_state(p, SCX_TASK_READY);
+		if (p->sched_class == &ext_sched_class)
+			scx_enable_task(to, p);
+	}
+}
+
+/**
+ * scx_punt_task - Hand a task to a failed sched without initialization
+ * @to: failed and bypassed sched taking custody of @p
+ * @p: task to punt
+ *
+ * Take @p off its current sched and put it on @to at %SCX_TASK_NONE. @to is
+ * dying and its teardown will re-home @p properly.
+ *
+ * Used when @to must take over @p but failed to initialize it. Bypass keeps
+ * scheduling decisions away from @to but @p can still trigger its task ops,
+ * which may confuse the BPF side. @to is dying anyway. The exit paths skip
+ * %NONE tasks (see __scx_disable_and_exit_task() and switched_from_scx()).
+ */
+static void scx_punt_task(struct scx_sched *to, struct task_struct *p)
+{
+	lockdep_assert_held(&p->pi_lock);
+	lockdep_assert_rq_held(task_rq(p));
+	WARN_ON_ONCE(!READ_ONCE(to->bypass_depth));
+
+	scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
+		scx_disable_and_exit_task(scx_task_sched(p), p);
+		scx_set_task_sched(p, to);
+	}
+}
+
 static void scx_fail_parent(struct scx_sched *sch,
 			    struct task_struct *failed, s32 fail_code)
 {
@@ -769,9 +820,9 @@ static void scx_fail_parent(struct scx_sched *sch,
 		  fail_code, failed->comm, failed->pid);
 
 	/*
-	 * Once $parent is bypassed, it's safe to put SCX_TASK_NONE tasks into
-	 * it. This may cause downstream failures on the BPF side but $parent is
-	 * dying anyway.
+	 * Once $parent is bypassed, tasks can be punted into it. This may
+	 * cause downstream failures on the BPF side but $parent is dying
+	 * anyway.
 	 */
 	scx_bypass(parent, true);
 
@@ -780,10 +831,7 @@ static void scx_fail_parent(struct scx_sched *sch,
 		if (scx_task_on_sched(parent, p))
 			continue;
 
-		scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
-			scx_disable_and_exit_task(sch, p);
-			scx_set_task_sched(p, parent);
-		}
+		scx_punt_task(parent, p);
 	}
 	scx_task_iter_stop(&sti);
 }
@@ -881,27 +929,7 @@ void scx_sub_disable(struct scx_sched *sch)
 			continue;
 		}
 
-		scoped_guard (sched_change, p, DEQUEUE_SAVE | DEQUEUE_MOVE) {
-			/*
-			 * $p is initialized for $parent and still attached to
-			 * @sch. Disable and exit for @sch, switch over to
-			 * $parent and override the state to READY to account
-			 * for $p having already been initialized.
-			 */
-			scx_disable_and_exit_task(sch, p);
-			scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
-			scx_set_task_state(p, SCX_TASK_INIT);
-			scx_set_task_sched(p, parent);
-			scx_set_task_state(p, SCX_TASK_READY);
-
-			/*
-			 * A task on a non-ext class, possible under an
-			 * %SCX_OPS_SWITCH_PARTIAL root, stays READY and is
-			 * enabled by switching_to_scx() if it switches over.
-			 */
-			if (p->sched_class == &ext_sched_class)
-				scx_enable_task(parent, p);
-		}
+		scx_rehome_task(parent, p);
 
 		task_rq_unlock(rq, p, &rf);
 		put_task_struct(p);
-- 
2.55.0

[PATCH 3/8] sched_ext: Relocate scx_cgroup_enabled

[Tejun Heo]

scx_cgroup_enabled is in the CONFIG_EXT_GROUP_SCHED block. The upcoming
cgroup migration re-homing needs the gate outside the block. Move the
definition and flag flips outside CONFIG_EXT_GROUP_SCHED. No functional
changes.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      | 18 ++++++++++++------
 kernel/sched/ext/internal.h |  1 +
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 145bda57b6fe..cd99947b4229 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -80,6 +80,14 @@ static bool scx_switching_all;
 DEFINE_STATIC_KEY_FALSE(__scx_switched_all);
 static DEFINE_STATIC_KEY_FALSE(__scx_tid_to_task_enabled);
 
+/*
+ * Gates cgroup ops delivery. Set at the end of the cgroup init phase of root
+ * enable and cleared before root disable starts tearing down tasks, both under
+ * scx_cgroup_lock(). Holding cgroup_lock() and seeing %true guarantees no race
+ * against root tearing down tasks.
+ */
+bool scx_cgroup_enabled;
+
 /*
  * True once SCX_OPS_TID_TO_TASK has been negotiated with the root scheduler
  * and the tid->task table is live. Wraps the static key so callers don't
@@ -4309,7 +4317,6 @@ bool scx_can_stop_tick(struct rq *rq)
 #ifdef CONFIG_EXT_GROUP_SCHED
 
 DEFINE_STATIC_PERCPU_RWSEM(scx_cgroup_ops_rwsem);
-static bool scx_cgroup_enabled;
 
 void scx_tg_init(struct task_group *tg)
 {
@@ -4689,8 +4696,6 @@ static void scx_cgroup_exit(struct scx_sched *sch)
 {
 	struct cgroup_subsys_state *css;
 
-	scx_cgroup_enabled = false;
-
 	/*
 	 * scx_tg_on/offline() are excluded through cgroup_lock(). If we walk
 	 * cgroups and exit all the inited ones, all online cgroups are exited.
@@ -4745,9 +4750,6 @@ static int scx_cgroup_init(struct scx_sched *sch)
 		tg->scx.flags |= SCX_TG_INITED;
 	}
 
-	WARN_ON_ONCE(scx_cgroup_enabled);
-	scx_cgroup_enabled = true;
-
 	return 0;
 }
 
@@ -5971,6 +5973,7 @@ static void scx_root_disable(struct scx_sched *sch)
 	 * doesn't race against scx_disable_and_exit_task().
 	 */
 	scx_cgroup_lock();
+	scx_cgroup_enabled = false;
 	scx_cgroup_exit(sch);
 	scx_cgroup_unlock();
 
@@ -7234,6 +7237,9 @@ static void scx_root_enable_workfn(struct kthread_work *work)
 	if (ret)
 		goto err_disable_unlock_all;
 
+	WARN_ON_ONCE(scx_cgroup_enabled);
+	scx_cgroup_enabled = true;
+
 	scx_task_iter_start(&sti, NULL);
 	while ((p = scx_task_iter_next_locked(&sti))) {
 		/*
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 4f4130f0d121..d76ac22019af 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1949,6 +1949,7 @@ __printf(3, 0) s32 scx_bstr_format(struct scx_sched *sch, struct scx_bstr_buf *b
 extern raw_spinlock_t scx_sched_lock;
 extern struct mutex scx_enable_mutex;
 extern struct percpu_rw_semaphore scx_fork_rwsem;
+extern bool scx_cgroup_enabled;
 extern raw_spinlock_t scx_exit_bstr_buf_lock;
 extern struct scx_bstr_buf scx_exit_bstr_buf;
 #ifdef CONFIG_EXT_SUB_SCHED
-- 
2.55.0

[PATCH 4/8] sched_ext: Re-home tasks on cgroup migration

[Tejun Heo]

A task's sched (p->scx.sched) must match its cgroup's owner
(cgrp->scx_sched). cgroup migration breaks the invariant:
scx_cgroup_move_task() only fires root's ops.cgroup_move() and never
re-homes the task, leading to wrong-sched scheduling and, once the stale
sched is freed, a use-after-free.

Hook into the new cgroup task migration events and re-home each task whose
destination cgroup is owned by a different sched. The events map naturally
to the transfer: MIGRATING runs the fallible init for the destination sched,
letting it reject the migration the same way ops.cgroup_prep_move() can,
MIGRATED does the re-home, which can't fail, and CANCELED undoes the init
when the migration falls through.

Pre-commit, the task's task_group still reflects the source, so
__scx_init_task() grows an explicit cgroup argument for the migration path
to hand ops.init_task() the destination cgroup.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      |  29 +++++---
 kernel/sched/ext/internal.h |   3 +-
 kernel/sched/ext/sub.c      | 135 ++++++++++++++++++++++++++++++++++--
 3 files changed, 152 insertions(+), 15 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index cd99947b4229..6bac68758704 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -3485,15 +3485,28 @@ static struct cgroup *tg_cgrp(struct task_group *tg)
 		return &cgrp_dfl_root.cgrp;
 }
 
-#define SCX_INIT_TASK_ARGS_CGROUP(tg)		.cgroup = tg_cgrp(tg),
+#define SCX_INIT_TASK_ARGS_CGROUP(cgrp)		.cgroup = (cgrp),
 
 #else	/* CONFIG_EXT_GROUP_SCHED */
 
-#define SCX_INIT_TASK_ARGS_CGROUP(tg)
+#define SCX_INIT_TASK_ARGS_CGROUP(cgrp)
 
 #endif	/* CONFIG_EXT_GROUP_SCHED */
 
-int __scx_init_task(struct scx_sched *sch, struct task_struct *p, bool fork)
+/**
+ * __scx_init_task - Initialize a task for a sched
+ * @sch: sched to initialize @p for
+ * @p: task of interest
+ * @cgrp: cgroup @p is joining, %NULL for @p's current task_group's cgroup
+ * @fork: %true if @p is being forked
+ *
+ * Pre-commit cgroup migration passes @cgrp explicitly as @p's task_group
+ * still reflects the source.
+ *
+ * Return 0 on success, -errno on failure.
+ */
+int __scx_init_task(struct scx_sched *sch, struct task_struct *p,
+		    struct cgroup *cgrp, bool fork)
 {
 	int ret;
 
@@ -3501,7 +3514,7 @@ int __scx_init_task(struct scx_sched *sch, struct task_struct *p, bool fork)
 
 	if (SCX_HAS_OP(sch, init_task)) {
 		struct scx_init_task_args args = {
-			SCX_INIT_TASK_ARGS_CGROUP(task_group(p))
+			SCX_INIT_TASK_ARGS_CGROUP(cgrp ?: tg_cgrp(task_group(p)))
 			.fork = fork,
 		};
 
@@ -3747,7 +3760,7 @@ int scx_fork(struct task_struct *p, struct kernel_clone_args *kargs)
 		struct scx_sched *sch = scx_root;
 #endif
 		scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
-		ret = __scx_init_task(sch, p, true);
+		ret = __scx_init_task(sch, p, NULL, true);
 		if (unlikely(ret)) {
 			scx_set_task_state(p, SCX_TASK_NONE);
 			return ret;
@@ -5969,8 +5982,8 @@ static void scx_root_disable(struct scx_sched *sch)
 	WRITE_ONCE(scx_switching_all, false);
 
 	/*
-	 * Shut down cgroup support before tasks so that the cgroup attach path
-	 * doesn't race against scx_disable_and_exit_task().
+	 * Shut down cgroup support before tasks so that the cgroup attach and
+	 * migration paths don't race against scx_disable_and_exit_task().
 	 */
 	scx_cgroup_lock();
 	scx_cgroup_enabled = false;
@@ -7263,7 +7276,7 @@ static void scx_root_enable_workfn(struct kthread_work *work)
 		scx_set_task_state(p, SCX_TASK_INIT_BEGIN);
 		scx_task_iter_unlock(&sti);
 
-		ret = __scx_init_task(sch, p, false);
+		ret = __scx_init_task(sch, p, NULL, false);
 
 		scx_task_iter_relock(&sti, p);
 
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index d76ac22019af..23fc95502ea6 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1920,7 +1920,8 @@ void scx_flush_dispatch_buf(struct scx_sched *sch, struct rq *rq);
 void scx_kick_cpu(struct scx_sched *sch, s32 cpu, u64 flags);
 void schedule_dsq_reenq(struct scx_sched *sch, struct scx_dispatch_q *dsq,
 			u64 reenq_flags, struct rq *locked_rq);
-int __scx_init_task(struct scx_sched *sch, struct task_struct *p, bool fork);
+int __scx_init_task(struct scx_sched *sch, struct task_struct *p,
+		    struct cgroup *cgrp, bool fork);
 void scx_enable_task(struct scx_sched *sch, struct task_struct *p);
 void __scx_disable_and_exit_task(struct scx_sched *sch, struct task_struct *p);
 void scx_sub_init_cancel_task(struct scx_sched *sch, struct task_struct *p);
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 5f7ac6696d17..393dbd00d2f7 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -906,7 +906,7 @@ void scx_sub_disable(struct scx_sched *sch)
 		 * parent. A child can't directly affect the parent through its
 		 * own failures.
 		 */
-		ret = __scx_init_task(parent, p, false);
+		ret = __scx_init_task(parent, p, NULL, false);
 		if (ret) {
 			scx_fail_parent(sch, p, ret);
 			put_task_struct(p);
@@ -1212,7 +1212,7 @@ void scx_sub_enable_workfn(struct kthread_work *work)
 		 * As $p is still on $parent, it can't be transitioned to INIT.
 		 * Let's worry about task state later. Use __scx_init_task().
 		 */
-		ret = __scx_init_task(sch, p, false);
+		ret = __scx_init_task(sch, p, NULL, false);
 		if (ret)
 			goto abort;
 
@@ -1336,6 +1336,99 @@ void scx_sub_enable_workfn(struct kthread_work *work)
 	cmd->ret = 0;
 }
 
+/**
+ * scx_cgroup_task_migrating - Prepare a task for a cgroup migration
+ * @ctx: migration being prepared
+ *
+ * A task's sched must match its cgroup's owner, so a migration that crosses a
+ * sched boundary re-homes the task once committed. Run the fallible part here,
+ * before the migration commits: initialize the task for the destination sched.
+ * A rejection fails the cgroup.procs write.
+ */
+static s32 scx_cgroup_task_migrating(struct cgroup_task_migrate_ctx *ctx)
+{
+	struct task_struct *p = ctx->task;
+	struct scx_sched *to;
+	int ret;
+
+	/*
+	 * Cleared under scx_cgroup_lock() before root disable starts tearing
+	 * down tasks. As cgroup_mutex is held, a set flag guarantees that the
+	 * teardown loop is not running concurrently.
+	 */
+	if (!scx_cgroup_enabled)
+		return NOTIFY_OK;
+
+	to = ctx->dst_dcgrp->scx_sched;
+	if (scx_task_on_sched(to, p))
+		return NOTIFY_OK;
+
+	ret = __scx_init_task(to, p, ctx->dst_dcgrp, false);
+	if (ret)
+		return notifier_from_errno(ret);
+
+	return NOTIFY_OK;
+}
+
+/**
+ * scx_cgroup_task_migrated - Re-home a task that changed cgroups
+ * @ctx: committed migration
+ *
+ * Move the task to its new cgroup's sched, which scx_cgroup_task_migrating()
+ * already initialized it for. Can't fail.
+ *
+ * This is safe against all phases of the destination sched's destruction. A
+ * disable resets cgroup ownership to the parent and re-homes tasks in one
+ * scx_cgroup_lock() section. If that section already ran, the destination would
+ * be the parent. Otherwise, the re-home loop is still ahead and guaranteed to
+ * visit the task, now in the destination cgroup.
+ */
+static void scx_cgroup_task_migrated(struct cgroup_task_migrate_ctx *ctx)
+{
+	struct task_struct *p = ctx->task;
+	struct scx_sched *to;
+	struct rq *rq;
+	struct rq_flags rf;
+
+	if (!scx_cgroup_enabled)
+		return;
+
+	to = ctx->dst_dcgrp->scx_sched;
+	if (scx_task_on_sched(to, p))
+		return;
+
+	rq = task_rq_lock(p, &rf);
+	scx_rehome_task(to, p);
+	task_rq_unlock(rq, p, &rf);
+}
+
+/**
+ * scx_cgroup_task_migrate_canceled - Undo migration preparation
+ * @ctx: canceled migration
+ *
+ * The migration failed after scx_cgroup_task_migrating() initialized the task
+ * for the destination sched. The task stays on its current sched in the source
+ * cgroup. Undo the destination's init.
+ */
+static void scx_cgroup_task_migrate_canceled(struct cgroup_task_migrate_ctx *ctx)
+{
+	struct task_struct *p = ctx->task;
+	struct scx_sched *to;
+	struct rq *rq;
+	struct rq_flags rf;
+
+	if (!scx_cgroup_enabled)
+		return;
+
+	to = ctx->dst_dcgrp->scx_sched;
+	if (scx_task_on_sched(to, p))
+		return;
+
+	rq = task_rq_lock(p, &rf);
+	scx_sub_init_cancel_task(to, p);
+	task_rq_unlock(rq, p, &rf);
+}
+
 static s32 scx_cgroup_lifetime_notify(struct notifier_block *nb,
 				      unsigned long action, void *data)
 {
@@ -1367,12 +1460,42 @@ static struct notifier_block scx_cgroup_lifetime_nb = {
 	.notifier_call = scx_cgroup_lifetime_notify,
 };
 
-static s32 __init scx_cgroup_lifetime_notifier_init(void)
+static s32 scx_cgroup_task_notify(struct notifier_block *nb,
+				  unsigned long action, void *data)
+{
+	struct cgroup_task_migrate_ctx *ctx = data;
+
+	switch (action) {
+	case CGROUP_TASK_MIGRATING:
+		return scx_cgroup_task_migrating(ctx);
+	case CGROUP_TASK_MIGRATED:
+		scx_cgroup_task_migrated(ctx);
+		break;
+	case CGROUP_TASK_MIGRATE_CANCELED:
+		scx_cgroup_task_migrate_canceled(ctx);
+		break;
+	}
+
+	return NOTIFY_OK;
+}
+
+static struct notifier_block scx_cgroup_task_nb = {
+	.notifier_call = scx_cgroup_task_notify,
+};
+
+static s32 __init scx_cgroup_notifier_init(void)
 {
-	return blocking_notifier_chain_register(&cgroup_lifetime_notifier,
-						&scx_cgroup_lifetime_nb);
+	s32 ret;
+
+	ret = blocking_notifier_chain_register(&cgroup_lifetime_notifier,
+					       &scx_cgroup_lifetime_nb);
+	if (ret)
+		return ret;
+
+	return blocking_notifier_chain_register(&cgroup_task_notifier,
+						&scx_cgroup_task_nb);
 }
-core_initcall(scx_cgroup_lifetime_notifier_init);
+core_initcall(scx_cgroup_notifier_init);
 
 static void scx_pstack_recursion(struct bpf_prog *prog, const char *op)
 {
-- 
2.55.0

[PATCH 5/8] sched_ext: Deliver cgroup ops to each task_group's sched

[Tejun Heo]

With sub-schedulers claiming cgroup subtrees, cgroup ops must be delivered
to each task_group's sched rather than always to root. Add tg->scx.sched to
track which sched initialized the task_group. It is set and cleared together
with SCX_TG_INITED.

Deliver the ops accordingly:

- ops.cgroup_exit() goes to the sched whose ops.cgroup_init() it pairs with.

- ops.cgroup_prep_move/move/cancel_move() go to the task's sched, and only
  for moves that don't re-home the task. A re-homing move is reported
  through the ops.exit_task/init_task() pair instead. The cgroups passed to
  the move ops can be outside the sched's inited set as the cpu controller
  can be coarser than the sub-scheduler topology.

- Knobs of a cgroup belong to the parent, so ops.set_weight/idle/bandwidth()
  go to the parent task_group's sched.

All task_groups currently resolve to the root sched, so no behavior changes
until sub-schedulers start claiming cgroups.

While at it, scx_cgroup_init() is restructured so both paths share the
recording.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/sched/ext.h   |   2 +
 kernel/sched/ext/ext.c      | 157 ++++++++++++++++++++++++++----------
 kernel/sched/ext/internal.h |  34 +++++++-
 3 files changed, 146 insertions(+), 47 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index cce42b21f5f5..a6db5d300f30 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -298,6 +298,8 @@ static inline bool scx_rcu_cpu_stall(const struct cpumask *stalled_mask) { retur
 
 struct scx_task_group {
 #ifdef CONFIG_EXT_GROUP_SCHED
+	struct scx_sched	*sched;
+
 	u32			flags;		/* SCX_TG_* */
 	u32			weight;
 	u64			bw_period_us;
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 6bac68758704..b542900da5a3 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -4339,6 +4339,51 @@ void scx_tg_init(struct task_group *tg)
 	tg->scx.idle = false;
 }
 
+/**
+ * scx_tg_sched - Resolve a task_group's sched
+ * @tg: task_group of interest
+ *
+ * Return the sched that @tg's ops.cgroup_init() succeeded on, %NULL if @tg
+ * isn't inited. An autogroup tg has no cgroup of its own and resolves to the
+ * root sched.
+ *
+ * Safe for callers read-locking the ops rwsem. tg->scx.sched rewrites
+ * write-lock it, and tg on/offline can't overlap such callers as a css's files
+ * are created after online and drained before offline.
+ */
+static struct scx_sched *scx_tg_sched(struct task_group *tg)
+{
+	lockdep_assert(lockdep_is_held(&cgroup_mutex) ||
+		       lockdep_is_held(&scx_cgroup_ops_rwsem));
+
+	if (!tg->css.cgroup)
+		tg = &root_task_group;
+	return tg->scx.sched;
+}
+
+/**
+ * scx_tg_knob_sched - Resolve the sched receiving a task_group's knob updates
+ * @tg: task_group of interest
+ *
+ * Knobs of a cgroup belong to the parent. Deliver the set_* ops to the
+ * parent task_group's sched, which equals @tg's own sched everywhere except
+ * at a sub-scheduler attach point, where the sub's parent sched receives
+ * them.
+ *
+ * The callers sit in @tg's cgroup file writes holding the ops rwsem read
+ * side. That extends scx_tg_sched()'s file-write argument to the parent's
+ * sched read: a parent css outlives its children's files.
+ */
+static struct scx_sched *scx_tg_knob_sched(struct task_group *tg)
+{
+	lockdep_assert(lockdep_is_held(&cgroup_mutex) ||
+		       lockdep_is_held(&scx_cgroup_ops_rwsem));
+
+	if (!tg->css.cgroup || !tg->css.parent)
+		return scx_tg_sched(&root_task_group);
+	return scx_tg_sched(css_tg(tg->css.parent));
+}
+
 int scx_tg_online(struct task_group *tg)
 {
 	struct scx_sched *sch = scx_root;
@@ -4359,8 +4404,10 @@ int scx_tg_online(struct task_group *tg)
 			if (ret)
 				ret = scx_ops_sanitize_err(sch, "cgroup_init", ret);
 		}
-		if (ret == 0)
+		if (ret == 0) {
+			tg->scx.sched = sch;
 			tg->scx.flags |= SCX_TG_ONLINE | SCX_TG_INITED;
+		}
 	} else {
 		tg->scx.flags |= SCX_TG_ONLINE;
 	}
@@ -4370,19 +4417,30 @@ int scx_tg_online(struct task_group *tg)
 
 void scx_tg_offline(struct task_group *tg)
 {
-	struct scx_sched *sch = scx_root;
+	struct scx_sched *sch = tg->scx.sched;
 
 	WARN_ON_ONCE(!(tg->scx.flags & SCX_TG_ONLINE));
 
-	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_exit) &&
-	    (tg->scx.flags & SCX_TG_INITED))
+	/* INITED implies non-NULL @sch, test before SCX_HAS_OP() derefs */
+	if (scx_cgroup_enabled && (tg->scx.flags & SCX_TG_INITED) &&
+	    SCX_HAS_OP(sch, cgroup_exit))
 		SCX_CALL_OP(sch, cgroup_exit, NULL, tg->css.cgroup);
+	tg->scx.sched = NULL;
 	tg->scx.flags &= ~(SCX_TG_ONLINE | SCX_TG_INITED);
 }
 
+/*
+ * @p's sched for the cgroup migration paths. Stable as re-homes happen either
+ * at CGROUP_TASK_MIGRATED of the same migration or under scx_cgroup_lock(),
+ * both while holding cgroup_mutex.
+ */
+static struct scx_sched *scx_cgroup_task_sched(struct task_struct *p)
+{
+	return rcu_dereference_protected(p->scx.sched, lockdep_is_held(&cgroup_mutex));
+}
+
 int scx_cgroup_can_attach(struct cgroup_taskset *tset)
 {
-	struct scx_sched *sch = scx_root;
 	struct cgroup_subsys_state *css;
 	struct task_struct *p;
 	int ret;
@@ -4391,6 +4449,7 @@ int scx_cgroup_can_attach(struct cgroup_taskset *tset)
 		return 0;
 
 	cgroup_taskset_for_each(p, css, tset) {
+		struct scx_sched *sch = scx_cgroup_task_sched(p);
 		struct cgroup *from = tg_cgrp(task_group(p));
 		struct cgroup *to = tg_cgrp(css_tg(css));
 
@@ -4404,11 +4463,22 @@ int scx_cgroup_can_attach(struct cgroup_taskset *tset)
 		if (from == to)
 			continue;
 
+		/*
+		 * The cgroup_move ops are delivered to @p's sched, and only for
+		 * moves that don't re-home @p. A re-homing move changes the dfl
+		 * cgroup's sched and is reported through the
+		 * exit_task/init_task pair that the re-homing generates.
+		 */
+		if (!sch || sch != task_css_set(p)->mg_dst_cset->dfl_cgrp->scx_sched)
+			continue;
+
 		if (SCX_HAS_OP(sch, cgroup_prep_move)) {
 			ret = SCX_CALL_OP_RET(sch, cgroup_prep_move, NULL,
 					      p, from, css->cgroup);
-			if (ret)
+			if (ret) {
+				ret = scx_ops_sanitize_err(sch, "cgroup_prep_move", ret);
 				goto err;
+			}
 		}
 
 		p->scx.cgrp_moving_from = from;
@@ -4418,41 +4488,41 @@ int scx_cgroup_can_attach(struct cgroup_taskset *tset)
 
 err:
 	cgroup_taskset_for_each(p, css, tset) {
-		if (SCX_HAS_OP(sch, cgroup_cancel_move) &&
-		    p->scx.cgrp_moving_from)
+		struct scx_sched *sch = scx_cgroup_task_sched(p);
+
+		/* cgrp_moving_from implies non-NULL @sch, test it first */
+		if (p->scx.cgrp_moving_from && SCX_HAS_OP(sch, cgroup_cancel_move))
 			SCX_CALL_OP(sch, cgroup_cancel_move, NULL,
 				    p, p->scx.cgrp_moving_from, css->cgroup);
 		p->scx.cgrp_moving_from = NULL;
 	}
 
-	return scx_ops_sanitize_err(sch, "cgroup_prep_move", ret);
+	return ret;
 }
 
 void scx_cgroup_move_task(struct task_struct *p)
 {
-	struct scx_sched *sch = scx_root;
+	struct scx_sched *sch;
 
 	if (!scx_cgroup_enabled)
 		return;
 
 	/*
-	 * scx_cgroup_can_attach() sets cgrp_moving_from only when the task's
-	 * cgroup changes. Migration keys off css rather than cgroup identity,
-	 * so it can hand an unchanged-cgroup task here with cgrp_moving_from
-	 * NULL. Nothing to report to the BPF scheduler then, so skip it and
-	 * keep prep_move and move paired. Cgroup ops run on the root sched,
-	 * dispatch on the explicit @sch.
+	 * Migration keys off css rather than cgroup identity, so it can hand an
+	 * unchanged-cgroup task here with cgrp_moving_from NULL. Nothing to
+	 * report to the BPF scheduler then, so skip it and keep prep_move and
+	 * move paired.
 	 */
-	if (SCX_HAS_OP(sch, cgroup_move) && p->scx.cgrp_moving_from)
-		__SCX_CALL_OP_TASK(sch, ops, cgroup_move, task_rq(p),
-				   p, p->scx.cgrp_moving_from,
-				   tg_cgrp(task_group(p)));
+	sch = scx_cgroup_task_sched(p);
+	if (p->scx.cgrp_moving_from && SCX_HAS_OP(sch, cgroup_move))
+		SCX_CALL_OP_TASK(sch, cgroup_move, task_rq(p),
+				 p, p->scx.cgrp_moving_from,
+				 tg_cgrp(task_group(p)));
 	p->scx.cgrp_moving_from = NULL;
 }
 
 void scx_cgroup_cancel_attach(struct cgroup_taskset *tset)
 {
-	struct scx_sched *sch = scx_root;
 	struct cgroup_subsys_state *css;
 	struct task_struct *p;
 
@@ -4460,8 +4530,10 @@ void scx_cgroup_cancel_attach(struct cgroup_taskset *tset)
 		return;
 
 	cgroup_taskset_for_each(p, css, tset) {
-		if (SCX_HAS_OP(sch, cgroup_cancel_move) &&
-		    p->scx.cgrp_moving_from)
+		struct scx_sched *sch = scx_cgroup_task_sched(p);
+
+		/* cgrp_moving_from implies non-NULL @sch, test it first */
+		if (p->scx.cgrp_moving_from && SCX_HAS_OP(sch, cgroup_cancel_move))
 			SCX_CALL_OP(sch, cgroup_cancel_move, NULL,
 				    p, p->scx.cgrp_moving_from, css->cgroup);
 		p->scx.cgrp_moving_from = NULL;
@@ -4473,7 +4545,7 @@ void scx_group_set_weight(struct task_group *tg, unsigned long weight)
 	struct scx_sched *sch;
 
 	percpu_down_read(&scx_cgroup_ops_rwsem);
-	sch = scx_root;
+	sch = scx_tg_knob_sched(tg);
 
 	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_weight) &&
 	    tg->scx.weight != weight)
@@ -4489,7 +4561,7 @@ void scx_group_set_idle(struct task_group *tg, bool idle)
 	struct scx_sched *sch;
 
 	percpu_down_read(&scx_cgroup_ops_rwsem);
-	sch = scx_root;
+	sch = scx_tg_knob_sched(tg);
 
 	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_idle))
 		SCX_CALL_OP(sch, cgroup_set_idle, NULL, tg_cgrp(tg), idle);
@@ -4506,7 +4578,7 @@ void scx_group_set_bandwidth(struct task_group *tg,
 	struct scx_sched *sch;
 
 	percpu_down_read(&scx_cgroup_ops_rwsem);
-	sch = scx_root;
+	sch = scx_tg_knob_sched(tg);
 
 	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_bandwidth) &&
 	    (tg->scx.bw_period_us != period_us ||
@@ -4718,6 +4790,7 @@ static void scx_cgroup_exit(struct scx_sched *sch)
 
 		if (!(tg->scx.flags & SCX_TG_INITED))
 			continue;
+		tg->scx.sched = NULL;
 		tg->scx.flags &= ~SCX_TG_INITED;
 
 		if (!sch->ops.cgroup_exit)
@@ -4738,28 +4811,26 @@ static int scx_cgroup_init(struct scx_sched *sch)
 	 */
 	css_for_each_descendant_pre(css, &root_task_group.css) {
 		struct task_group *tg = css_tg(css);
-		struct scx_cgroup_init_args args = {
-			.weight = tg->scx.weight,
-			.bw_period_us = tg->scx.bw_period_us,
-			.bw_quota_us = tg->scx.bw_quota_us,
-			.bw_burst_us = tg->scx.bw_burst_us,
-		};
 
-		if ((tg->scx.flags &
-		     (SCX_TG_ONLINE | SCX_TG_INITED)) != SCX_TG_ONLINE)
+		if ((tg->scx.flags & (SCX_TG_ONLINE | SCX_TG_INITED)) != SCX_TG_ONLINE)
 			continue;
 
-		if (!sch->ops.cgroup_init) {
-			tg->scx.flags |= SCX_TG_INITED;
-			continue;
-		}
+		if (sch->ops.cgroup_init) {
+			struct scx_cgroup_init_args args = {
+				.weight = tg->scx.weight,
+				.bw_period_us = tg->scx.bw_period_us,
+				.bw_quota_us = tg->scx.bw_quota_us,
+				.bw_burst_us = tg->scx.bw_burst_us,
+			};
 
-		ret = SCX_CALL_OP_RET(sch, cgroup_init, NULL,
-				      css->cgroup, &args);
-		if (ret) {
-			scx_error(sch, "ops.cgroup_init() failed (%d)", ret);
-			return ret;
+			ret = SCX_CALL_OP_RET(sch, cgroup_init, NULL, css->cgroup, &args);
+			if (ret) {
+				scx_error(sch, "ops.cgroup_init() failed (%d)", ret);
+				return ret;
+			}
 		}
+
+		tg->scx.sched = sch;
 		tg->scx.flags |= SCX_TG_INITED;
 	}
 
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 23fc95502ea6..ad90d4645a33 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -645,8 +645,14 @@ struct sched_ext_ops {
 	 * @cgrp: cgroup being initialized
 	 * @args: init arguments, see the struct definition
 	 *
-	 * Either the BPF scheduler is being loaded or @cgrp created, initialize
-	 * @cgrp for sched_ext. This operation may block.
+	 * Initialize @cgrp for sched_ext, delivered to @cgrp's sched either
+	 * when the BPF scheduler is being loaded or when @cgrp is created. This
+	 * operation may block.
+	 *
+	 * When the BPF scheduler is being loaded or cgroups are being handed
+	 * over, @cgrp may already have been removed by userspace: a removed
+	 * cgroup stays schedulable until its dying tasks finish their final
+	 * context switches.
 	 *
 	 * Return 0 for success, -errno for failure. An error return while
 	 * loading will abort loading of the BPF scheduler. During cgroup
@@ -659,8 +665,13 @@ struct sched_ext_ops {
 	 * @cgroup_exit: Exit a cgroup
 	 * @cgrp: cgroup being exited
 	 *
-	 * Either the BPF scheduler is being unloaded or @cgrp destroyed, exit
-	 * @cgrp for sched_ext. This operation my block.
+	 * Exit @cgrp for sched_ext, delivered to the sched whose
+	 * ops.cgroup_init() it pairs with, either when the BPF scheduler is
+	 * being unloaded or when @cgrp is destroyed. This operation may block.
+	 *
+	 * For a destroyed @cgrp, delivery follows the last scheduling event on
+	 * it: a removed cgroup stays schedulable until its dying tasks finish
+	 * their final context switches.
 	 */
 	void (*cgroup_exit)(struct cgroup *cgrp);
 
@@ -673,6 +684,12 @@ struct sched_ext_ops {
 	 * Prepare @p for move from cgroup @from to @to. This operation may
 	 * block and can be used for allocations.
 	 *
+	 * The cgroup_move ops are delivered to @p's sched, and only for moves
+	 * that don't re-home @p. A re-homing move is reported through
+	 * ops.exit_task() and ops.init_task() instead. @from and @to can
+	 * reference cgroups the sched never received ops.cgroup_init() for, as
+	 * the cpu controller can be coarser than the sub-scheduler topology.
+	 *
 	 * Return 0 for success, -errno for failure. An error return aborts the
 	 * migration.
 	 */
@@ -708,6 +725,11 @@ struct sched_ext_ops {
 	 * @weight: new weight [1..10000]
 	 *
 	 * Update @cgrp's weight to @weight.
+	 *
+	 * Knobs of a cgroup belong to the parent, so the set_* ops are
+	 * delivered to @cgrp's parent's sched. That sched may never have seen
+	 * ops.cgroup_init() for @cgrp - at a sub-scheduler attach point, the
+	 * parent sched tracks @cgrp through ops.sub_attach() instead.
 	 */
 	void (*cgroup_set_weight)(struct cgroup *cgrp, u32 weight);
 
@@ -728,6 +750,8 @@ struct sched_ext_ops {
 	 * burst temporarily. The specific control mechanism and thus the
 	 * interpretation of @period_us and burstiness is up to the BPF
 	 * scheduler.
+	 *
+	 * Delivery follows the same rule as cgroup_set_weight().
 	 */
 	void (*cgroup_set_bandwidth)(struct cgroup *cgrp,
 				     u64 period_us, u64 quota_us, u64 burst_us);
@@ -740,6 +764,8 @@ struct sched_ext_ops {
 	 * Update @cgrp's idle state to @idle. This callback is invoked when
 	 * a cgroup transitions between idle and non-idle states, allowing the
 	 * BPF scheduler to adjust its behavior accordingly.
+	 *
+	 * Delivery follows the same rule as cgroup_set_weight().
 	 */
 	void (*cgroup_set_idle)(struct cgroup *cgrp, bool idle);
 
-- 
2.55.0

[PATCH 6/8] sched_ext: Hand over cgroups at sub-scheduler enable/disable

[Tejun Heo]

Sub-schedulers don't get cgroups yet: every task_group is inited on the root
sched and the routing added by the previous patches always resolves to it.
Add the handover: an enabling sub-scheduler takes over the cgroups in its
subtree and a disabling one returns them to its parent.

scx_cgroup_claim_subtree() runs while the sub enables, after the subtree's
cgrp->scx_sched's are set and before any task is claimed. It inits each
subtree task_group on the sub, exits it from the parent and updates
tg->scx.sched. A failed ops.cgroup_init() unwinds the sub-side inits and
aborts the enable with the parent untouched.

Disabling reverses it with scx_cgroup_return_subtree(): exit each cgroup
from the sub, then re-init it on the parent with the current tg->scx.*
values, resyncing weight and bandwidth changes made while the sub had it.
When a re-init fails, the parent is failed and the remaining task_groups
still transfer uninited and get no cgroup ops - the same punting done for
tasks. The dying parent's own disable moves them onward.

The handover walks include dying but not yet offlined task_groups, the same
as root's bulk walks: a removed cgroup keeps hosting scheduling events until
its dying tasks finish their final context switches, and its
ops.cgroup_exit() must follow the last of them. tg on/offlining is excluded
through cgroup_lock(), so either ordering against an rmdir of a subtree
cgroup delivers balanced init/exit pairs.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 include/linux/sched/ext.h   |  13 +++
 kernel/sched/ext/ext.c      |  45 +++++---
 kernel/sched/ext/internal.h |   6 ++
 kernel/sched/ext/sub.c      | 199 ++++++++++++++++++++++++++++++++++++
 4 files changed, 250 insertions(+), 13 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index a6db5d300f30..78b2f289cb98 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -298,6 +298,19 @@ static inline bool scx_rcu_cpu_stall(const struct cpumask *stalled_mask) { retur
 
 struct scx_task_group {
 #ifdef CONFIG_EXT_GROUP_SCHED
+	/*
+	 * The sched this tg is on, NULL if none. SCX_TG_INITED tracks whether
+	 * ops.cgroup_init() succeeded on it. When a child sched exits and its
+	 * tgs move to the parent, a failed init leaves the tg on the parent
+	 * with INITED clear (see scx_cgroup_return_subtree()).
+	 *
+	 * This is tracked separately from cgrp->scx_sched because the tg
+	 * hierarchy can diverge from the cgroup2 hierarchy in both lifetime and
+	 * shape. A tg stays online past its cgroup's removal while the
+	 * cgrp->scx_sched rewrites visit only live cgroups, leaving a removed
+	 * cgroup's pointer stale. The cpu controller can also be mounted on
+	 * cgroup1.
+	 */
 	struct scx_sched	*sched;
 
 	u32			flags;		/* SCX_TG_* */
diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index b542900da5a3..58cd971e5fc5 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -4347,6 +4347,11 @@ void scx_tg_init(struct task_group *tg)
  * isn't inited. An autogroup tg has no cgroup of its own and resolves to the
  * root sched.
  *
+ * When a child sched exits, its task_groups are moved to the parent and
+ * re-inited on it. A failed re-init fails the parent in turn and leaves the
+ * task_group without a sched it's inited on, resolving to %NULL. See
+ * scx_cgroup_return_subtree().
+ *
  * Safe for callers read-locking the ops rwsem. tg->scx.sched rewrites
  * write-lock it, and tg on/offline can't overlap such callers as a css's files
  * are created after online and drained before offline.
@@ -4358,7 +4363,8 @@ static struct scx_sched *scx_tg_sched(struct task_group *tg)
 
 	if (!tg->css.cgroup)
 		tg = &root_task_group;
-	return tg->scx.sched;
+	/* INITED means ops.cgroup_init() succeeded on @tg->scx.sched */
+	return (tg->scx.flags & SCX_TG_INITED) ? tg->scx.sched : NULL;
 }
 
 /**
@@ -4370,6 +4376,9 @@ static struct scx_sched *scx_tg_sched(struct task_group *tg)
  * at a sub-scheduler attach point, where the sub's parent sched receives
  * them.
  *
+ * Return %NULL if the parent task_group has no sched. That can happen when the
+ * parent's ops.cgroup_init() fails while a sub-scheduler is being disabled.
+ *
  * The callers sit in @tg's cgroup file writes holding the ops rwsem read
  * side. That extends scx_tg_sched()'s file-write argument to the parent's
  * sched read: a parent css outlives its children's files.
@@ -4386,12 +4395,24 @@ static struct scx_sched *scx_tg_knob_sched(struct task_group *tg)
 
 int scx_tg_online(struct task_group *tg)
 {
-	struct scx_sched *sch = scx_root;
 	int ret = 0;
 
 	WARN_ON_ONCE(tg->scx.flags & (SCX_TG_ONLINE | SCX_TG_INITED));
 
 	if (scx_cgroup_enabled) {
+		struct scx_sched *sch;
+
+		/*
+		 * The cgroup lifetime notifier populates cgrp->scx_sched before
+		 * css_online, but only on the default hierarchy. Sub-scheds are
+		 * attached to the cgroup2 hierarchy, so a cgroup1 task_group
+		 * always belongs to the root sched.
+		 */
+		if (cgroup_on_dfl(tg->css.cgroup))
+			sch = tg->css.cgroup->scx_sched;
+		else
+			sch = scx_tg_sched(&root_task_group);
+
 		if (SCX_HAS_OP(sch, cgroup_init)) {
 			struct scx_cgroup_init_args args =
 				{ .weight = tg->scx.weight,
@@ -4547,7 +4568,7 @@ void scx_group_set_weight(struct task_group *tg, unsigned long weight)
 	percpu_down_read(&scx_cgroup_ops_rwsem);
 	sch = scx_tg_knob_sched(tg);
 
-	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_weight) &&
+	if (scx_cgroup_enabled && sch && SCX_HAS_OP(sch, cgroup_set_weight) &&
 	    tg->scx.weight != weight)
 		SCX_CALL_OP(sch, cgroup_set_weight, NULL, tg_cgrp(tg), weight);
 
@@ -4563,7 +4584,7 @@ void scx_group_set_idle(struct task_group *tg, bool idle)
 	percpu_down_read(&scx_cgroup_ops_rwsem);
 	sch = scx_tg_knob_sched(tg);
 
-	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_idle))
+	if (scx_cgroup_enabled && sch && SCX_HAS_OP(sch, cgroup_set_idle))
 		SCX_CALL_OP(sch, cgroup_set_idle, NULL, tg_cgrp(tg), idle);
 
 	/* Update the task group's idle state */
@@ -4580,7 +4601,7 @@ void scx_group_set_bandwidth(struct task_group *tg,
 	percpu_down_read(&scx_cgroup_ops_rwsem);
 	sch = scx_tg_knob_sched(tg);
 
-	if (scx_cgroup_enabled && SCX_HAS_OP(sch, cgroup_set_bandwidth) &&
+	if (scx_cgroup_enabled && sch && SCX_HAS_OP(sch, cgroup_set_bandwidth) &&
 	    (tg->scx.bw_period_us != period_us ||
 	     tg->scx.bw_quota_us != quota_us ||
 	     tg->scx.bw_burst_us != burst_us))
@@ -4788,15 +4809,13 @@ static void scx_cgroup_exit(struct scx_sched *sch)
 	css_for_each_descendant_post(css, &root_task_group.css) {
 		struct task_group *tg = css_tg(css);
 
-		if (!(tg->scx.flags & SCX_TG_INITED))
-			continue;
+		/* also clear the sched of tgs whose ops.cgroup_init() failed */
 		tg->scx.sched = NULL;
-		tg->scx.flags &= ~SCX_TG_INITED;
-
-		if (!sch->ops.cgroup_exit)
-			continue;
-
-		SCX_CALL_OP(sch, cgroup_exit, NULL, css->cgroup);
+		if (tg->scx.flags & SCX_TG_INITED) {
+			tg->scx.flags &= ~SCX_TG_INITED;
+			if (sch->ops.cgroup_exit)
+				SCX_CALL_OP(sch, cgroup_exit, NULL, css->cgroup);
+		}
 	}
 }
 
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index ad90d4645a33..0c0a8fdaa2c4 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -649,6 +649,11 @@ struct sched_ext_ops {
 	 * when the BPF scheduler is being loaded or when @cgrp is created. This
 	 * operation may block.
 	 *
+	 * Cgroup handovers also generate these ops: an enabling sub-scheduler
+	 * receives ops.cgroup_init() for every cgroup in its subtree while the
+	 * previous sched receives ops.cgroup_exit(), and disabling reverses the
+	 * two.
+	 *
 	 * When the BPF scheduler is being loaded or cgroups are being handed
 	 * over, @cgrp may already have been removed by userspace: a removed
 	 * cgroup stays schedulable until its dying tasks finish their final
@@ -1722,6 +1727,7 @@ enum scx_kick_flags {
 enum scx_tg_flags {
 	SCX_TG_ONLINE		= 1U << 0,
 	SCX_TG_INITED		= 1U << 1,
+	SCX_TG_SUB_INIT		= 1U << 2,	/* see scx_cgroup_claim_subtree() */
 };
 
 enum scx_enable_state {
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 393dbd00d2f7..8d8737149bc0 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -836,6 +836,191 @@ static void scx_fail_parent(struct scx_sched *sch,
 	scx_task_iter_stop(&sti);
 }
 
+#ifdef CONFIG_EXT_GROUP_SCHED
+/**
+ * scx_cgroup_claim_subtree - Claim the subtree's cgroups for an enabling sub
+ * @sch: sub-scheduler being enabled
+ *
+ * Called while enabling @sch, after the subtree's cgrp->scx_sched's are pointed
+ * at @sch and before any task is claimed. This mirrors root enable's
+ * cgroups-before-tasks order. The ops.init_task() args are task_group-granular
+ * and can still reference a cgroup outside the handed-over set when the cpu
+ * controller is coarser than the sub topology or mounted on cgroup1.
+ *
+ * First init each of the parent sched's subtree cgroups on @sch, and only then
+ * exit them from the parent, so that a failed init can be unwound with the
+ * parent untouched. The both-inited transient is invisible outside
+ * scx_cgroup_lock(). %SCX_TG_SUB_INIT tracks the first pass's progress.
+ * %SCX_TG_INITED stays set throughout, except for a task_group whose
+ * ops.cgroup_init() failed on the parent (see scx_cgroup_return_subtree()):
+ * there is nothing to exit from the parent and %SCX_TG_INITED is set back with
+ * the transfer.
+ *
+ * Dying but not yet offlined task_groups are included: a removed cgroup keeps
+ * hosting scheduling events until its dying tasks finish their final context
+ * switches, so it still needs to be inited on a sched, and its offline-time
+ * ops.cgroup_exit() follows the last of those events.
+ *
+ * Return 0 on success, -errno on failure. On failure, @sch has been
+ * scx_error()'d and is left with no cgroups.
+ */
+static s32 scx_cgroup_claim_subtree(struct scx_sched *sch)
+{
+	struct cgroup *sub_cgrp = sch_cgroup(sch);
+	struct cgroup_subsys_state *ecss = cgroup_e_css(sub_cgrp, &cpu_cgrp_subsys);
+	struct scx_sched *parent = scx_parent(sch);
+	struct cgroup_subsys_state *css;
+	int ret;
+
+	css_for_each_descendant_pre(css, ecss) {
+		struct task_group *tg = css_tg(css);
+		struct scx_cgroup_init_args args = {
+			.weight = tg->scx.weight,
+			.bw_period_us = tg->scx.bw_period_us,
+			.bw_quota_us = tg->scx.bw_quota_us,
+			.bw_burst_us = tg->scx.bw_burst_us,
+		};
+
+		if (tg->scx.sched != parent ||
+		    !cgroup_is_descendant(css->cgroup, sub_cgrp))
+			continue;
+
+		if (SCX_HAS_OP(sch, cgroup_init)) {
+			ret = SCX_CALL_OP_RET(sch, cgroup_init, NULL, css->cgroup, &args);
+			if (ret) {
+				scx_error(sch, "ops.cgroup_init() failed (%d)", ret);
+				goto err;
+			}
+		}
+		tg->scx.flags |= SCX_TG_SUB_INIT;
+	}
+
+	css_for_each_descendant_post(css, ecss) {
+		struct task_group *tg = css_tg(css);
+
+		/*
+		 * SUB_INIT is pass 1's progress mark: pass 2 and the err path
+		 * must visit exactly the tgs pass 1 inited.
+		 */
+		if (!(tg->scx.flags & SCX_TG_SUB_INIT))
+			continue;
+
+		/* skip the exit if the parent's ops.cgroup_init() failed */
+		if ((tg->scx.flags & SCX_TG_INITED) && SCX_HAS_OP(parent, cgroup_exit))
+			SCX_CALL_OP(parent, cgroup_exit, NULL, css->cgroup);
+		tg->scx.sched = sch;
+		tg->scx.flags |= SCX_TG_INITED;
+		tg->scx.flags &= ~SCX_TG_SUB_INIT;
+	}
+
+	return 0;
+
+err:
+	css_for_each_descendant_post(css, ecss) {
+		struct task_group *tg = css_tg(css);
+
+		if (!(tg->scx.flags & SCX_TG_SUB_INIT))
+			continue;
+
+		if (SCX_HAS_OP(sch, cgroup_exit))
+			SCX_CALL_OP(sch, cgroup_exit, NULL, css->cgroup);
+		tg->scx.flags &= ~SCX_TG_SUB_INIT;
+	}
+	return ret;
+}
+
+/**
+ * scx_cgroup_return_subtree - Return the subtree's cgroups to the parent sched
+ * @sch: sub-scheduler being disabled
+ *
+ * Called while disabling @sch, after the subtree's cgrp->scx_sched's are reset
+ * to the parent sched and before tasks are re-homed, mirroring root disable's
+ * cgroups-before-tasks teardown order. The reverse of
+ * scx_cgroup_claim_subtree(): exit @sch's cgroups from @sch, then init them on
+ * the parent with the current tg->scx.* values, resyncing settings that changed
+ * while @sch had them.
+ *
+ * When an init on the parent fails, the parent is failed - the same policy as
+ * task re-homing. The remaining task_groups are punted: they move to the parent
+ * anyway with %SCX_TG_INITED cleared, as ops.cgroup_init() failed or never ran
+ * for them. A punted task_group gets no cgroup ops. The dying parent's own
+ * disable moves it one sched up, initing it there. Root ends the chain: root
+ * teardown drops cgroup ops entirely and the next enable's bulk init re-inits
+ * every online task_group.
+ *
+ * The task re-home that follows still delivers ops.init_task() to the dying
+ * parent, including for tasks in punted cgroups it never inited - tolerated
+ * like the downstream failures of task punting (see scx_punt_task()).
+ */
+static void scx_cgroup_return_subtree(struct scx_sched *sch)
+{
+	struct cgroup *sub_cgrp = sch_cgroup(sch);
+	struct cgroup_subsys_state *ecss = cgroup_e_css(sub_cgrp, &cpu_cgrp_subsys);
+	struct scx_sched *parent = scx_parent(sch);
+	struct cgroup_subsys_state *css;
+	bool parent_failed = false;
+	int ret;
+
+	css_for_each_descendant_post(css, ecss) {
+		struct task_group *tg = css_tg(css);
+
+		if (tg->scx.sched != sch ||
+		    !cgroup_is_descendant(css->cgroup, sub_cgrp))
+			continue;
+
+		/* skip the exit if @sch's ops.cgroup_init() failed for the tg */
+		if ((tg->scx.flags & SCX_TG_INITED) && SCX_HAS_OP(sch, cgroup_exit))
+			SCX_CALL_OP(sch, cgroup_exit, NULL, css->cgroup);
+		tg->scx.sched = parent;
+		tg->scx.flags |= SCX_TG_SUB_INIT;
+	}
+
+	css_for_each_descendant_pre(css, ecss) {
+		struct task_group *tg = css_tg(css);
+		struct scx_cgroup_init_args args = {
+			.weight = tg->scx.weight,
+			.bw_period_us = tg->scx.bw_period_us,
+			.bw_quota_us = tg->scx.bw_quota_us,
+			.bw_burst_us = tg->scx.bw_burst_us,
+		};
+
+		/* the first pass must have transferred everything */
+		WARN_ON_ONCE(tg->scx.sched == sch);
+
+		/*
+		 * SUB_INIT distinguishes the tgs pass 1 moved. The sched test
+		 * can't: a tg punted to the parent by an earlier failure would
+		 * also match.
+		 */
+		if (!(tg->scx.flags & SCX_TG_SUB_INIT))
+			continue;
+		tg->scx.flags &= ~(SCX_TG_SUB_INIT | SCX_TG_INITED);
+
+		/*
+		 * A re-init on $parent failed. The task_groups from here on are
+		 * punted: they stay on the dying $parent with INITED clear and
+		 * move onward when it disables.
+		 */
+		if (parent_failed)
+			continue;
+
+		if (SCX_HAS_OP(parent, cgroup_init)) {
+			ret = SCX_CALL_OP_RET(parent, cgroup_init, NULL, css->cgroup, &args);
+			if (ret) {
+				scx_error(parent, "ops.cgroup_init() failed (%d) while disabling a sub-scheduler",
+					  ret);
+				parent_failed = true;
+				continue;
+			}
+		}
+		tg->scx.flags |= SCX_TG_INITED;
+	}
+}
+#else
+static inline s32 scx_cgroup_claim_subtree(struct scx_sched *sch) { return 0; }
+static inline void scx_cgroup_return_subtree(struct scx_sched *sch) {}
+#endif
+
 void scx_sub_disable(struct scx_sched *sch)
 {
 	struct scx_sched *parent = scx_parent(sch);
@@ -871,6 +1056,12 @@ void scx_sub_disable(struct scx_sched *sch)
 
 	set_cgroup_sched(sch_cgroup(sch), parent);
 
+	/*
+	 * Return the subtree's cgroups before re-homing tasks so that any
+	 * ops.init_task() on $parent only sees cgroups it has initialized.
+	 */
+	scx_cgroup_return_subtree(sch);
+
 	scx_task_iter_start(&sti, sch->cgrp);
 	while ((p = scx_task_iter_next_locked(&sti))) {
 		struct rq *rq;
@@ -1172,6 +1363,14 @@ void scx_sub_enable_workfn(struct kthread_work *work)
 		goto err_unlock_and_disable;
 	}
 
+	/*
+	 * Take over the subtree's cgroups before any task is claimed,
+	 * mirroring root enable's cgroups-before-tasks order.
+	 */
+	ret = scx_cgroup_claim_subtree(sch);
+	if (ret)
+		goto err_unlock_and_disable;
+
 	/*
 	 * Initialize tasks for the new child $sch without exiting them for
 	 * $parent so that the tasks can always be reverted back to $parent
-- 
2.55.0

[PATCH 7/8] tools/sched_ext: scx_qmap - Consume cgroup weights through set_weight

[Tejun Heo]

With the set_* ops delivered to the parent's sched, a parent qmap instance
now receives cgroup_set_weight for its child subs' attach points. Update
the matching sub_sched_ctx weight and redistribute() in-kernel, and drop
the userspace feed_weights() polling. This exercises the knob routing end
to end.

The self weight is fixed at 100: a cgroup's weight is its parent's knob and
not the scheduler's own business. This drops the self-weight polling and the
repartition PROG_RUN poke with it.

sub_attach seeds the slot with the cgroup's current weight, read through
bpf_cgroup_from_id(), so a weight set before the sub attaches is picked up.
A write racing the attach can still be lost until the next value-changing
cpu.weight write. Acceptable for a demo.

While at it, add a traced ops.cgroup_move() so tests can observe move
delivery.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 tools/sched_ext/scx_qmap.bpf.c |  80 +++++++++++++++++++++-----
 tools/sched_ext/scx_qmap.c     | 102 +--------------------------------
 tools/sched_ext/scx_qmap.h     |   3 +-
 3 files changed, 69 insertions(+), 116 deletions(-)

diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 09aee49120c2..8b8b14309466 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -1019,10 +1019,32 @@ s32 BPF_STRUCT_OPS(qmap_cgroup_init, struct cgroup *cgrp, struct scx_cgroup_init
 	return 0;
 }
 
+static void redistribute(void);
+
 void BPF_STRUCT_OPS(qmap_cgroup_set_weight, struct cgroup *cgrp, u32 weight)
 {
+	u64 cgid = cgrp->kn->id;
+	s32 i;
+
+	QMAP_TOUCH_ARENA();
+
 	if (print_msgs)
-		bpf_printk("CGRP SET %llu weight=%u", cgrp->kn->id, weight);
+		bpf_printk("CGRP SET %llu weight=%u", cgid, weight);
+
+	/*
+	 * Knobs belong to the parent, so this op carries the child subs'
+	 * attach point weights. Adjust the matching sub's share of the cid
+	 * partition. Other cgroups don't participate in the split.
+	 */
+	for (i = 0; i < MAX_SUB_SCHEDS; i++) {
+		if (qa.sub_sched_ctxs[i].cgroup_id != cgid)
+			continue;
+		if (qa.sub_sched_ctxs[i].weight != weight) {
+			qa.sub_sched_ctxs[i].weight = weight;
+			redistribute();
+		}
+		break;
+	}
 }
 
 void BPF_STRUCT_OPS(qmap_cgroup_set_bandwidth, struct cgroup *cgrp,
@@ -1033,6 +1055,14 @@ void BPF_STRUCT_OPS(qmap_cgroup_set_bandwidth, struct cgroup *cgrp,
 			   cgrp->kn->id, period_us, quota_us, burst_us);
 }
 
+void BPF_STRUCT_OPS(qmap_cgroup_move, struct task_struct *p,
+		    struct cgroup *from, struct cgroup *to)
+{
+	if (print_msgs)
+		bpf_printk("CGRP MOVE %d %llu -> %llu",
+			   p->pid, from->kn->id, to->kn->id);
+}
+
 void BPF_STRUCT_OPS(qmap_update_idle, s32 cid, bool idle)
 {
 	QMAP_TOUCH_ARENA();
@@ -1336,10 +1366,12 @@ __noinline void compute_partition(void)
 	}
 
 	/*
-	 * Snapshot membership and weights so the sum_w and share loops agree.
-	 * A mid-compute change would otherwise wrap nr_shared negative.
+	 * Snapshot membership and weights so the sum_w and share loops agree. A
+	 * mid-compute change would otherwise wrap nr_shared negative. The self
+	 * weight is fixed at the default: a cgroup's weight is its parent's
+	 * knob, not the scheduler's own business.
 	 */
-	self_w = qa.self_weight ?: 100;
+	self_w = 100;
 	bpf_for(i, 0, MAX_SUB_SCHEDS) {
 		cgid_snap[i] = qa.sub_sched_ctxs[i].cgroup_id;
 		w_snap[i] = cgid_snap[i] ? (qa.sub_sched_ctxs[i].weight ?: 100) : 0;
@@ -1565,7 +1597,7 @@ __noinline void apply_partition(void)
 
 /*
  * Recompute the split off the node's held caps and apply it. The contexts this
- * runs from (the sub-sched callbacks, the userspace poke, the rr timer) are not
+ * runs from (the sub-sched and cgroup callbacks, the rr timer) are not
  * serialized by the kernel, so a single runner does the work. A caller that
  * finds the guard held leaves part_pending set; the holder drains it before
  * releasing, with the rr timer as a backstop.
@@ -1592,14 +1624,6 @@ static void redistribute(void)
 	part_end();
 }
 
-/* userspace pokes this (PROG_RUN) to resplit after a cpu.weight change */
-SEC("syscall")
-int repartition(void *ctx)
-{
-	redistribute();
-	return 0;
-}
-
 /*
  * Userspace pokes this (PROG_RUN) to bring alloc_ns[] current before reading
  * it for the stats display. Skipping when the partition guard is held is
@@ -1849,6 +1873,33 @@ void BPF_STRUCT_OPS(qmap_exit, struct scx_exit_info *ei)
 	UEI_RECORD(uei, ei);
 }
 
+/*
+ * Seed a new sub slot with the cgroup's current weight. The kernel delivers
+ * ops.cgroup_set_weight() only on value-changing writes, so a weight set
+ * before the sub attached would otherwise go unnoticed.
+ */
+static u32 cgrp_cur_weight(u64 cgid)
+{
+	struct cgroup_subsys_state *css;
+	struct cgroup *cgrp;
+	u32 weight = 100;
+
+	cgrp = bpf_cgroup_from_id(cgid);
+	if (!cgrp)
+		return weight;
+
+	css = BPF_CORE_READ(cgrp, subsys[cpu_cgrp_id]);
+	if (css) {
+		struct task_group *tg = container_of(css, struct task_group, css);
+		u32 w = BPF_CORE_READ(tg, scx.weight);
+
+		if (w)
+			weight = w;
+	}
+	bpf_cgroup_release(cgrp);
+	return weight;
+}
+
 s32 BPF_STRUCT_OPS(qmap_sub_attach, struct scx_sub_attach_args *args)
 {
 	s32 i;
@@ -1862,7 +1913,7 @@ s32 BPF_STRUCT_OPS(qmap_sub_attach, struct scx_sub_attach_args *args)
 			continue;
 
 		qa.sub_sched_ctxs[i].cgroup_id = args->ops->sub_cgroup_id;
-		qa.sub_sched_ctxs[i].weight = 100;	/* until userspace feeds it */
+		qa.sub_sched_ctxs[i].weight = cgrp_cur_weight(args->ops->sub_cgroup_id);
 		qa.nr_sub_scheds++;
 		bpf_printk("attaching sub-sched[%d] on %s", i, args->cgroup_path);
 		redistribute();
@@ -1926,6 +1977,7 @@ SCX_OPS_CID_DEFINE(qmap_ops,
 	       .cgroup_init		= (void *)qmap_cgroup_init,
 	       .cgroup_set_weight	= (void *)qmap_cgroup_set_weight,
 	       .cgroup_set_bandwidth	= (void *)qmap_cgroup_set_bandwidth,
+	       .cgroup_move		= (void *)qmap_cgroup_move,
 	       .sub_attach		= (void *)qmap_sub_attach,
 	       .sub_detach		= (void *)qmap_sub_detach,
 	       .sub_caps_updated	= (void *)qmap_sub_caps_updated,
diff --git a/tools/sched_ext/scx_qmap.c b/tools/sched_ext/scx_qmap.c
index dda3ddf5b749..105745f51b8e 100644
--- a/tools/sched_ext/scx_qmap.c
+++ b/tools/sched_ext/scx_qmap.c
@@ -24,11 +24,6 @@
 #include "scx_qmap.h"
 #include "scx_qmap.bpf.skel.h"
 
-/* kernfs file-handle type for open_by_handle_at(), from linux/exportfs.h */
-#ifndef FILEID_KERNFS
-#define FILEID_KERNFS	0xfe
-#endif
-
 const char help_fmt[] =
 "A simple five-level FIFO queue sched_ext scheduler.\n"
 "\n"
@@ -72,7 +67,7 @@ const char help_fmt[] =
 "  -I            Turn on SCX_OPS_ALWAYS_ENQ_IMMED\n"
 "  -F COUNT      IMMED stress: force every COUNT'th enqueue to a busy local DSQ (use with -I)\n"
 "  -C MODE       cid-override test (shuffle|bad-dup|bad-range|bad-mono)\n"
-"  -i SEC        Stats and weight-refresh interval, seconds (default 5)\n"
+"  -i SEC        Stats interval, seconds (default 5)\n"
 "  -R MS         Round-robin period for time-shared cpus, ms (default 200)\n"
 "  -J MODE       Fault injection (wrong-cid: dispatch to a cid not held)\n"
 "  -v            Print libbpf debug messages\n"
@@ -93,83 +88,6 @@ static void sigint_handler(int dummy)
 	exit_req = 1;
 }
 
-/*
- * Open a cgroup directory directly from its id. In cgroup2 the cgroup id is the
- * kernfs node id, so a FILEID_KERNFS handle built from the id resolves to the
- * directory via open_by_handle_at() against the cgroup mount.
- */
-static int open_cgroup_by_id(u64 cgid)
-{
-	static int mnt_fd = -1;
-	struct {
-		struct file_handle fh;
-		u64 id;
-	} h;
-
-	if (mnt_fd < 0) {
-		mnt_fd = open("/sys/fs/cgroup", O_RDONLY | O_DIRECTORY);
-		if (mnt_fd < 0)
-			return -1;
-	}
-	h.fh.handle_bytes = sizeof(h.id);
-	h.fh.handle_type = FILEID_KERNFS;
-	h.id = cgid;
-	return open_by_handle_at(mnt_fd, &h.fh, O_RDONLY | O_DIRECTORY);
-}
-
-/* read a cgroup's cpu.weight (1-10000) by id, 0 if unavailable */
-static u32 read_cgroup_weight(u64 cgid)
-{
-	char buf[32];
-	int dfd, wfd;
-	u32 w = 0;
-	ssize_t n;
-
-	dfd = open_cgroup_by_id(cgid);
-	if (dfd < 0)
-		return 0;
-	wfd = openat(dfd, "cpu.weight", O_RDONLY);
-	close(dfd);
-	if (wfd < 0)
-		return 0;
-	n = read(wfd, buf, sizeof(buf) - 1);
-	close(wfd);
-	if (n > 0) {
-		buf[n] = '\0';
-		w = strtoul(buf, NULL, 10);
-	}
-	return w;
-}
-
-/* read each direct child's cpu.weight into the arena, true if any changed */
-static bool feed_weights(struct qmap_arena *qa)
-{
-	bool changed = false;
-	int i;
-
-	for (i = 0; i < MAX_SUB_SCHEDS; i++) {
-		u64 cgid = qa->sub_sched_ctxs[i].cgroup_id;
-		u32 w;
-
-		if (!cgid)
-			continue;
-		/* racy against slot reuse but weight is advisory and self-corrects */
-		w = read_cgroup_weight(cgid);
-		if (w && w != qa->sub_sched_ctxs[i].weight) {
-			qa->sub_sched_ctxs[i].weight = w;
-			changed = true;
-		}
-	}
-	return changed;
-}
-
-static void invoke_repartition(struct scx_qmap *skel)
-{
-	LIBBPF_OPTS(bpf_test_run_opts, opts);
-
-	bpf_prog_test_run_opts(bpf_program__fd(skel->progs.repartition), &opts);
-}
-
 static void invoke_flush_alloc(struct scx_qmap *skel)
 {
 	LIBBPF_OPTS(bpf_test_run_opts, opts);
@@ -281,7 +199,7 @@ static void print_hier(struct qmap_arena *qa, struct hier_prev *prev, u64 own_cg
 
 	format_cid_ranges(qa, CID_SELF, ranges, sizeof(ranges));
 	printf("hier   : %-4s %10llu %4u %6.2f %8s  %s\n", "self",
-	       (unsigned long long)own_cgid, qa->self_weight,
+	       (unsigned long long)own_cgid, 100,
 	       secs > 0 ? (qa->self_alloc_ns - prev->self_alloc_ns) / (secs * 1e9) : 0.0,
 	       "-", ranges);
 	prev->self_alloc_ns = qa->self_alloc_ns;
@@ -505,8 +423,6 @@ int main(int argc, char **argv)
 	while (!exit_req && !UEI_EXITED(skel, uei)) {
 		long nr_enqueued = qa->nr_enqueued;
 		long nr_dispatched = qa->nr_dispatched;
-		u32 self_weight;
-		bool repart;
 
 		printf("---- %s ----\n",
 		       tstamp(tbuf, sizeof(tbuf)));
@@ -531,20 +447,6 @@ int main(int argc, char **argv)
 			       qa->cpuperf_target_avg,
 			       qa->cpuperf_target_max);
 
-		self_weight = own_cgid ? read_cgroup_weight(own_cgid) : 100;
-		if (!self_weight)
-			self_weight = 100;
-
-		repart = feed_weights(qa);
-
-		if (self_weight != qa->self_weight) {
-			qa->self_weight = self_weight;
-			repart = true;
-		}
-
-		if (repart)
-			invoke_repartition(skel);
-
 		invoke_flush_alloc(skel);
 		print_hier(qa, &hprev, own_cgid);
 		fflush(stdout);
diff --git a/tools/sched_ext/scx_qmap.h b/tools/sched_ext/scx_qmap.h
index 87642d21fce3..6db2ea4bdc85 100644
--- a/tools/sched_ext/scx_qmap.h
+++ b/tools/sched_ext/scx_qmap.h
@@ -90,7 +90,7 @@ struct task_ctx;
 /* per-direct-child state for the sub-scheduler */
 struct sub_sched_ctx {
 	u64 cgroup_id;
-	u32 weight;			/* cpu.weight, fed by userspace */
+	u32 weight;			/* cpu.weight, seeded at attach, then set_weight */
 	u64 nr_dsps;
 	struct qmap_cmask granted_cids;	/* cids granted excl to this child */
 	struct qmap_cmask prev_granted;	/* last grant, for delta calculation */
@@ -147,7 +147,6 @@ struct qmap_arena {
 
 	struct sub_sched_ctx sub_sched_ctxs[MAX_SUB_SCHEDS]; /* per-child context */
 	u64 nr_sub_scheds;		/* number of attached children */
-	u32 self_weight;		/* this node's cpu.weight, fed by userspace */
 
 	/* bpf-internal per-cid state */
 	u8 cid_shared[SCX_QMAP_MAX_CPUS]; /* per cid: 1 if held shared (ENQ_IMMED-only) */
-- 
2.55.0

[PATCH 8/8] tools/sched_ext: scx_qmap - Add init fault injection modes

[Tejun Heo]

Add -J init-fail which makes ops.init_task() fail with -ENOMEM for tasks
whose comm starts with "qmfail", and -J cgrp-init-fail which does the same
in ops.cgroup_init() for cgroups named "qmfail*".

The former exercises the migration veto path: the cgroup.procs write must
fail with the injected errno while the destination sched stays up and the
task stays put. The latter exercises the ownership-return failure path: a
parent failing to re-init a returned cgroup leaves it unowned, and moves and
set_* ops against it must be skipped instead of dereferencing the missing
owner. Matching on "qmfail" names keeps the injecting scheduler's own enable
unaffected.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 tools/sched_ext/scx_qmap.bpf.c | 15 +++++++++++++++
 tools/sched_ext/scx_qmap.c     |  8 +++++++-
 tools/sched_ext/scx_qmap.h     |  2 ++
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 8b8b14309466..925ae1a1d440 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -888,6 +888,10 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init_task, struct task_struct *p,
 	struct task_ctx_stor_val *v;
 	task_ctx_t *taskc;
 
+	if (qa.inject_mode == QMAP_INJ_INIT_FAIL &&
+	    !bpf_strncmp(p->comm, 6, "qmfail"))
+		return -ENOMEM;
+
 	if (p->tgid == disallow_tgid)
 		p->scx.disallow = true;
 
@@ -1012,10 +1016,21 @@ void BPF_STRUCT_OPS(qmap_dump_task, struct scx_dump_ctx *dctx, struct task_struc
 
 s32 BPF_STRUCT_OPS(qmap_cgroup_init, struct cgroup *cgrp, struct scx_cgroup_init_args *args)
 {
+	QMAP_TOUCH_ARENA();
+
 	if (print_msgs)
 		bpf_printk("CGRP INIT %llu weight=%u period=%lu quota=%ld burst=%lu",
 			   cgrp->kn->id, args->weight, args->bw_period_us,
 			   args->bw_quota_us, args->bw_burst_us);
+
+	if (qa.inject_mode == QMAP_INJ_CGRP_INIT_FAIL) {
+		char name[7] = {};
+
+		bpf_probe_read_kernel_str(name, sizeof(name), cgrp->kn->name);
+		if (!bpf_strncmp(name, 6, "qmfail"))
+			return -ENOMEM;
+	}
+
 	return 0;
 }
 
diff --git a/tools/sched_ext/scx_qmap.c b/tools/sched_ext/scx_qmap.c
index 105745f51b8e..46892b4bb448 100644
--- a/tools/sched_ext/scx_qmap.c
+++ b/tools/sched_ext/scx_qmap.c
@@ -69,7 +69,9 @@ const char help_fmt[] =
 "  -C MODE       cid-override test (shuffle|bad-dup|bad-range|bad-mono)\n"
 "  -i SEC        Stats interval, seconds (default 5)\n"
 "  -R MS         Round-robin period for time-shared cpus, ms (default 200)\n"
-"  -J MODE       Fault injection (wrong-cid: dispatch to a cid not held)\n"
+"  -J MODE       Fault injection (wrong-cid: dispatch to a cid not held,\n"
+"                init-fail/cgrp-init-fail: fail init_task/cgroup_init for\n"
+"                \"qmfail*\" comms/cgroups)\n"
 "  -v            Print libbpf debug messages\n"
 "  -h            Display this help and exit\n";
 
@@ -391,6 +393,10 @@ int main(int argc, char **argv)
 		case 'J':
 			if (!strcmp(optarg, "wrong-cid"))
 				inject_mode = QMAP_INJ_WRONG_CID;
+			else if (!strcmp(optarg, "init-fail"))
+				inject_mode = QMAP_INJ_INIT_FAIL;
+			else if (!strcmp(optarg, "cgrp-init-fail"))
+				inject_mode = QMAP_INJ_CGRP_INIT_FAIL;
 			else
 				inject_mode = strtoul(optarg, NULL, 0);
 			break;
diff --git a/tools/sched_ext/scx_qmap.h b/tools/sched_ext/scx_qmap.h
index 6db2ea4bdc85..cc2840e7aa3c 100644
--- a/tools/sched_ext/scx_qmap.h
+++ b/tools/sched_ext/scx_qmap.h
@@ -64,6 +64,8 @@ struct qmap_fifo {
 enum qmap_inject {
 	QMAP_INJ_OFF		= 0,
 	QMAP_INJ_WRONG_CID	= 1,	/* dispatch to a cid we don't hold */
+	QMAP_INJ_INIT_FAIL	= 2,	/* fail init_task for "qmfail*" comms */
+	QMAP_INJ_CGRP_INIT_FAIL	= 3,	/* fail cgroup_init for "qmfail*" cgroups */
 };
 
 /*
-- 
2.55.0