[PATCH] selftests/sched_ext: Validate select_cpu_and mask constraints

[PATCH] selftests/sched_ext: Validate select_cpu_and mask constraints

[Tianyi Chen]

The syscall test checks whether BPF_PROG_TEST_RUN succeeds, but only
prints the CPU selection result returned by the BPF program.

Exercise empty and affinity-disjoint custom masks and require -EBUSY.
Also exercise a legal singleton mask, allowing its CPU to be busy.
Interpret retval as signed and reject unexpected errors or selection
outside the custom mask or the task's affinity.

Choose CPUs from the runner's affinity, restore it after the disjoint
case, and skip only that case when fewer than two CPUs are available.

The test passes in a two-CPU VM with full and single-CPU affinity.
Forcing selection of CPU 0 against an empty mask makes it fail with
-ERANGE; the restored test passes again.

Assisted-by: LLM
Signed-off-by: Tianyi Chen <hi@tychen.cc>
---
 .../selftests/sched_ext/allowed_cpus.bpf.c    | 24 +++++-
 .../selftests/sched_ext/allowed_cpus.c        | 77 ++++++++++++++++---
 2 files changed, 86 insertions(+), 15 deletions(-)

diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
index 9dd72d0da29..08c2fe0e1c3 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.bpf.c
@@ -147,23 +147,41 @@ void BPF_STRUCT_OPS(allowed_cpus_exit, struct scx_exit_info *ei)
 }
 
 struct task_cpu_arg {
-	pid_t pid;
+	u64 pid;
+	s64 custom_cpu;
 };
 
 SEC("syscall")
 int select_cpu_from_user(struct task_cpu_arg *input)
 {
 	struct task_struct *p;
-	int cpu;
+	struct bpf_cpumask *mask;
+	s32 cpu;
 
 	p = bpf_task_from_pid(input->pid);
 	if (!p)
 		return -EINVAL;
 
+	mask = bpf_cpumask_create();
+	if (!mask) {
+		bpf_task_release(p);
+		return -ENOMEM;
+	}
+
+	/* A negative custom_cpu leaves the custom mask empty. */
+	if (input->custom_cpu >= 0)
+		bpf_cpumask_set_cpu(input->custom_cpu, mask);
+
 	bpf_rcu_read_lock();
-	cpu = scx_bpf_select_cpu_and(p, bpf_get_smp_processor_id(), 0, p->cpus_ptr, 0);
+	cpu = scx_bpf_select_cpu_and(p, bpf_get_smp_processor_id(), 0,
+				     cast_mask(mask), 0);
+	if (cpu >= 0 &&
+	    (!bpf_cpumask_test_cpu(cpu, cast_mask(mask)) ||
+	     !bpf_cpumask_test_cpu(cpu, p->cpus_ptr)))
+		cpu = -ERANGE;
 	bpf_rcu_read_unlock();
 
+	bpf_cpumask_release(mask);
 	bpf_task_release(p);
 
 	return cpu;
diff --git a/tools/testing/selftests/sched_ext/allowed_cpus.c b/tools/testing/selftests/sched_ext/allowed_cpus.c
index 093f285ab4b..cc08c42ad95 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.c
@@ -2,7 +2,9 @@
 /*
  * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
  */
+#define _GNU_SOURCE
 #include <bpf/bpf.h>
+#include <sched.h>
 #include <scx/common.h>
 #include <sys/wait.h>
 #include <unistd.h>
@@ -23,17 +25,19 @@ static enum scx_test_status setup(void **ctx)
 	return SCX_TEST_PASS;
 }
 
-static int test_select_cpu_from_user(const struct allowed_cpus *skel)
+static int test_select_cpu_from_user(const struct allowed_cpus *skel,
+				     const char *name, int custom_cpu,
+				     bool expect_busy)
 {
 	int fd, ret;
-	__u64 args[1];
+	__s32 cpu;
+	__u64 args[] = { getpid(), (__u64)(__s64)custom_cpu };
 
 	LIBBPF_OPTS(bpf_test_run_opts, attr,
 		.ctx_in = args,
 		.ctx_size_in = sizeof(args),
 	);
 
-	args[0] = getpid();
 	fd = bpf_program__fd(skel->progs.select_cpu_from_user);
 	if (fd < 0)
 		return fd;
@@ -42,7 +46,13 @@ static int test_select_cpu_from_user(const struct allowed_cpus *skel)
 	if (ret < 0)
 		return ret;
 
-	fprintf(stderr, "%s: CPU %d\n", __func__, attr.retval);
+	/* test_run returns the signed BPF result through an unsigned field. */
+	cpu = (__s32)attr.retval;
+	if ((expect_busy && cpu != -EBUSY) ||
+	    (!expect_busy && cpu != -EBUSY && cpu != custom_cpu)) {
+		SCX_ERR("%s: unexpected CPU selection result %d", name, cpu);
+		return -EINVAL;
+	}
 
 	return 0;
 }
@@ -50,21 +60,64 @@ static int test_select_cpu_from_user(const struct allowed_cpus *skel)
 static enum scx_test_status run(void *ctx)
 {
 	struct allowed_cpus *skel = ctx;
+	enum scx_test_status status = SCX_TEST_FAIL;
+	cpu_set_t original, pinned;
+	int first = -1, second = -1, cpu;
 	struct bpf_link *link;
 
+	SCX_FAIL_IF(sched_getaffinity(0, sizeof(original), &original),
+		    "Failed to get affinity (%d)", errno);
+	for (cpu = 0; cpu < CPU_SETSIZE; cpu++) {
+		if (!CPU_ISSET(cpu, &original))
+			continue;
+		if (first < 0)
+			first = cpu;
+		else {
+			second = cpu;
+			break;
+		}
+	}
+	SCX_FAIL_IF(first < 0, "No CPU in affinity mask");
+
 	link = bpf_map__attach_struct_ops(skel->maps.allowed_cpus_ops);
 	SCX_FAIL_IF(!link, "Failed to attach scheduler");
 
-	/* Pick an idle CPU from user-space */
-	SCX_FAIL_IF(test_select_cpu_from_user(skel), "Failed to pick idle CPU");
-
-	/* Just sleeping is fine, plenty of scheduling events happening */
+	if (test_select_cpu_from_user(skel, "empty mask", -1, true))
+		goto out;
+
+	/* A legal candidate may be busy; selection need not succeed. */
+	if (test_select_cpu_from_user(skel, "legal candidate", first, false))
+		goto out;
+
+	if (second >= 0) {
+		CPU_ZERO(&pinned);
+		CPU_SET(first, &pinned);
+		if (sched_setaffinity(0, sizeof(pinned), &pinned)) {
+			SCX_ERR("Failed to pin task (%d)", errno);
+			goto out;
+		}
+		if (test_select_cpu_from_user(skel, "disjoint masks", second, true))
+			goto restore;
+	} else {
+		fprintf(stderr, "Skipping disjoint masks: need two allowed CPUs\n");
+	}
+
+	/* Just sleeping is fine, plenty of scheduling events happening. */
 	sleep(1);
-
-	SCX_EQ(skel->data->uei.kind, EXIT_KIND(SCX_EXIT_NONE));
+	if (skel->data->uei.kind != EXIT_KIND(SCX_EXIT_NONE)) {
+		SCX_ERR("Scheduler exited unexpectedly");
+		goto restore;
+	}
+	status = SCX_TEST_PASS;
+
+restore:
+	if (second >= 0 && sched_setaffinity(0, sizeof(original), &original)) {
+		SCX_ERR("Failed to restore affinity (%d)", errno);
+		status = SCX_TEST_FAIL;
+	}
+out:
 	bpf_link__destroy(link);
-
-	return SCX_TEST_PASS;
+	return status;
 }
 
 static void cleanup(void *ctx)
-- 
2.55.0

Re: [PATCH] selftests/sched_ext: Validate select_cpu_and mask constraints

[Tejun Heo]

Hello, Tianyi.

On Sun, Sep 06, 2026 at 10:40:29PM +0800, Tianyi Chen wrote:
> +	if (cpu >= 0 &&
> +	    (!bpf_cpumask_test_cpu(cpu, cast_mask(mask)) ||
> +	     !bpf_cpumask_test_cpu(cpu, p->cpus_ptr)))
> +		cpu = -ERANGE;

The syscall runs with migration disabled. With CONFIG_PREEMPT_RCU,
preemption between selection and this check can make
migrate_disable_switch() narrow p->cpus_ptr to the running CPU, rejecting a
valid selection on another CPU. Please check the requested affinity instead.

> +	SCX_FAIL_IF(sched_getaffinity(0, sizeof(original), &original),
> +		    "Failed to get affinity (%d)", errno);

Please use dynamically allocated masks. The fixed cpu_set_t makes
sched_getaffinity() fail with EINVAL when nr_cpu_ids exceeds 1024, even if
the task only uses low-numbered CPUs.

Thanks.
-- 
tejun