[PATCH 3/3] tools/sched_ext: scx_pair: Verify task cgroup before dispatch

[Wanwu Li]

pair_enqueue() pushes a task's pid into the FIFO for the cgroup it
was in at enqueue time.  scx_pair has no .dequeue callback and no
.cgroup_move callback, so when a task migrates to another cgroup
between enqueue and dispatch, its stale pid remains in the old
cgroup's FIFO.

try_dispatch() pops such a stale pid, bpf_task_from_pid() returns
the (alive) task, and scx_bpf_dsq_insert() dispatches it under the
old cgroup's context.  This violates the pair scheduler's core
invariant: paired CPUs must only run tasks from the same cgroup.

Fix it by reading the task's current cgroup after
bpf_task_from_pid() and comparing it against the cgid the pair is
dispatching for.  A mismatch means the task migrated away; drop it
and retry, mirroring the existing lost-task path.

Fixes: f0262b102c7c ("tools/sched_ext: add scx_pair scheduler")
Signed-off-by: Wanwu Li <liwanwu@xxxxxxxxxx>
---
 tools/sched_ext/scx_pair.bpf.c | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tools/sched_ext/scx_pair.bpf.c b/tools/sched_ext/scx_pair.bpf.c
index 0d61b7b812db..05658d5dc796 100644
--- a/tools/sched_ext/scx_pair.bpf.c
+++ b/tools/sched_ext/scx_pair.bpf.c
@@ -519,6 +519,25 @@ static int try_dispatch(s32 cpu)
 
 	p = bpf_task_from_pid(pid);
 	if (p) {
+		struct cgroup *task_cgrp;
+		u64 task_cgid;
+
+		/*
+		 * Without a .dequeue callback, a task that migrated to
+		 * another cgroup after being enqueued leaves a stale pid
+		 * behind. Dispatching it here would run it under the wrong
+		 * cgroup. Drop it and retry.
+		 */
+		task_cgrp = scx_bpf_task_cgroup(p);
+		task_cgid = task_cgrp->kn->id;
+		bpf_cgroup_release(task_cgrp);
+
+		if (task_cgid != cgid) {
+			bpf_task_release(p);
+			__sync_fetch_and_add(&nr_missing, 1);
+			return -EAGAIN;
+		}
+
 		__sync_fetch_and_add(&nr_dispatched, 1);
 		scx_bpf_dsq_insert(p, SCX_DSQ_GLOBAL, SCX_SLICE_DFL, 0);
 		bpf_task_release(p);
-- 
2.25.1