Re: [PATCH 2/2] sched_ext: Maintain an online cid mask in the scheduler arena

[Andrea Righi <arighi@nvidia.com>]

Hi Tejun,

On Mon, Sep 14, 2026 at 01:42:59PM -1000, Tejun Heo wrote:
> Schedulers on the default cid mapping treat [0, nr_online_cids) as the
> online set and restart on hotplug. Schedulers that install their own mapping
> with scx_bpf_cid_override() have no way to learn which cids are online: the
> count no longer identifies members and the CPU-form cpumask is unusable from
> cid programs. This is an obvious hole in the cid API.
> 
> Add scx_bpf_online_cmask(), a kernel-maintained cmask in the scheduler's
> arena, allocated alongside the per-CPU scratch masks and populated after the
> cid mapping is finalized and before ops.init(), for child schedulers too.
> The pointer stays valid through ops.exit() with no reference to take. It is
> the arena offset as a void pointer, the same form struct_ops arena arguments
> arrive in. The verifier types the void return as a scalar for the program's
> arena cast.
> 
> The mask follows the SCX hotplug notifications: seeded from cpu_active_mask
> and updated before ops.cid_online/offline() runs, so it lags cpu_online_mask
> only inside a hotplug transition. Updates walk the scheduler list under the
> lock that also serializes unlinking. Reads are live, not atomic snapshots.
> Root initialization excludes hotplug.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>
> ---
...
> +/**
> + * scx_bpf_online_cmask - Return the online cid mask in the scheduler arena
> + * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
> + *
> + * Return a read-only cmask covering [0, scx_bpf_nr_cids()) which the kernel
> + * keeps in the scheduler's arena, or NULL if the calling program is not
> + * associated with a live cid-form scheduler. The mask follows the SCX hotplug

Documentation nit (feel free to ignore): since the read-only property isn't
enforced by the verifier, should we rephrase this part as the following (or
something along these lines):

  Return a kernel-maintained cmask covering [0, scx_bpf_nr_cids()), or NULL
  if the calling program is not associated with a live cid-form scheduler.
  The returned cmask should be treated as read-only, even though arena memory
  remains writable by the BPF scheduler.

And still keep const return type as source-level documentation/checking.

Other that that, looks good to me.

Reviewed-by: Andrea Righi <arighi@nvidia.com>

Thanks,
-Andrea

> + * notifications: a cid's bit is updated before ops.cid_online/offline() runs
> + * for it. The pointer is valid from ops.init() through ops.exit(). Root
> + * ops.init() runs with hotplug excluded. Other contexts can observe concurrent
> + * updates.
> + */
> +__bpf_kfunc const void *scx_bpf_online_cmask(const struct bpf_prog_aux *aux)
> +{
> +	struct scx_sched *sch;
> +	struct scx_cmask *online;
> +
> +	guard(rcu)();
> +
> +	sch = scx_prog_sched(aux);
> +	if (unlikely(!sch))
> +		return NULL;
> +	online = sch->online_cmask;
> +	if (unlikely(!online))
> +		return NULL;
> +
> +	/* BPF rebases by the low 32 bits, like __arena callback args */
> +	return (void *)((unsigned long)online - sch->arena_kern_base);
> +}
> +
>  /**
>   * scx_bpf_this_cid - Return the cid of the CPU this program is running on
>   *
> @@ -10709,6 +10770,7 @@ BTF_ID_FLAGS(func, scx_bpf_nr_node_ids)
>  BTF_ID_FLAGS(func, scx_bpf_nr_cpu_ids)
>  BTF_ID_FLAGS(func, scx_bpf_nr_cids)
>  BTF_ID_FLAGS(func, scx_bpf_nr_online_cids)
> +BTF_ID_FLAGS(func, scx_bpf_online_cmask, KF_IMPLICIT_ARGS | KF_ARENA_RET)
>  BTF_ID_FLAGS(func, scx_bpf_this_cid)
>  BTF_ID_FLAGS(func, scx_bpf_get_possible_cpumask, KF_ACQUIRE)
>  BTF_ID_FLAGS(func, scx_bpf_get_online_cpumask, KF_ACQUIRE)
> diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
> index 2586879e9084..067cddf47220 100644
> --- a/kernel/sched/ext/internal.h
> +++ b/kernel/sched/ext/internal.h
> @@ -1572,6 +1572,7 @@ struct scx_sched {
>  	 * and passes it to the callback's __arena argument.
>  	 */
>  	struct scx_cmask * __percpu *set_cmask_scratch;
> +	struct scx_cmask *online_cmask;
>  
>  	DECLARE_BITMAP(has_op, SCX_OPI_END);
>  
> @@ -2098,7 +2099,7 @@ void scx_disable_and_exit_task(struct scx_sched *sch, struct task_struct *p);
>  void scx_cgroup_lock(void);
>  void scx_cgroup_unlock(void);
>  #endif
> -s32 scx_set_cmask_scratch_alloc(struct scx_sched *sch);
> +s32 scx_alloc_kern_arena_objs(struct scx_sched *sch);
>  void scx_disable_bypass_dsp(struct scx_sched *sch);
>  void scx_bypass(struct scx_sched *sch, bool bypass);
>  s32 scx_link_sched(struct scx_sched *sch);
> diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
> index 385302d19914..f7aeb1488566 100644
> --- a/kernel/sched/ext/sub.c
> +++ b/kernel/sched/ext/sub.c
> @@ -1805,6 +1805,12 @@ void scx_sub_enable_workfn(struct kthread_work *work)
>  		goto err_disable;
>  	}
>  
> +	scoped_guard(cpus_read_lock) {
> +		ret = scx_alloc_kern_arena_objs(sch);
> +		if (ret)
> +			goto err_disable;
> +	}
> +
>  	if (sch->ops.init) {
>  		ret = SCX_CALL_OP_RET(sch, init, NULL);
>  		if (ret) {
> @@ -1815,10 +1821,6 @@ void scx_sub_enable_workfn(struct kthread_work *work)
>  		sch->exit_info->flags |= SCX_EFLAG_INITIALIZED;
>  	}
>  
> -	ret = scx_set_cmask_scratch_alloc(sch);
> -	if (ret)
> -		goto err_disable;
> -
>  	struct scx_sub_attach_args sub_attach_args = {
>  		.ops = &sch->ops,
>  		.cgroup_path = sch->cgrp_path,
> diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
> index 76f5e025e107..2ddb01a059fd 100644
> --- a/tools/sched_ext/include/scx/common.bpf.h
> +++ b/tools/sched_ext/include/scx/common.bpf.h
> @@ -113,6 +113,7 @@ s32 scx_bpf_this_cid(void) __ksym __weak;
>  struct task_struct *scx_bpf_cid_curr(s32 cid) __ksym __weak;
>  u32 scx_bpf_nr_cids(void) __ksym __weak;
>  u32 scx_bpf_nr_online_cids(void) __ksym __weak;
> +const void __arena *scx_bpf_online_cmask(void) __ksym __weak;
>  u32 scx_bpf_cidperf_cap(s32 cid) __ksym __weak;
>  u32 scx_bpf_cidperf_cur(s32 cid) __ksym __weak;
>  s32 scx_bpf_cidperf_set(s32 cid, u32 perf) __ksym __weak;
> -- 
> 2.55.0
>