tools/sched_ext: Handle migration-disabled tasks in scx_central

When a task calls migrate_disable(), p->cpus_ptr is not updated until
migrate_disable_switch() runs during context switch, so dispatch_to_cpu()
may dequeue such a task and dispatch it to a CPU it cannot run on.

Extend the mismatch check in dispatch_to_cpu() to also test
is_migration_disabled() alongside the cpumask check, so tasks in this
window are bounced to the fallback DSQ.

Suggested-by: Andrea Righi <arighi@nvidia.com>
Suggested-by: Tejun Heo <tj@kernel.org>
Suggested-by: Kuba Piecuch <jpiecuch@google.com>
Signed-off-by: Cheng-Yang Chou <yphbchou0911@gmail.com>
Reviewed-by: Kuba Piecuch <jpiecuch@google.com>
Signed-off-by: Tejun Heo <tj@kernel.org>

diff --git a/tools/sched_ext/scx_central.bpf.c b/tools/sched_ext/scx_central.bpf.c
index 4efcce099bd52399b564aa4a899eb02569496d32..64dd60b3e92239b8097e020640cd452ae9dcf5af 100644 (file)
--- a/tools/sched_ext/scx_central.bpf.c
+++ b/tools/sched_ext/scx_central.bpf.c
@@ -149,10 +149,14 @@ static bool dispatch_to_cpu(s32 cpu)
                }
 
                /*
-                * If we can't run the task at the top, do the dumb thing and
-                * bounce it to the fallback dsq.
+                * If we can't run the task at the top for whatever reason,
+                * bounce it to the fallback dsq. Also check
+                * is_migration_disabled() explicitly as p->cpus_ptr may not
+                * reflect the migration-disabled state yet if
+                * migrate_disable_switch() hasn't run.
                 */
-               if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) {
+               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);
                        scx_bpf_dsq_insert(p, FALLBACK_DSQ_ID, SCX_SLICE_INF, 0);
                        bpf_task_release(p);