[PATCH RESEND v2] selftests/sched_ext: Validate select_cpu_and mask constraints

[PATCH RESEND v2] selftests/sched_ext: Validate select_cpu_and mask constraints

[Tianyi Chen]

[-- Attachment #1: Type: text/plain, Size: 106 bytes --]

The unchanged patch is attached as text/plain to preserve long lines
and whitespace during mail delivery.

[-- Attachment #2: scx_patch.patch --]
[-- Type: text/plain, Size: 8170 bytes --]

From: Tianyi Chen <hi@tychen.cc>
Date: Sun, 06 Sep 2026 22:34:48 +0800

The syscall test only prints the CPU selection result after checking
whether BPF_PROG_TEST_RUN succeeds.

Require -EBUSY for empty and affinity-disjoint custom masks, and allow
a legal singleton candidate to be selected or busy. Reject unexpected
errors and selections outside the custom mask or configured affinity.
Check cpus_mask because migration disabling can temporarily narrow
cpus_ptr under CONFIG_PREEMPT_RCU.

Grow dynamically allocated affinity masks until sched_getaffinity()
succeeds, retaining that allocation size for pinning and restoration.
Restore affinity on all exits after pinning, and skip only the disjoint
case when fewer than two CPUs are allowed.

Assisted-by: LLM
Signed-off-by: Tianyi Chen <hi@tychen.cc>
---
Resending from my Gmail address because delivery from my previous address
failed for the kernel.org and kernel mailing-list recipients. The patch
is unchanged from v2 and still applies to sched_ext/for-7.4.

Changes in v2:
- Check configured affinity through p->cpus_mask, avoiding temporary
  cpus_ptr narrowing when migration is disabled.
- Use zeroed, dynamically allocated affinity masks, retrying EINVAL with
  a larger size and handling restoration and cleanup on failure paths.
- Rebase onto sched_ext/for-7.4.

Validation:
- Build and allowed_cpus tests passed with CONFIG_PREEMPT_RCU, including
  single-CPU affinity and ten runs under load with preempt=full.
- A guest with 2048 possible CPUs and two online CPUs reproduced v1's
  EINVAL. V2 grew the mask from 128 to 256 bytes and passed with both
  full and single-CPU affinity.
- Injected affinity-read, pinning, restoration and BPF test-run errors
  exercised failure handling, including restoration after a BPF failure.
  Forced mask growth with malloc perturbation also passed.

v1: https://lore.kernel.org/r/20260906144029.848978-1-hi@tychen.cc
Review: https://lore.kernel.org/r/53d04980eecb0e5f6eeea25a5cf0dcf5@kernel.org

 .../selftests/sched_ext/allowed_cpus.bpf.c    |  24 +++-
 .../selftests/sched_ext/allowed_cpus.c        | 129 ++++++++++++++++--
 2 files changed, 135 insertions(+), 18 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..f14d7e5bef9 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_mask)))
+		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..773699d120e 100644
--- a/tools/testing/selftests/sched_ext/allowed_cpus.c
+++ b/tools/testing/selftests/sched_ext/allowed_cpus.c
@@ -2,7 +2,10 @@
 /*
  * Copyright (c) 2025 Andrea Righi <arighi@nvidia.com>
  */
+#define _GNU_SOURCE
 #include <bpf/bpf.h>
+#include <limits.h>
+#include <sched.h>
 #include <scx/common.h>
 #include <sys/wait.h>
 #include <unistd.h>
@@ -23,17 +26,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,29 +47,123 @@ 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;
 }
 
+/* Grow until the mask covers the kernel's CPU range, including offline CPUs. */
+static int alloc_affinity(cpu_set_t **mask, size_t *size)
+{
+	int nr_cpus = CPU_SETSIZE;
+	cpu_set_t *cpus;
+	int err;
+
+	for (;;) {
+		*size = CPU_ALLOC_SIZE(nr_cpus);
+		cpus = CPU_ALLOC(nr_cpus);
+		if (!cpus)
+			return -ENOMEM;
+		CPU_ZERO_S(*size, cpus);
+		if (!sched_getaffinity(0, *size, cpus)) {
+			*mask = cpus;
+			return nr_cpus;
+		}
+		err = errno;
+		CPU_FREE(cpus);
+		if (err != EINVAL)
+			return -err;
+		if (nr_cpus > INT_MAX / 2)
+			return -EOVERFLOW;
+		nr_cpus *= 2;
+	}
+}
+
 static enum scx_test_status run(void *ctx)
 {
 	struct allowed_cpus *skel = ctx;
-	struct bpf_link *link;
+	enum scx_test_status status = SCX_TEST_FAIL;
+	cpu_set_t *original = NULL, *pinned = NULL;
+	bool affinity_changed = false;
+	size_t size;
+	int first = -1, second = -1, cpu, nr_cpus;
+	struct bpf_link *link = NULL;
+
+	nr_cpus = alloc_affinity(&original, &size);
+	if (nr_cpus < 0) {
+		SCX_ERR("Failed to get affinity (%d)", -nr_cpus);
+		goto out;
+	}
+	pinned = CPU_ALLOC(nr_cpus);
+	if (!pinned) {
+		SCX_ERR("Failed to allocate affinity mask");
+		goto out;
+	}
+	for (cpu = 0; cpu < nr_cpus; cpu++) {
+		if (!CPU_ISSET_S(cpu, size, original))
+			continue;
+		if (first < 0) {
+			first = cpu;
+		} else {
+			second = cpu;
+			break;
+		}
+	}
+	if (first < 0) {
+		SCX_ERR("No CPU in affinity mask");
+		goto out;
+	}
 
 	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 (!link) {
+		SCX_ERR("Failed to attach scheduler");
+		goto out;
+	}
+
+	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_S(size, pinned);
+		CPU_SET_S(first, size, pinned);
+		if (sched_setaffinity(0, size, pinned)) {
+			SCX_ERR("Failed to pin task (%d)", errno);
+			goto out;
+		}
+		affinity_changed = true;
+		if (test_select_cpu_from_user(skel, "disjoint masks", second, true))
+			goto out;
+	} 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 out;
+	}
+	status = SCX_TEST_PASS;
+
+out:
+	if (affinity_changed && sched_setaffinity(0, size, original)) {
+		SCX_ERR("Failed to restore affinity (%d)", errno);
+		status = SCX_TEST_FAIL;
+	}
 	bpf_link__destroy(link);
-
-	return SCX_TEST_PASS;
+	CPU_FREE(pinned);
+	CPU_FREE(original);
+	return status;
 }
 
 static void cleanup(void *ctx)
-- 
2.55.0

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

[Tejun Heo]

Hello, Tianyi.

Applied to sched_ext/for-7.4.

Thanks.
--
tejun