scx_rustland: use scx_bpf_dispatch_cancel()

Use scx_bpf_dispatch_cancel() to invalidate dispatches on wrong per-CPU
DSQ, due to cpumask race conditions.

This prevents dispatching tasks to CPU that cannot be used according to
the task's cpumask.

With this applied the scheduler passed all the `stress-ng --race-sched`
stress tests.

Link: sched-ext/sched_ext#135
Signed-off-by: Andrea Righi <[email protected]>

scheds/rust/scx_rustland/src/bpf/main.bpf.c

@@ -299,6 +299,9 @@ static void dispatch_task(struct task_struct *p, u64 dsq_id, u64 enq_flags)
 	switch (dsq_id) {
 	case SCX_DSQ_LOCAL:
 	case SHARED_DSQ:
+		scx_bpf_dispatch(p, dsq_id, slice, enq_flags);
+		dbg_msg("dispatch: pid=%d (%s) dsq=%llu",
+			p->pid, p->comm, dsq_id);
 		break;
 	default:
 		/*
@@ -311,18 +314,26 @@ static void dispatch_task(struct task_struct *p, u64 dsq_id, u64 enq_flags)
 		 * in the meantime), we can simply ignore the task, as it has
 		 * been dequeued and re-enqueued, so it will pick a valid CPU.
 		 */
+		scx_bpf_dispatch(p, dsq_id, slice, enq_flags);
+
 		cpu = dsq_to_cpu(dsq_id);
 		if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) {
 			cpu = scx_bpf_task_cpu(p);
 			if (!bpf_cpumask_test_cpu(cpu, p->cpus_ptr))
-				return;
-			dsq_id = dsq_to_cpu(cpu);
+				dsq_id = SHARED_DSQ;
+			else
+				dsq_id = cpu_to_dsq(cpu);
+
+			scx_bpf_dispatch_cancel();
+			scx_bpf_dispatch(p, dsq_id, slice, enq_flags);
 		}
+		dbg_msg("dispatch: pid=%d (%s) dsq=%llu",
+			p->pid, p->comm, dsq_id);
+
+		/* Wake up the target CPU (only if idle) */
 		__COMPAT_scx_bpf_kick_cpu_IDLE(cpu);
 		break;
 	}
-	dbg_msg("dispatch: pid=%d (%s) dsq=%llu", p->pid, p->comm, dsq_id);
-	scx_bpf_dispatch(p, dsq_id, slice, enq_flags);
 }
 
 /*