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

Hi Tejun,

Please disregard the previous Gmail resend in this thread. While
investigating the same delivery path for an MM series, I confirmed that
outbound delivery wraps long lines and damages patch whitespace.

I have resent the unchanged v2 as a text/plain attachment:
https://lore.kernel.org/r/CACGbirTEtYnxiZW0EnOCHbZt66=t-A_Nppp0ChRfM9G2ZD5+aQ@mail.gmail.com

The externally received attachment preserves the patch and applies
through b4 and git am to the prepared tree.

Sorry for the extra resend.

Thanks,
Tianyi

Hi Tejun,

Following up on your review:
https://lore.kernel.org/r/53d04980eecb0e5f6eeea25a5cf0dcf5@kernel.org

My earlier v2 submission and reply were rejected for the kernel.org and
kernel mailing-list recipients, so I have resent v2 from this address at
the start of this thread. The code is unchanged, and for-7.4 has not moved.

V2 checks the selected CPU against p->cpus_mask and uses dynamically
allocated affinity masks, growing them when sched_getaffinity() returns
EINVAL. The masks are zeroed before use, and affinity restoration and
resource cleanup are handled on failure paths.

I reproduced the fixed-size mask failure in a guest with 2048 possible
CPUs and two online CPUs. V2 passes there, including with single-CPU
affinity. Testing with CONFIG_PREEMPT_RCU and injected failure paths
also passed.

Thanks,
Tianyi