[PATCHSET sched_ext/for-7.4] sched_ext: Add scx_bpf_dsq_reenq_task()

[PATCHSET sched_ext/for-7.4] sched_ext: Add scx_bpf_dsq_reenq_task()

[Tejun Heo]

Hello,

A scheduler that loses the cpus serving a DSQ needs the tasks stranded on it
placed again, and only those. scx_bpf_dsq_reenq() re-enqueues the whole DSQ:
every task goes back through ops.enqueue() and is re-placed and reordered.
The only per-task primitive, scx_bpf_dsq_move(), is a placement decision,
not a re-enqueue. A sub-scheduler sweeping its DSQs after a cap revoke is
one such caller.

Patch 1 factors the rq lock switch out of scx_dsq_move() so that the new
kfunc can share it. Patch 2 adds 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, in
DSQ order with SCX_ENQ_REENQ. The other tasks keep their positions.

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.

Patch 3 makes scx_qmap re-enqueue some of the tasks stranded on its shared
DSQs with the new kfunc instead of moving them to the rescue path, so that
both kfuncs are exercised. They go back through ops.enqueue(), which rescues
them, and the other queued tasks keep their places.

Based on sched_ext/for-7.4 (9c0af54cb49c).

This patchset contains the following 3 patches.

 0001 sched_ext: Factor out the rq lock switch of scx_dsq_move()
 0002 sched_ext: Add scx_bpf_dsq_reenq_task() to re-enqueue single tasks from a DSQ iteration
 0003 sched_ext: scx_qmap: Also re-enqueue stranded tasks with scx_bpf_dsq_reenq_task()

The patchset is also available in the following git branch:

 git://git.kernel.org/pub/scm/linux/kernel/git/tj/sched_ext.git dsq-reenq-task

diffstat follows. Thanks.

 include/linux/sched/ext.h                |   1 +
 kernel/sched/ext/ext.c                   | 142 ++++++++++++++++++++++++-------
 kernel/sched/ext/internal.h              |   2 +
 tools/sched_ext/include/scx/common.bpf.h |   1 +
 tools/sched_ext/include/scx/compat.bpf.h |   6 ++
 tools/sched_ext/scx_qmap.bpf.c           |  50 ++++++-----
 6 files changed, 153 insertions(+), 49 deletions(-)

--
tejun

[PATCH 1/3] sched_ext: Factor out the rq lock switch of scx_dsq_move()

[Tejun Heo]

Factor the rq lock switch on entry to scx_dsq_move() and its undo on exit
out into kf_switch_rq_lock() and kf_restore_rq_lock() so that other kfuncs
operating on a task from a DSQ iteration can share them. No functional
change.

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

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index b9f80fb29255..3e707b00e78b 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -9745,6 +9745,40 @@ static const struct btf_kfunc_id_set scx_kfunc_set_enqueue_dispatch = {
 	.filter			= scx_kfunc_context_filter,
 };
 
+/*
+ * Lock @rq from a kfunc called either from ops.dispatch(), which holds the
+ * dispatched rq's lock, or from a context holding no rq lock. Under core
+ * scheduling, dispatch can run for a sibling rq, so the locked rq is not
+ * necessarily this CPU's. Return the rq locked on entry, NULL if none.
+ */
+static struct rq *kf_switch_rq_lock(struct rq *rq, unsigned long *flags)
+{
+	struct rq *entry_rq;
+
+	local_irq_save(*flags);
+
+	entry_rq = scx_locked_rq();
+	if (entry_rq) {
+		if (entry_rq != rq)
+			switch_rq_lock(entry_rq, rq);
+	} else {
+		raw_spin_rq_lock(rq);
+	}
+	return entry_rq;
+}
+
+/* undo kf_switch_rq_lock(), @rq is the rq locked now */
+static void kf_restore_rq_lock(struct rq *rq, struct rq *entry_rq, unsigned long flags)
+{
+	if (entry_rq) {
+		if (entry_rq != rq)
+			switch_rq_lock(rq, entry_rq);
+	} else {
+		scx_rq_lock_drop(rq);
+		raw_spin_rq_unlock_irqrestore(rq, flags);
+	}
+}
+
 static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 			 struct task_struct *p, u64 dsq_id, u64 enq_flags,
 			 bool priq)
@@ -9779,27 +9813,8 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 	if (unlikely(READ_ONCE(sch->aborting)))
 		return false;
 
-	/*
-	 * Can be called from either ops.dispatch() holding the dispatched rq's
-	 * lock or any context where no rq lock is held. If latter, lock @p's
-	 * task_rq which we'll likely need anyway.
-	 */
 	src_rq = task_rq(p);
-
-	local_irq_save(flags);
-
-	/*
-	 * Under core scheduling, dispatch can run for a sibling rq, so the
-	 * locked rq is not necessarily this CPU's.
-	 */
-	locked_rq = scx_locked_rq();
-
-	if (locked_rq) {
-		if (locked_rq != src_rq)
-			switch_rq_lock(locked_rq, src_rq);
-	} else {
-		raw_spin_rq_lock(src_rq);
-	}
+	locked_rq = kf_switch_rq_lock(src_rq, &flags);
 
 	p_rq = src_rq;
 	raw_spin_lock(&src_dsq->lock);
@@ -9838,13 +9853,7 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
 	p_rq = move_task_between_dsqs(sch, p, enq_flags, src_dsq, dst_dsq);
 	dispatched = true;
 out:
-	if (locked_rq) {
-		if (locked_rq != p_rq)
-			switch_rq_lock(p_rq, locked_rq);
-	} else {
-		scx_rq_lock_drop(p_rq);
-		raw_spin_rq_unlock_irqrestore(p_rq, flags);
-	}
+	kf_restore_rq_lock(p_rq, locked_rq, flags);
 
 	kit->cursor.flags &= ~(__SCX_DSQ_ITER_HAS_SLICE |
 			       __SCX_DSQ_ITER_HAS_VTIME);
-- 
2.56.0

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

[Tejun Heo]

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);
 }
 
 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.
+ *
+ * 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.
+ */
+__bpf_kfunc bool scx_bpf_dsq_reenq_task(struct bpf_iter_scx_dsq *it__iter,
+					struct task_struct *p)
+{
+	struct bpf_iter_scx_dsq_kern *kit = (void *)it__iter;
+	struct scx_dispatch_q *dsq = kit->dsq;
+	struct rq *src_rq, *entry_rq;
+	unsigned long flags;
+	bool marked = false;
+
+	/* NULL after a failed bpf_iter_scx_dsq_new(), see scx_dsq_move() */
+	if (unlikely(!dsq))
+		return false;
+
+	/* the cursor check needs @p's rq locked */
+	src_rq = task_rq(p);
+	entry_rq = kf_switch_rq_lock(src_rq, &flags);
+
+	scoped_guard (raw_spinlock, &dsq->lock) {
+		if (nldsq_cursor_lost_task(&kit->cursor, src_rq, dsq, p))
+			break;
+
+		if (unlikely(!scx_task_on_sched(dsq->sched, p))) {
+			scx_error(dsq->sched, "scx_bpf_dsq_reenq_task() on %s[%d] but the task belongs to a different scheduler",
+				  p->comm, p->pid);
+			break;
+		}
+
+		p->scx.dsq_flags |= SCX_TASK_DSQ_REENQ;
+		/* the walk is scheduled once, by bpf_iter_scx_dsq_destroy() */
+		kit->cursor.flags |= __SCX_DSQ_ITER_HAS_REENQ;
+		marked = true;
+	}
+
+	kf_restore_rq_lock(src_rq, entry_rq, flags);
+	return marked;
+}
+
 /**
  * scx_bpf_reenqueue_local___v2 - Re-enqueue tasks on a local DSQ
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index b408993af794..29e69a83d876 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1946,6 +1946,8 @@ enum scx_reenq_flags {
 
 	/* internal: kernel-issued on cap revoke, not accepted from BPF */
 	SCX_REENQ_CAP_REVOKE	= 1LLU << 1,
+	/* internal: tasks marked by scx_bpf_dsq_reenq_task() */
+	SCX_REENQ_MARKED	= 1LLU << 2,
 
 	__SCX_REENQ_FILTER_MASK	= 0xffffLLU,
 
diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
index 651bd4952647..9376d5704a47 100644
--- a/tools/sched_ext/include/scx/common.bpf.h
+++ b/tools/sched_ext/include/scx/common.bpf.h
@@ -75,6 +75,7 @@ struct task_struct *scx_bpf_dsq_peek(u64 dsq_id) __ksym __weak;
 int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id, u64 flags) __ksym __weak;
 struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it) __ksym __weak;
 void bpf_iter_scx_dsq_destroy(struct bpf_iter_scx_dsq *it) __ksym __weak;
+bool scx_bpf_dsq_reenq_task(struct bpf_iter_scx_dsq *it, struct task_struct *p) __ksym __weak;
 void scx_bpf_exit_bstr(s64 exit_code, char *fmt, unsigned long long *data, u32 data__sz) __ksym __weak;
 void scx_bpf_error_bstr(char *fmt, unsigned long long *data, u32 data_len) __ksym;
 void scx_bpf_dump_bstr(char *fmt, unsigned long long *data, u32 data_len) __ksym __weak;
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index c26024a6e2b1..254882bd8ce6 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -498,6 +498,12 @@ static inline bool __COMPAT_has_generic_reenq(void)
 	return bpf_ksym_exists(scx_bpf_dsq_reenq___compat);
 }
 
+/* v7.4: Add scx_bpf_dsq_reenq_task(). */
+static inline bool __COMPAT_has_dsq_reenq_task(void)
+{
+	return bpf_ksym_exists(scx_bpf_dsq_reenq_task);
+}
+
 /*
  * v6.19: The new void variant can be called from anywhere while the older v1
  * variant can only be called from ops.cpu_release(). The double ___ prefixes on
-- 
2.56.0

[PATCH 3/3] sched_ext: scx_qmap: Also re-enqueue stranded tasks with scx_bpf_dsq_reenq_task()

[Tejun Heo]

scan_shared_dsq() rescues a task stranded on SHARED_DSQ or HIGHPRI_DSQ by
moving it to the rescue path with scx_bpf_dsq_move(). For a task with an odd
pid, mark it with scx_bpf_dsq_reenq_task() instead, which sends it back
through qmap_enqueue() and its rescue, so that both kfuncs are exercised.

scan_shared_dsq() runs from qmap_dispatch() and from the monitor timer, so
the kfunc can be called both with the dispatch rq locked and with no rq lock
held.

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 tools/sched_ext/scx_qmap.bpf.c | 50 +++++++++++++++++++++-------------
 1 file changed, 31 insertions(+), 19 deletions(-)

diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 61143a2d79f8..e8e451be862d 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -651,10 +651,35 @@ static void update_core_sched_head_seq(struct task_struct *p)
 		qa.core_sched_head_seqs[idx] = taskc->core_sched_seq;
 }
 
+/*
+ * Send @p, stranded on the DSQ @it iterates, to the kernel rescue path. To
+ * exercise both kfuncs, move tasks with an even pid there with
+ * scx_bpf_dsq_move() and re-enqueue the others with scx_bpf_dsq_reenq_task(),
+ * which takes them through qmap_enqueue()'s rescue. The choice is per task so
+ * that a marked task isn't moved by another scan before the re-enqueue reaches
+ * it.
+ */
+static __always_inline void rescue_stranded(struct bpf_iter_scx_dsq *it,
+					    struct task_struct *p, task_ctx_t *taskc)
+{
+	s32 c;
+
+	if (p->pid & 1) {
+		scx_bpf_dsq_reenq_task(it, p);
+		return;
+	}
+
+	c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
+	if (c >= 0 && c < scx_bpf_nr_cids()) {
+		__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
+		scx_bpf_dsq_move(it, p, SCX_DSQ_LOCAL_ON | c, SCX_ENQ_RESCUE);
+	}
+}
+
 /*
  * One pass over SHARED_DSQ: rescue stranded tasks and boost highpri ones. A
  * task whose cids were lost while it was queued in the fifos would strand on
- * SHARED_DSQ, which is consumed only on self cids it can't run on - move it to
+ * SHARED_DSQ, which is consumed only on self cids it can't run on - send it to
  * the kernel rescue path. One whose cids were lost after the highpri cull is
  * likewise rescued out of HIGHPRI_DSQ below.
  *
@@ -663,9 +688,9 @@ static void update_core_sched_head_seq(struct task_struct *p)
  * consuming them first. This makes minor difference only when dsp_batch is
  * larger than 1.
  *
- * scx_bpf_dsq_move[_vtime]() are allowed both from ops.dispatch() and
- * non-rq-lock holding BPF programs. As demonstration, this function is called
- * from qmap_dispatch() and monitor_timerfn().
+ * scx_bpf_dsq_move[_vtime]() and scx_bpf_dsq_reenq_task() are allowed both from
+ * ops.dispatch() and non-rq-lock holding BPF programs. As demonstration, this
+ * function is called from qmap_dispatch() and monitor_timerfn().
  */
 static bool scan_shared_dsq(bool from_timer)
 {
@@ -677,19 +702,13 @@ static bool scan_shared_dsq(bool from_timer)
 	bpf_for_each(scx_dsq, p, SHARED_DSQ, 0) {
 		static u64 highpri_seq;
 		task_ctx_t *taskc;
-		s32 c;
 
 		if (!(taskc = lookup_task_ctx(p)))
 			return false;
 
 		/* stranded? rescue - it can't be dispatched here either way */
 		if (!cmask_intersects(&taskc->cpus_allowed, &qa.self_cids.mask)) {
-			c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
-			if (c >= 0 && c < scx_bpf_nr_cids()) {
-				__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
-				scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | c,
-						 SCX_ENQ_RESCUE);
-			}
+			rescue_stranded(BPF_FOR_EACH_ITER, p, taskc);
 			continue;
 		}
 
@@ -723,19 +742,12 @@ static bool scan_shared_dsq(bool from_timer)
 						      &qa.usable_cids.mask,
 						      this_cid + 1);
 		if (cid >= nr_cids) {
-			s32 c;
-
 			/* self cids lack caps in effect yet, leave it queued */
 			if (cmask_intersects(&taskc->cpus_allowed, &qa.self_cids.mask))
 				continue;
 
 			/* stranded after the cull - rescue it from here */
-			c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
-			if (c >= 0 && c < nr_cids) {
-				__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
-				scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | c,
-						 SCX_ENQ_RESCUE);
-			}
+			rescue_stranded(BPF_FOR_EACH_ITER, p, taskc);
 			continue;
 		}
 
-- 
2.56.0

Re: [PATCH 1/3] sched_ext: Factor out the rq lock switch of scx_dsq_move()

[Andrea Righi]

On Sat, Oct 17, 2026 at 12:37:36PM -1000, Tejun Heo wrote:
> Factor the rq lock switch on entry to scx_dsq_move() and its undo on exit
> out into kf_switch_rq_lock() and kf_restore_rq_lock() so that other kfuncs
> operating on a task from a DSQ iteration can share them. No functional
> change.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>

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

Thanks,
-Andrea

> ---
>  kernel/sched/ext/ext.c | 63 ++++++++++++++++++++++++------------------
>  1 file changed, 36 insertions(+), 27 deletions(-)
> 
> diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
> index b9f80fb29255..3e707b00e78b 100644
> --- a/kernel/sched/ext/ext.c
> +++ b/kernel/sched/ext/ext.c
> @@ -9745,6 +9745,40 @@ static const struct btf_kfunc_id_set scx_kfunc_set_enqueue_dispatch = {
>  	.filter			= scx_kfunc_context_filter,
>  };
>  
> +/*
> + * Lock @rq from a kfunc called either from ops.dispatch(), which holds the
> + * dispatched rq's lock, or from a context holding no rq lock. Under core
> + * scheduling, dispatch can run for a sibling rq, so the locked rq is not
> + * necessarily this CPU's. Return the rq locked on entry, NULL if none.
> + */
> +static struct rq *kf_switch_rq_lock(struct rq *rq, unsigned long *flags)
> +{
> +	struct rq *entry_rq;
> +
> +	local_irq_save(*flags);
> +
> +	entry_rq = scx_locked_rq();
> +	if (entry_rq) {
> +		if (entry_rq != rq)
> +			switch_rq_lock(entry_rq, rq);
> +	} else {
> +		raw_spin_rq_lock(rq);
> +	}
> +	return entry_rq;
> +}
> +
> +/* undo kf_switch_rq_lock(), @rq is the rq locked now */
> +static void kf_restore_rq_lock(struct rq *rq, struct rq *entry_rq, unsigned long flags)
> +{
> +	if (entry_rq) {
> +		if (entry_rq != rq)
> +			switch_rq_lock(rq, entry_rq);
> +	} else {
> +		scx_rq_lock_drop(rq);
> +		raw_spin_rq_unlock_irqrestore(rq, flags);
> +	}
> +}
> +
>  static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
>  			 struct task_struct *p, u64 dsq_id, u64 enq_flags,
>  			 bool priq)
> @@ -9779,27 +9813,8 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
>  	if (unlikely(READ_ONCE(sch->aborting)))
>  		return false;
>  
> -	/*
> -	 * Can be called from either ops.dispatch() holding the dispatched rq's
> -	 * lock or any context where no rq lock is held. If latter, lock @p's
> -	 * task_rq which we'll likely need anyway.
> -	 */
>  	src_rq = task_rq(p);
> -
> -	local_irq_save(flags);
> -
> -	/*
> -	 * Under core scheduling, dispatch can run for a sibling rq, so the
> -	 * locked rq is not necessarily this CPU's.
> -	 */
> -	locked_rq = scx_locked_rq();
> -
> -	if (locked_rq) {
> -		if (locked_rq != src_rq)
> -			switch_rq_lock(locked_rq, src_rq);
> -	} else {
> -		raw_spin_rq_lock(src_rq);
> -	}
> +	locked_rq = kf_switch_rq_lock(src_rq, &flags);
>  
>  	p_rq = src_rq;
>  	raw_spin_lock(&src_dsq->lock);
> @@ -9838,13 +9853,7 @@ static bool scx_dsq_move(struct bpf_iter_scx_dsq_kern *kit,
>  	p_rq = move_task_between_dsqs(sch, p, enq_flags, src_dsq, dst_dsq);
>  	dispatched = true;
>  out:
> -	if (locked_rq) {
> -		if (locked_rq != p_rq)
> -			switch_rq_lock(p_rq, locked_rq);
> -	} else {
> -		scx_rq_lock_drop(p_rq);
> -		raw_spin_rq_unlock_irqrestore(p_rq, flags);
> -	}
> +	kf_restore_rq_lock(p_rq, locked_rq, flags);
>  
>  	kit->cursor.flags &= ~(__SCX_DSQ_ITER_HAS_SLICE |
>  			       __SCX_DSQ_ITER_HAS_VTIME);
> -- 
> 2.56.0
> 

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

[Andrea Righi]

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

Re: [PATCH 3/3] sched_ext: scx_qmap: Also re-enqueue stranded tasks with scx_bpf_dsq_reenq_task()

[Andrea Righi]

On Sat, Oct 17, 2026 at 12:37:38PM -1000, Tejun Heo wrote:
> scan_shared_dsq() rescues a task stranded on SHARED_DSQ or HIGHPRI_DSQ by
> moving it to the rescue path with scx_bpf_dsq_move(). For a task with an odd
> pid, mark it with scx_bpf_dsq_reenq_task() instead, which sends it back
> through qmap_enqueue() and its rescue, so that both kfuncs are exercised.
> 
> scan_shared_dsq() runs from qmap_dispatch() and from the monitor timer, so
> the kfunc can be called both with the dispatch rq locked and with no rq lock
> held.
> 
> Signed-off-by: Tejun Heo <tj@kernel.org>
> ---
>  tools/sched_ext/scx_qmap.bpf.c | 50 +++++++++++++++++++++-------------
>  1 file changed, 31 insertions(+), 19 deletions(-)
> 
> diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
> index 61143a2d79f8..e8e451be862d 100644
> --- a/tools/sched_ext/scx_qmap.bpf.c
> +++ b/tools/sched_ext/scx_qmap.bpf.c
> @@ -651,10 +651,35 @@ static void update_core_sched_head_seq(struct task_struct *p)
>  		qa.core_sched_head_seqs[idx] = taskc->core_sched_seq;
>  }
>  
> +/*
> + * Send @p, stranded on the DSQ @it iterates, to the kernel rescue path. To
> + * exercise both kfuncs, move tasks with an even pid there with
> + * scx_bpf_dsq_move() and re-enqueue the others with scx_bpf_dsq_reenq_task(),
> + * which takes them through qmap_enqueue()'s rescue. The choice is per task so
> + * that a marked task isn't moved by another scan before the re-enqueue reaches
> + * it.
> + */
> +static __always_inline void rescue_stranded(struct bpf_iter_scx_dsq *it,
> +					    struct task_struct *p, task_ctx_t *taskc)
> +{
> +	s32 c;
> +
> +	if (p->pid & 1) {

nit: if ((p->pid & 1) && __COMPAT_has_dsq_reenq_task())

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

Thanks,
-Andrea

> +		scx_bpf_dsq_reenq_task(it, p);
> +		return;
> +	}
> +
> +	c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
> +	if (c >= 0 && c < scx_bpf_nr_cids()) {
> +		__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
> +		scx_bpf_dsq_move(it, p, SCX_DSQ_LOCAL_ON | c, SCX_ENQ_RESCUE);
> +	}
> +}
> +
>  /*
>   * One pass over SHARED_DSQ: rescue stranded tasks and boost highpri ones. A
>   * task whose cids were lost while it was queued in the fifos would strand on
> - * SHARED_DSQ, which is consumed only on self cids it can't run on - move it to
> + * SHARED_DSQ, which is consumed only on self cids it can't run on - send it to
>   * the kernel rescue path. One whose cids were lost after the highpri cull is
>   * likewise rescued out of HIGHPRI_DSQ below.
>   *
> @@ -663,9 +688,9 @@ static void update_core_sched_head_seq(struct task_struct *p)
>   * consuming them first. This makes minor difference only when dsp_batch is
>   * larger than 1.
>   *
> - * scx_bpf_dsq_move[_vtime]() are allowed both from ops.dispatch() and
> - * non-rq-lock holding BPF programs. As demonstration, this function is called
> - * from qmap_dispatch() and monitor_timerfn().
> + * scx_bpf_dsq_move[_vtime]() and scx_bpf_dsq_reenq_task() are allowed both from
> + * ops.dispatch() and non-rq-lock holding BPF programs. As demonstration, this
> + * function is called from qmap_dispatch() and monitor_timerfn().
>   */
>  static bool scan_shared_dsq(bool from_timer)
>  {
> @@ -677,19 +702,13 @@ static bool scan_shared_dsq(bool from_timer)
>  	bpf_for_each(scx_dsq, p, SHARED_DSQ, 0) {
>  		static u64 highpri_seq;
>  		task_ctx_t *taskc;
> -		s32 c;
>  
>  		if (!(taskc = lookup_task_ctx(p)))
>  			return false;
>  
>  		/* stranded? rescue - it can't be dispatched here either way */
>  		if (!cmask_intersects(&taskc->cpus_allowed, &qa.self_cids.mask)) {
> -			c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
> -			if (c >= 0 && c < scx_bpf_nr_cids()) {
> -				__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
> -				scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | c,
> -						 SCX_ENQ_RESCUE);
> -			}
> +			rescue_stranded(BPF_FOR_EACH_ITER, p, taskc);
>  			continue;
>  		}
>  
> @@ -723,19 +742,12 @@ static bool scan_shared_dsq(bool from_timer)
>  						      &qa.usable_cids.mask,
>  						      this_cid + 1);
>  		if (cid >= nr_cids) {
> -			s32 c;
> -
>  			/* self cids lack caps in effect yet, leave it queued */
>  			if (cmask_intersects(&taskc->cpus_allowed, &qa.self_cids.mask))
>  				continue;
>  
>  			/* stranded after the cull - rescue it from here */
> -			c = cmask_next_set_wrap(&taskc->cpus_allowed, 0);
> -			if (c >= 0 && c < nr_cids) {
> -				__sync_fetch_and_add(&qa.nr_rescue_dsp, 1);
> -				scx_bpf_dsq_move(BPF_FOR_EACH_ITER, p, SCX_DSQ_LOCAL_ON | c,
> -						 SCX_ENQ_RESCUE);
> -			}
> +			rescue_stranded(BPF_FOR_EACH_ITER, p, taskc);
>  			continue;
>  		}
>  
> -- 
> 2.56.0
> 

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

[Tejun Heo]

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.

v2: Note in the kerneldoc that another iteration's re-enqueue may take a
marked task early, and that a task on another scheduler returns false
(Andrea).

Signed-off-by: Tejun Heo <tj@kernel.org>
Reviewed-by: Andrea Righi <arighi@nvidia.com>
---
 include/linux/sched/ext.h                |  1 +
 kernel/sched/ext/ext.c                   | 82 +++++++++++++++++++++++-
 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, 89 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..14fe63428036 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);
 }
 
 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,64 @@ __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. The mark belongs to @p, so a re-enqueue
+ * scheduled by another iteration of the DSQ may re-enqueue @p before this
+ * iteration ends. @p must be on the DSQ's scheduler.
+ *
+ * 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, @p is on another scheduler, or the iterator holds no
+ * DSQ.
+ */
+__bpf_kfunc bool scx_bpf_dsq_reenq_task(struct bpf_iter_scx_dsq *it__iter,
+					struct task_struct *p)
+{
+	struct bpf_iter_scx_dsq_kern *kit = (void *)it__iter;
+	struct scx_dispatch_q *dsq = kit->dsq;
+	struct rq *src_rq, *entry_rq;
+	unsigned long flags;
+	bool marked = false;
+
+	/* NULL after a failed bpf_iter_scx_dsq_new(), see scx_dsq_move() */
+	if (unlikely(!dsq))
+		return false;
+
+	/* the cursor check needs @p's rq locked */
+	src_rq = task_rq(p);
+	entry_rq = kf_switch_rq_lock(src_rq, &flags);
+
+	scoped_guard (raw_spinlock, &dsq->lock) {
+		if (nldsq_cursor_lost_task(&kit->cursor, src_rq, dsq, p))
+			break;
+
+		if (unlikely(!scx_task_on_sched(dsq->sched, p))) {
+			scx_error(dsq->sched, "scx_bpf_dsq_reenq_task() on %s[%d] but the task belongs to a different scheduler",
+				  p->comm, p->pid);
+			break;
+		}
+
+		p->scx.dsq_flags |= SCX_TASK_DSQ_REENQ;
+		/* the walk is scheduled once, by bpf_iter_scx_dsq_destroy() */
+		kit->cursor.flags |= __SCX_DSQ_ITER_HAS_REENQ;
+		marked = true;
+	}
+
+	kf_restore_rq_lock(src_rq, entry_rq, flags);
+	return marked;
+}
+
 /**
  * scx_bpf_reenqueue_local___v2 - Re-enqueue tasks on a local DSQ
  * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index b408993af794..29e69a83d876 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1946,6 +1946,8 @@ enum scx_reenq_flags {
 
 	/* internal: kernel-issued on cap revoke, not accepted from BPF */
 	SCX_REENQ_CAP_REVOKE	= 1LLU << 1,
+	/* internal: tasks marked by scx_bpf_dsq_reenq_task() */
+	SCX_REENQ_MARKED	= 1LLU << 2,
 
 	__SCX_REENQ_FILTER_MASK	= 0xffffLLU,
 
diff --git a/tools/sched_ext/include/scx/common.bpf.h b/tools/sched_ext/include/scx/common.bpf.h
index 651bd4952647..9376d5704a47 100644
--- a/tools/sched_ext/include/scx/common.bpf.h
+++ b/tools/sched_ext/include/scx/common.bpf.h
@@ -75,6 +75,7 @@ struct task_struct *scx_bpf_dsq_peek(u64 dsq_id) __ksym __weak;
 int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id, u64 flags) __ksym __weak;
 struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it) __ksym __weak;
 void bpf_iter_scx_dsq_destroy(struct bpf_iter_scx_dsq *it) __ksym __weak;
+bool scx_bpf_dsq_reenq_task(struct bpf_iter_scx_dsq *it, struct task_struct *p) __ksym __weak;
 void scx_bpf_exit_bstr(s64 exit_code, char *fmt, unsigned long long *data, u32 data__sz) __ksym __weak;
 void scx_bpf_error_bstr(char *fmt, unsigned long long *data, u32 data_len) __ksym;
 void scx_bpf_dump_bstr(char *fmt, unsigned long long *data, u32 data_len) __ksym __weak;
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index c26024a6e2b1..254882bd8ce6 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -498,6 +498,12 @@ static inline bool __COMPAT_has_generic_reenq(void)
 	return bpf_ksym_exists(scx_bpf_dsq_reenq___compat);
 }
 
+/* v7.4: Add scx_bpf_dsq_reenq_task(). */
+static inline bool __COMPAT_has_dsq_reenq_task(void)
+{
+	return bpf_ksym_exists(scx_bpf_dsq_reenq_task);
+}
+
 /*
  * v6.19: The new void variant can be called from anywhere while the older v1
  * variant can only be called from ops.cpu_release(). The double ___ prefixes on
-- 
2.56.0

Re: [PATCH 3/3] sched_ext: scx_qmap: Also re-enqueue stranded tasks with scx_bpf_dsq_reenq_task()

[Tejun Heo]

Hello, Andrea.

On Sun, Oct 18, 2026 at 01:06:39AM +0200, Andrea Righi wrote:
> > +	if (p->pid & 1) {
>
> nit: if ((p->pid & 1) && __COMPAT_has_dsq_reenq_task())

scx_qmap is already 7.4-only. It implements ops.sub_cid_sched_updated() and
ops.sub_child_ecaps_updated(), which older kernels don't have, so it doesn't
load there anyway.

Given what qmap is used for, it'd probably make sense to rip out all the
compat stuff from it and just pair it with the latest kernel by default.

Thanks.

--
tejun

Re: [PATCHSET sched_ext/for-7.4] sched_ext: Add scx_bpf_dsq_reenq_task()

[Tejun Heo]

On Sat, Oct 17, 2026 at 12:37:35PM -1000, Tejun Heo wrote:
>  0001 sched_ext: Factor out the rq lock switch of scx_dsq_move()
>  0002 sched_ext: Add scx_bpf_dsq_reenq_task() to re-enqueue single tasks from a DSQ iteration
>  0003 sched_ext: scx_qmap: Also re-enqueue stranded tasks with scx_bpf_dsq_reenq_task()

Applied 1-3 to sched_ext/for-7.4 with the following changes:

- 2 is the v2 posted as a reply to the original.
- Added to 1 and 3:

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

Thanks.

--
tejun

Re: [PATCH 3/3] sched_ext: scx_qmap: Also re-enqueue stranded tasks with scx_bpf_dsq_reenq_task()

[Andrea Righi]

On Sat, Oct 17, 2026 at 01:35:13PM -1000, Tejun Heo wrote:
> Hello, Andrea.
> 
> On Sun, Oct 18, 2026 at 01:06:39AM +0200, Andrea Righi wrote:
> > > +	if (p->pid & 1) {
> >
> > nit: if ((p->pid & 1) && __COMPAT_has_dsq_reenq_task())
> 
> scx_qmap is already 7.4-only. It implements ops.sub_cid_sched_updated() and
> ops.sub_child_ecaps_updated(), which older kernels don't have, so it doesn't
> load there anyway.
> 
> Given what qmap is used for, it'd probably make sense to rip out all the
> compat stuff from it and just pair it with the latest kernel by default.

Ah right, it doesn't even load with on kernels. Then ripping all the compat bits
is probably a good thing to do.

Thanks,
-Andrea