[PATCH sched_ext/for-7.2-fixes] selftests/sched_ext: Handle sleeping task affinity changes in numa test

[Kuba Piecuch <jpiecuch@google.com>]

When a sleeping task's affinity is changed, task_cpu(p) can be outside
of p->cpus_ptr until after select_task_rq() selects a new runqueue for
the task during wakeup.

Thus, the task's NUMA node determined by numa_select_cpu() can be
completely outside of the task's cpumask, leading to
scx_pick_{idle,any}_cpu_node() failing to find an eligible CPU and
returning -EBUSY. This leads to the numa.bpf.c scheduler abnormally
exiting with the following message in dmesg:

sched_ext: numa: invalid CPU -16
   scx_bpf_cpu_node+0x120/0x190
   bpf_prog_0a34b8e0f515771f_numa_select_cpu+0x108/0x14e
   bpf__sched_ext_ops_select_cpu+0x4f/0xb4
   select_task_rq_scx+0xb0/0x210
   select_task_rq+0xa0/0xd0
   __try_to_wake_up+0x196/0x650
   complete_all+0x76/0x100
   migration_cpu_stop+0x22b/0x300
   cpu_stopper_thread+0xc1/0x180
   smpboot_thread_fn+0x16b/0x230
   kthread+0x2d7/0x350
   ret_from_fork+0x1c2/0x350
   ret_from_fork_asm+0x1a/0x30

Make numa_select_cpu() robust against this case by returning @prev_cpu
if no CPU could be found in the selected NUMA node _and_ we have reason
to believe that the task's affinity was changed while it was sleeping.

Fixes: 5ae5161820e5 ("selftests/sched_ext: Add NUMA-aware scheduler test")

Signed-off-by: Kuba Piecuch <jpiecuch@google.com>
---
 tools/testing/selftests/sched_ext/numa.bpf.c | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tools/testing/selftests/sched_ext/numa.bpf.c b/tools/testing/selftests/sched_ext/numa.bpf.c
index 78cc49a7f9a6..6b4515c28aa0 100644
--- a/tools/testing/selftests/sched_ext/numa.bpf.c
+++ b/tools/testing/selftests/sched_ext/numa.bpf.c
@@ -34,7 +34,8 @@ static bool is_cpu_idle(s32 cpu, int node)
 s32 BPF_STRUCT_OPS(numa_select_cpu,
 		   struct task_struct *p, s32 prev_cpu, u64 wake_flags)
 {
-	int node = __COMPAT_scx_bpf_cpu_node(scx_bpf_task_cpu(p));
+	s32 task_cpu = scx_bpf_task_cpu(p);
+	int node = __COMPAT_scx_bpf_cpu_node(task_cpu);
 	s32 cpu;
 
 	/*
@@ -48,6 +49,16 @@ s32 BPF_STRUCT_OPS(numa_select_cpu,
 		cpu = __COMPAT_scx_bpf_pick_any_cpu_node(p->cpus_ptr, node,
 						__COMPAT_SCX_PICK_IDLE_IN_NODE);
 
+	/*
+	 * @task_cpu may be outside of p->cpus_ptr if @p's affinity
+	 * changed while it was sleeping. This means it's possible for
+	 * p->cpus_ptr to not include any CPUs from @node.
+	 * If we failed to find a cpu in @node, check if @task_cpu
+	 * is outside of p->cpus_ptr and just return @prev_cpu if it is.
+	 */
+	if (cpu < 0 && !bpf_cpumask_test_cpu(task_cpu, p->cpus_ptr))
+		return prev_cpu;
+
 	if (is_cpu_idle(cpu, node))
 		scx_bpf_error("CPU %d should be marked as busy", cpu);
 
-- 
2.55.0.229.g6434b31f56-goog