Re: [PATCH 2/3] sched_ext: Add scx_bpf_dsq_reenq_task() to re-enqueue single tasks from a DSQ iteration

[Andrea Righi <arighi@nvidia.com>]

On Sat, Oct 17, 2026 at 12:37:37PM -1000, Tejun Heo wrote:
> scx_bpf_dsq_reenq() re-enqueues a whole DSQ, but a scheduler that loses the
> cpus serving a DSQ, such as a sub-scheduler after a cap revoke, needs only
> the stranded tasks placed again.
> 
> Add scx_bpf_dsq_reenq_task(), which marks a task from a DSQ iteration. When
> the iteration ends, the DSQ's deferred re-enqueue walk is scheduled once and
> re-enqueues only the marked tasks. The kfunc locks the task's rq and the DSQ
> like scx_bpf_dsq_move() and is callable from the same contexts. The mark is
> cleared when the task leaves the DSQ, so a task consumed before the walk is
> not re-enqueued.
> 
> An iteration's DSQ can be destroyed before the iteration ends, so
> schedule_dsq_reenq() now skips a destroyed DSQ instead of aborting the
> scheduler. scx_bpf_dsq_reenq() racing scx_bpf_destroy_dsq() takes the same
> path.
> 
> Declare the kfunc for the tools and add a compat probe.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>
> ---
>  include/linux/sched/ext.h                |  1 +
>  kernel/sched/ext/ext.c                   | 79 +++++++++++++++++++++++-
>  kernel/sched/ext/internal.h              |  2 +
>  tools/sched_ext/include/scx/common.bpf.h |  1 +
>  tools/sched_ext/include/scx/compat.bpf.h |  6 ++
>  5 files changed, 86 insertions(+), 3 deletions(-)
> 
> diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
> index ae696d11dac9..efd635a99309 100644
> --- a/include/linux/sched/ext.h
> +++ b/include/linux/sched/ext.h
> @@ -157,6 +157,7 @@ enum scx_ent_flags {
>  /* scx_entity.dsq_flags */
>  enum scx_ent_dsq_flags {
>  	SCX_TASK_DSQ_ON_PRIQ	= 1 << 0, /* task is queued on the priority queue of a dsq */
> +	SCX_TASK_DSQ_REENQ	= 1 << 1, /* marked by scx_bpf_dsq_reenq_task() */
>  };
>  
>  enum scx_dsq_lnode_flags {
> diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
> index 3e707b00e78b..2d274b54904e 100644
> --- a/kernel/sched/ext/ext.c
> +++ b/kernel/sched/ext/ext.c
> @@ -555,11 +555,13 @@ enum scx_dsq_iter_flags {
>  	/* iterate in the reverse dispatch order */
>  	SCX_DSQ_ITER_REV		= 1U << 16,
>  
> +	__SCX_DSQ_ITER_HAS_REENQ	= 1U << 29,
>  	__SCX_DSQ_ITER_HAS_SLICE	= 1U << 30,
>  	__SCX_DSQ_ITER_HAS_VTIME	= 1U << 31,
>  
>  	__SCX_DSQ_ITER_USER_FLAGS	= SCX_DSQ_ITER_REV,
>  	__SCX_DSQ_ITER_ALL_FLAGS	= __SCX_DSQ_ITER_USER_FLAGS |
> +					  __SCX_DSQ_ITER_HAS_REENQ |
>  					  __SCX_DSQ_ITER_HAS_SLICE |
>  					  __SCX_DSQ_ITER_HAS_VTIME,
>  };
> @@ -1292,7 +1294,9 @@ void schedule_dsq_reenq(struct scx_sched *sch, struct scx_dispatch_q *dsq,
>  			WRITE_ONCE(dru->flags, dru->flags | reenq_flags);
>  		}
>  	} else {
> -		scx_error(sch, "DSQ 0x%llx not allowed for reenq", dsq->id);
> +		/* destroyed since the caller resolved it, nothing to reenq */
> +		if (READ_ONCE(dsq->id) != SCX_DSQ_INVALID)
> +			scx_error(sch, "DSQ 0x%llx not allowed for reenq", dsq->id);
>  		return;
>  	}
>  
> @@ -1918,6 +1922,7 @@ void scx_task_unlink_from_dsq(struct task_struct *p, struct scx_dispatch_q *dsq)
>  		RB_CLEAR_NODE(&p->scx.dsq_priq);
>  		p->scx.dsq_flags &= ~SCX_TASK_DSQ_ON_PRIQ;
>  	}
> +	p->scx.dsq_flags &= ~SCX_TASK_DSQ_REENQ;
>  
>  	list_del_init(&p->scx.dsq_list.node);
>  	dsq_dec_nr(dsq, p);
> @@ -5095,7 +5100,9 @@ static void process_deferred_reenq_locals(struct rq *rq)
>  static bool user_task_should_reenq(struct task_struct *p, u64 reenq_flags, u32 *reason)
>  {
>  	*reason = SCX_TASK_REENQ_KFUNC;
> -	return reenq_flags & SCX_REENQ_ANY;
> +	if (reenq_flags & SCX_REENQ_ANY)
> +		return true;
> +	return (reenq_flags & SCX_REENQ_MARKED) && (p->scx.dsq_flags & SCX_TASK_DSQ_REENQ);

With SCX_REENQ_MARKED this can now return false, and reenq_user() skips the task
while still holding dsq->lock (and the rq lock, with IRQs
disabled). If we have too many tasks queued, it means we're walking all of them
with IRQs off.

This is not a blocker, but as a follow-up maybe we could count the visited tasks
and drop dsq->lock periodically? (the cursor should keep the position)

>  }
>  
>  static void reenq_user(struct rq *rq, struct scx_dispatch_q *dsq, u64 reenq_flags)
> @@ -10087,6 +10094,8 @@ BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_slice, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_vtime, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_move, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_move_vtime, KF_RCU)
> +/* also in scx_kfunc_ids_unlocked, same contexts as scx_bpf_dsq_move() */
> +BTF_ID_FLAGS(func, scx_bpf_dsq_reenq_task, KF_RCU)
>  #ifdef CONFIG_EXT_SUB_SCHED
>  BTF_ID_FLAGS(func, scx_bpf_sub_dispatch, KF_IMPLICIT_ARGS)
>  #endif
> @@ -10202,6 +10211,7 @@ BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_slice, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_move_set_vtime, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_move, KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_dsq_move_vtime, KF_RCU)
> +BTF_ID_FLAGS(func, scx_bpf_dsq_reenq_task, KF_RCU)
>  /* also in scx_kfunc_ids_select_cpu: also callable from ops.select_cpu()/ops.enqueue() */
>  BTF_ID_FLAGS(func, __scx_bpf_select_cpu_and, KF_IMPLICIT_ARGS | KF_RCU)
>  BTF_ID_FLAGS(func, scx_bpf_select_cpu_and, KF_RCU)
> @@ -10610,7 +10620,8 @@ __bpf_kfunc struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *i
>   * bpf_iter_scx_dsq_destroy - Destroy a DSQ iterator
>   * @it: iterator to destroy
>   *
> - * Undo bpf_iter_scx_dsq_new().
> + * Undo bpf_iter_scx_dsq_new(). If the iteration marked tasks with
> + * scx_bpf_dsq_reenq_task(), schedule their re-enqueue.
>   */
>  __bpf_kfunc void bpf_iter_scx_dsq_destroy(struct bpf_iter_scx_dsq *it)
>  {
> @@ -10626,6 +10637,13 @@ __bpf_kfunc void bpf_iter_scx_dsq_destroy(struct bpf_iter_scx_dsq *it)
>  		list_del_init(&kit->cursor.node);
>  		raw_spin_unlock_irqrestore(&kit->dsq->lock, flags);
>  	}
> +
> +	if (kit->cursor.flags & __SCX_DSQ_ITER_HAS_REENQ) {
> +		guard(preempt)();
> +		schedule_dsq_reenq(kit->dsq->sched, kit->dsq, SCX_REENQ_MARKED,
> +				   scx_locked_rq());
> +	}
> +
>  	kit->dsq = NULL;
>  }
>  
> @@ -10714,6 +10732,61 @@ __bpf_kfunc void scx_bpf_dsq_reenq(u64 dsq_id, u64 reenq_flags,
>  	schedule_dsq_reenq(sch, dsq, reenq_flags, locked_rq);
>  }
>  
> +/**
> + * scx_bpf_dsq_reenq_task - Re-enqueue a task from a DSQ iteration
> + * @it__iter: DSQ iterator in progress
> + * @p: task to re-enqueue
> + *
> + * Mark @p, which is on the DSQ iterated by @it__iter and was queued before the
> + * iteration started, for re-enqueue. When the iteration ends, the DSQ's marked
> + * tasks go through ops.enqueue() with %SCX_ENQ_REENQ in DSQ order and the other
> + * tasks on the DSQ keep their positions. A marked task that leaves the DSQ in
> + * the meantime is not re-enqueued. @p must be on the DSQ's scheduler.

nit: the mark belongs to the task, if two CPUs iterate the same DSQ
concurrently, the walk scheduled by one of them can re-enqueue tasks marked by
the other before its iteration ends. It's harmless, but maybe worth a sentence
here.

> + *
> + * Re-enqueues are performed asynchronously. Can be called from ops.dispatch()
> + * or any BPF context which doesn't hold a rq lock (e.g. BPF timers or SYSCALL
> + * programs).
> + *
> + * Return %true if @p was marked, %false if @p left the DSQ or was queued after
> + * the iteration started, or the iterator holds no DSQ.

nit: also %false if @p belongs to a different scheduler (after scx_error()).

None of the above are blockers, so:

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

Thanks,
-Andrea