sched_ext: Pass the initial cmask to cid-form ops.enable()

The cid-form API has an obvious hole. A task's cid mask is only visible
through ops.set_cmask(), which fires on affinity changes and class switches
but not when a task enters a scheduler through fork, sub-sched enable or
re-home, and there is no p->cpus_ptr equivalent to fall back on. Schedulers
work around it by seeding the mask in ops.init_task() from p->cpus_ptr cid
by cid, which is subtly wrong: on sub-sched enable and re-home, an affinity
change between init_task() and enable() is delivered to the sched the task
is still on, and nothing corrects the new sched's copy afterwards.

Fix it by adding struct scx_enable_args to cid-form ops.enable() carrying
the task's cmask, built in the per-cpu scratch under the rq lock as the task
enters the scheduler, and calling set_cmask() with the same mask right after
enable(), ahead of set_weight(). A scheduler can then track affinity in
set_cmask() alone, and scx_qmap drops its init_task() seed. set_cmask() no
longer fires for a cid-form task before it is enabled, and the class-switch
republish in switching_to_scx() is limited to the cpu form.

This changes the cid-form ops.enable() signature, which is fine as the
cid-form API is still considered unreleased. An args struct rather than a
bare cmask argument leaves room for more initial state without another
signature change, and the cmask travels as a plain arena address because BTF
can't type arena struct members yet.

v2: The cmask travels as a u64 arena address, cmask_arena_addr, instead of a
kernel-typed pointer, with the typing limitation and the planned typed alias
documented (Sashiko review).

v3: The initial set_cmask() is delivered before set_weight() so that the
mask is in place when weight-dependent state is derived (Andrea Righi).
Selftest added.

Signed-off-by: Tejun Heo <tj@kernel.org>
Reviewed-by: Andrea Righi <arighi@nvidia.com>

1 files changed, 0 insertions, 3 deletions

diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 67b7c01ca..2f3653199 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -961,9 +961,6 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(qmap_init_task, struct task_struct *p,
 	taskc->highpri = false;
 	taskc->core_sched_seq = 0;
 	cmask_init(&taskc->cpus_allowed, 0, scx_bpf_nr_cids());
-	bpf_rcu_read_lock();
-	cmask_from_cpumask(&taskc->cpus_allowed, p->cpus_ptr);
-	bpf_rcu_read_unlock();
 
 	v = bpf_task_storage_get(&task_ctx_stor, p, NULL,
 				 BPF_LOCAL_STORAGE_GET_F_CREATE);