[PATCH] sched_ext: Allow ops.cgroup_set_weight/idle() to be sleepable

[Tao Cui]

From: Tao Cui <cuitao@xxxxxxxxxx>

ops.cgroup_set_weight() and ops.cgroup_set_idle() are delivered from
scx_group_set_weight() and scx_group_set_idle(), which run from the
cpu.weight (and v1 cpu.shares) and cpu.idle cgroup interface write paths
in process context. Both hold percpu_down_read(&scx_cgroup_ops_rwsem),
whose read side may sleep. The call sites are therefore sleepable, like
ops.cgroup_set_bandwidth(), which was recently added to the sleepable
allow-list.

bpf_scx_check_member() rejects a sleepable program on any member not on
its allow-list, so a BPF scheduler cannot allocate -- which is sleepable
-- when a cgroup's weight or idle state changes at runtime. Add
cgroup_set_weight() and cgroup_set_idle() to the allow-list so these
callbacks can allocate on demand, and document that they may block.

Also add the matching compatibility markers, so userspace can detect
this support via BTF, mirroring
scx_compat_marker_cgroup_set_bandwidth_may_sleep().

To size the alternative, a scheduler that gives each cgroup a
dedicated idle DSQ must create it in ops.cgroup_init() for every
cgroup up front. In a VM with 2000 cgroups that is 2000+ standing
DSQs, each a struct scx_dispatch_q plus a per-CPU area. With the
allow-list entries the same scheduler can create the DSQ lazily on
the first cpu.idle=1 write of a cgroup: 4 allocations for the 4
cgroups marked idle at runtime, and repeated cpu.idle writes do not re-allocate. Verified with a probe scheduler on an
unpatched kernel (load rejected with -EINVAL) and on a patched one.

Signed-off-by: Tao Cui <cuitao@xxxxxxxxxx>
---
 kernel/sched/ext/ext.c      | 18 ++++++++++++++++++
 kernel/sched/ext/internal.h |  9 +++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index b646711a45fe..53b888f5a6c5 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -8079,7 +8079,9 @@ static int bpf_scx_check_member(const struct btf_type *t,
 	case offsetof(struct sched_ext_ops, cgroup_init):
 	case offsetof(struct sched_ext_ops, cgroup_exit):
 	case offsetof(struct sched_ext_ops, cgroup_prep_move):
+	case offsetof(struct sched_ext_ops, cgroup_set_weight):
 	case offsetof(struct sched_ext_ops, cgroup_set_bandwidth):
+	case offsetof(struct sched_ext_ops, cgroup_set_idle):
 #endif
 	case offsetof(struct sched_ext_ops, cpu_online):
 	case offsetof(struct sched_ext_ops, cpu_offline):
@@ -11055,3 +11057,19 @@ __initcall(scx_init);
 #ifdef CONFIG_EXT_GROUP_SCHED
 DEFINE_SCX_COMPAT_MARKER(cgroup_set_bandwidth_may_sleep);
 #endif	/* CONFIG_EXT_GROUP_SCHED */
+
+/*
+ * scx_compat_marker_cgroup_set_weight_may_sleep: advertises that
+ * ops.cgroup_set_weight() may be implemented as a sleepable callback.
+ */
+#ifdef CONFIG_EXT_GROUP_SCHED
+DEFINE_SCX_COMPAT_MARKER(cgroup_set_weight_may_sleep);
+#endif	/* CONFIG_EXT_GROUP_SCHED */
+
+/*
+ * scx_compat_marker_cgroup_set_idle_may_sleep: advertises that
+ * ops.cgroup_set_idle() may be implemented as a sleepable callback.
+ */
+#ifdef CONFIG_EXT_GROUP_SCHED
+DEFINE_SCX_COMPAT_MARKER(cgroup_set_idle_may_sleep);
+#endif	/* CONFIG_EXT_GROUP_SCHED */
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 53e136a47924..f81d03de2d3c 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -736,7 +736,7 @@ struct sched_ext_ops {
 	 * @cgrp: cgroup whose weight is being updated
 	 * @weight: new weight [1..10000]
 	 *
-	 * Update @cgrp's weight to @weight.
+	 * Update @cgrp's weight to @weight. This operation may block.
 	 *
 	 * 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
@@ -773,9 +773,10 @@ struct sched_ext_ops {
 	 * @cgrp: cgroup whose idle state is being updated
 	 * @idle: whether the cgroup is entering or exiting idle state
 	 *
-	 * 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.
+	 * Update @cgrp's idle state to @idle. This operation may block. 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().
 	 */
-- 
2.43.0