[PATCH 4/4] selftests/sched_ext: Add test for scx_bpf_select_cpu_pref()

[Andrea Righi <arighi@nvidia.com>]

Add a selftest to validate the behavior of the built-in idle CPU
selection policy with preferred CPUs, using scx_bpf_select_cpu_pref().

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 tools/testing/selftests/sched_ext/Makefile    |  1 +
 .../selftests/sched_ext/pref_cpus.bpf.c       | 95 +++++++++++++++++++
 tools/testing/selftests/sched_ext/pref_cpus.c | 58 +++++++++++
 3 files changed, 154 insertions(+)
 create mode 100644 tools/testing/selftests/sched_ext/pref_cpus.bpf.c
 create mode 100644 tools/testing/selftests/sched_ext/pref_cpus.c

diff --git a/tools/testing/selftests/sched_ext/Makefile b/tools/testing/selftests/sched_ext/Makefile
index f4531327b8e76..44fd180111389 100644
--- a/tools/testing/selftests/sched_ext/Makefile
+++ b/tools/testing/selftests/sched_ext/Makefile
@@ -173,6 +173,7 @@ auto-test-targets :=			\
 	maybe_null			\
 	minimal				\
 	numa				\
+	pref_cpus			\
 	prog_run			\
 	reload_loop			\
 	select_cpu_dfl			\
diff --git a/tools/testing/selftests/sched_ext/pref_cpus.bpf.c b/tools/testing/selftests/sched_ext/pref_cpus.bpf.c
new file mode 100644
index 0000000000000..460f5a54f9749
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/pref_cpus.bpf.c
@@ -0,0 +1,95 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * A scheduler that validates the behavior of scx_bpf_select_cpu_pref() by
+ * selecting idle CPUs strictly within a subset of preferred CPUs.
+ *
+ * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
+ */
+
+#include <scx/common.bpf.h>
+
+char _license[] SEC("license") = "GPL";
+
+UEI_DEFINE(uei);
+
+const volatile unsigned int __COMPAT_SCX_PICK_IDLE_IN_PREF;
+
+private(PREF_CPUS) struct bpf_cpumask __kptr * preferred_cpumask;
+
+s32 BPF_STRUCT_OPS(pref_cpus_select_cpu,
+		   struct task_struct *p, s32 prev_cpu, u64 wake_flags)
+{
+	const struct cpumask *preferred;
+	s32 cpu;
+
+	preferred = cast_mask(preferred_cpumask);
+	if (!preferred) {
+		scx_bpf_error("preferred domain not initialized");
+		return -EINVAL;
+	}
+
+	/*
+	 * Select an idle CPU strictly within the preferred domain.
+	 */
+	cpu = scx_bpf_select_cpu_pref(p, preferred, prev_cpu, wake_flags,
+				      __COMPAT_SCX_PICK_IDLE_IN_PREF);
+	if (cpu >= 0) {
+		if (scx_bpf_test_and_clear_cpu_idle(cpu))
+			scx_bpf_error("CPU %d should be marked as busy", cpu);
+
+		if (__COMPAT_SCX_PICK_IDLE_IN_PREF &&
+		    bpf_cpumask_subset(preferred, p->cpus_ptr) &&
+		    !bpf_cpumask_test_cpu(cpu, preferred))
+			scx_bpf_error("CPU %d not in the preferred domain for %d (%s)",
+				      cpu, p->pid, p->comm);
+
+		scx_bpf_dsq_insert(p, SCX_DSQ_LOCAL, SCX_SLICE_DFL, 0);
+
+		return cpu;
+	}
+
+	return prev_cpu;
+}
+
+s32 BPF_STRUCT_OPS_SLEEPABLE(pref_cpus_init)
+{
+	struct bpf_cpumask *mask;
+
+	mask = bpf_cpumask_create();
+	if (!mask)
+		return -ENOMEM;
+
+	mask = bpf_kptr_xchg(&preferred_cpumask, mask);
+	if (mask)
+		bpf_cpumask_release(mask);
+
+	bpf_rcu_read_lock();
+
+	/*
+	 * Assign the first online CPU to the preferred domain.
+	 */
+	mask = preferred_cpumask;
+	if (mask) {
+		const struct cpumask *online = scx_bpf_get_online_cpumask();
+
+		bpf_cpumask_set_cpu(bpf_cpumask_first(online), mask);
+		scx_bpf_put_cpumask(online);
+	}
+
+	bpf_rcu_read_unlock();
+
+	return 0;
+}
+
+void BPF_STRUCT_OPS(pref_cpus_exit, struct scx_exit_info *ei)
+{
+	UEI_RECORD(uei, ei);
+}
+
+SEC(".struct_ops.link")
+struct sched_ext_ops pref_cpus_ops = {
+	.select_cpu		= (void *)pref_cpus_select_cpu,
+	.init			= (void *)pref_cpus_init,
+	.exit			= (void *)pref_cpus_exit,
+	.name			= "pref_cpus",
+};
diff --git a/tools/testing/selftests/sched_ext/pref_cpus.c b/tools/testing/selftests/sched_ext/pref_cpus.c
new file mode 100644
index 0000000000000..75a09a355e1db
--- /dev/null
+++ b/tools/testing/selftests/sched_ext/pref_cpus.c
@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
+ */
+#include <bpf/bpf.h>
+#include <scx/common.h>
+#include <sys/wait.h>
+#include <unistd.h>
+#include "pref_cpus.bpf.skel.h"
+#include "scx_test.h"
+
+static enum scx_test_status setup(void **ctx)
+{
+	struct pref_cpus *skel;
+
+	skel = pref_cpus__open();
+	SCX_FAIL_IF(!skel, "Failed to open");
+	SCX_ENUM_INIT(skel);
+	skel->rodata->__COMPAT_SCX_PICK_IDLE_IN_PREF = SCX_PICK_IDLE_IN_PREF;
+	SCX_FAIL_IF(pref_cpus__load(skel), "Failed to load skel");
+
+	*ctx = skel;
+
+	return SCX_TEST_PASS;
+}
+
+static enum scx_test_status run(void *ctx)
+{
+	struct pref_cpus *skel = ctx;
+	struct bpf_link *link;
+
+	link = bpf_map__attach_struct_ops(skel->maps.pref_cpus_ops);
+	SCX_FAIL_IF(!link, "Failed to attach scheduler");
+
+	/* Just sleeping is fine, plenty of scheduling events happening */
+	sleep(1);
+
+	SCX_EQ(skel->data->uei.kind, EXIT_KIND(SCX_EXIT_NONE));
+	bpf_link__destroy(link);
+
+	return SCX_TEST_PASS;
+}
+
+static void cleanup(void *ctx)
+{
+	struct pref_cpus *skel = ctx;
+
+	pref_cpus__destroy(skel);
+}
+
+struct scx_test pref_cpus = {
+	.name = "pref_cpus",
+	.description = "Verify scx_bpf_select_cpu_pref()",
+	.setup = setup,
+	.run = run,
+	.cleanup = cleanup,
+};
+REGISTER_SCX_TEST(&pref_cpus)
-- 
2.48.1