Re: [PATCH v2 sched_ext/for-7.1] sched_ext: Invalidate dispatch decisions on CPU affinity changes

[Cheng-Yang Chou]

Hi Tejun, Andrea, and Kuba

On Mon, Mar 23, 2026 at 01:13:20PM -1000, Tejun Heo wrote:
> > The simple way to do this is to do scx_bpf_dsq_insert() at the very beginning,
> > once we know which task we would like to dispatch, and cancel the pending
> > dispatch via scx_bpf_dispatch_cancel() if any of the pre-dispatch checks fail
> > on the BPF side. This way, the "critical section" includes BPF-side checks, and
> > SCX will ignore the dispatch if there was a dequeue/enqueue racing with the
> > critical section.
> > 
> > With this solution, we can throw an error if task_can_run_on_remote_rq() is
> > false, because we know that there was no racing cpumask change (if there was,
> > it would have been caught earlier, in finish_dispatch()).
> 
> Yeah, I think this makes more sense. qseq is already there to provide
> protection against these events. It's just that the capturing of qseq is too
> late. If insert/cancel is too ugly, we can introduce another kfunc to
> capture the qseq - scx_bpf_dsq_insert_begin() or something like that - and
> stash it in a per-cpu variable. That way, qseq would be cover the "current"
> queued instance and the existing qseq mechanism would be able to reliably
> ignore the ones that lost race to dequeue.

Since this has been stale for a while, I prepared a patch to implement
scx_bpf_dsq_insert_begin() as suggested.

Is anyone else working on this? If not, I'm happy to send the formal
patch to fix this.

-- 
Cheers,
Cheng-Yang

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 0a53a0dd64bf..0215a21a02db 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -7933,6 +7933,7 @@ static void scx_dsq_insert_commit(struct scx_sched *sch, struct task_struct *p,
 {
 	struct scx_dsp_ctx *dspc = &this_cpu_ptr(sch->pcpu)->dsp_ctx;
 	struct task_struct *ddsp_task;
+	unsigned long qseq;
 
 	ddsp_task = __this_cpu_read(direct_dispatch_task);
 	if (ddsp_task) {
@@ -7945,9 +7946,16 @@ static void scx_dsq_insert_commit(struct scx_sched *sch, struct task_struct *p,
 		return;
 	}
 
+	if (dspc->insert_begin_valid) {
+		qseq = dspc->insert_begin_qseq;
+		dspc->insert_begin_valid = false;
+	} else {
+		qseq = atomic_long_read(&p->scx.ops_state) & SCX_OPSS_QSEQ_MASK;
+	}
+
 	dspc->buf[dspc->cursor++] = (struct scx_dsp_buf_ent){
 		.task = p,
-		.qseq = atomic_long_read(&p->scx.ops_state) & SCX_OPSS_QSEQ_MASK,
+		.qseq = qseq,
 		.dsq_id = dsq_id,
 		.enq_flags = enq_flags,
 	};
@@ -7955,6 +7963,39 @@ static void scx_dsq_insert_commit(struct scx_sched *sch, struct task_struct *p,
 
 __bpf_kfunc_start_defs();
 
+/**
+ * scx_bpf_dsq_insert_begin - Snapshot qseq before a dispatch decision
+ * @p: task_struct being considered for dispatch
+ * @aux: implicit BPF argument to access bpf_prog_aux hidden from BPF progs
+ *
+ * Capture @p's qseq before the BPF scheduler reads @p's properties (e.g.
+ * cpus_ptr) to make a dispatch decision. The snapshot is used by the
+ * subsequent scx_bpf_dsq_insert() call, extending the race detection window
+ * to cover any BPF-side checks between this call and the insert. If a
+ * concurrent dequeue/re-enqueue races within this window, finish_dispatch()
+ * detects the qseq mismatch and discards the stale dispatch.
+ */
+__bpf_kfunc void scx_bpf_dsq_insert_begin(struct task_struct *p,
+					  const struct bpf_prog_aux *aux)
+{
+	struct scx_sched *sch;
+	struct scx_dsp_ctx *dspc;
+
+	guard(rcu)();
+
+	sch = scx_prog_sched(aux);
+	if (unlikely(!sch))
+		return;
+
+	if (!scx_kf_allowed(sch, SCX_KF_ENQUEUE | SCX_KF_DISPATCH))
+		return;
+
+	dspc = &this_cpu_ptr(sch->pcpu)->dsp_ctx;
+	dspc->insert_begin_qseq = atomic_long_read(&p->scx.ops_state) &
+				  SCX_OPSS_QSEQ_MASK;
+	dspc->insert_begin_valid = true;
+}
+
 /**
  * scx_bpf_dsq_insert - Insert a task into the FIFO queue of a DSQ
  * @p: task_struct to insert
@@ -8134,6 +8175,7 @@ __bpf_kfunc void scx_bpf_dsq_insert_vtime(struct task_struct *p, u64 dsq_id,
 __bpf_kfunc_end_defs();
 
 BTF_KFUNCS_START(scx_kfunc_ids_enqueue_dispatch)
+BTF_ID_FLAGS(func, scx_bpf_dsq_insert_begin, KF_IMPLICIT_ARGS | KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_dsq_insert, KF_IMPLICIT_ARGS | KF_RCU)
 BTF_ID_FLAGS(func, scx_bpf_dsq_insert___v2, KF_IMPLICIT_ARGS | KF_RCU)
 BTF_ID_FLAGS(func, __scx_bpf_dsq_insert_vtime, KF_IMPLICIT_ARGS | KF_RCU)
diff --git a/kernel/sched/ext_internal.h b/kernel/sched/ext_internal.h
index 4a7ffc7f55d2..adc4f1c01b56 100644
--- a/kernel/sched/ext_internal.h
+++ b/kernel/sched/ext_internal.h
@@ -989,6 +989,8 @@ struct scx_dsp_ctx {
 	struct rq		*rq;
 	u32			cursor;
 	u32			nr_tasks;
+	unsigned long		insert_begin_qseq;
+	bool			insert_begin_valid;
 	struct scx_dsp_buf_ent	buf[];
 };
 
diff --git a/tools/sched_ext/scx_central.bpf.c b/tools/sched_ext/scx_central.bpf.c
index 64dd60b3e922..fb68a7d7e201 100644
--- a/tools/sched_ext/scx_central.bpf.c
+++ b/tools/sched_ext/scx_central.bpf.c
@@ -155,6 +155,8 @@ static bool dispatch_to_cpu(s32 cpu)
 		* reflect the migration-disabled state yet if
 		* migrate_disable_switch() hasn't run.
 		*/
+		scx_bpf_dsq_insert_begin(p);
+
 		if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr) ||
 		   (is_migration_disabled(p) && scx_bpf_task_cpu(p) != cpu)) {
 			__sync_fetch_and_add(&nr_mismatches, 1);
-- 
-- 
2.48.1