Merge pull request #132 from sched-ext/init_flag

Add fork field to struct scx_init_task_args

include/linux/sched/ext.h

@@ -131,6 +131,11 @@ struct scx_init_task_args {
 	/* the cgroup the task is joining */
 	struct cgroup		*cgroup;
 #endif
+	/*
+	 * Set if ops.init_task() is being invoked on the fork path, as opposed
+	 * to the scheduler transition path.
+	 */
+	bool			fork;
 };
 
 /* argument container for ops.exit_task() */

kernel/sched/ext.c

@@ -2367,7 +2367,7 @@ static void scx_set_task_state(struct task_struct *p, enum scx_task_state state)
 	p->scx.flags |= state << SCX_TASK_STATE_SHIFT;
 }
 
-static int scx_ops_init_task(struct task_struct *p, struct task_group *tg)
+static int scx_ops_init_task(struct task_struct *p, struct task_group *tg, bool fork)
 {
 	int ret;
 
@@ -2376,6 +2376,7 @@ static int scx_ops_init_task(struct task_struct *p, struct task_group *tg)
 	if (SCX_HAS_OP(init_task)) {
 		struct scx_init_task_args args = {
 			SCX_INIT_TASK_ARGS_CGROUP(tg)
+			.fork = fork,
 		};
 
 		ret = SCX_CALL_OP_RET(SCX_KF_SLEEPABLE, init_task, p, &args);
@@ -2490,7 +2491,7 @@ int scx_fork(struct task_struct *p)
 	percpu_rwsem_assert_held(&scx_fork_rwsem);
 
 	if (scx_enabled())
-		return scx_ops_init_task(p, task_group(p));
+		return scx_ops_init_task(p, task_group(p), true);
 	else
 		return 0;
 }
@@ -3774,7 +3775,7 @@ static int scx_ops_enable(struct sched_ext_ops *ops)
 		get_task_struct(p);
 		spin_unlock_irq(&scx_tasks_lock);
 
-		ret = scx_ops_init_task(p, task_group(p));
+		ret = scx_ops_init_task(p, task_group(p), false);
 		if (ret) {
 			put_task_struct(p);
 			spin_lock_irq(&scx_tasks_lock);

tools/testing/selftests/scx/init_enable_count.bpf.c

@@ -13,13 +13,19 @@
 char _license[] SEC("license") = "GPL";
 
 u64 init_task_cnt, exit_task_cnt, enable_cnt, disable_cnt;
+u64 init_fork_cnt, init_transition_cnt;
 volatile const bool switch_all;
 
 s32 BPF_STRUCT_OPS_SLEEPABLE(cnt_init_task, struct task_struct *p,
 			     struct scx_init_task_args *args)
 {
 	__sync_fetch_and_add(&init_task_cnt, 1);
 
+	if (args->fork)
+		__sync_fetch_and_add(&init_fork_cnt, 1);
+	else
+		__sync_fetch_and_add(&init_transition_cnt, 1);
+
 	return 0;
 }
 

tools/testing/selftests/scx/init_enable_count.c

@@ -35,15 +35,39 @@ static enum scx_test_status run_test(bool global)
 {
 	struct init_enable_count *skel;
 	struct bpf_link *link;
-	const u32 num_children = 5;
+	const u32 num_children = 5, num_pre_forks = 1024;
 	int ret, i, status;
 	struct sched_param param = {};
-	pid_t pids[num_children];
+	pid_t pids[num_pre_forks];
 
 	skel = open_load_prog(global);
+
+	/*
+	 * Fork a bunch of children before we attach the scheduler so that we
+	 * ensure (at least in practical terms) that there are more tasks that
+	 * transition from SCHED_OTHER -> SCHED_EXT than there are tasks that
+	 * take the fork() path either below or in other processes.
+	 */
+	for (i = 0; i < num_pre_forks; i++) {
+		pids[i] = fork();
+		SCX_FAIL_IF(pids[i] < 0, "Failed to fork child");
+		if (pids[i] == 0) {
+			sleep(1);
+			exit(0);
+		}
+	}
+
 	link = bpf_map__attach_struct_ops(skel->maps.init_enable_count_ops);
 	SCX_FAIL_IF(!link, "Failed to attach struct_ops");
 
+	for (i = 0; i < num_pre_forks; i++) {
+		SCX_FAIL_IF(waitpid(pids[i], &status, 0) != pids[i],
+			    "Failed to wait for pre-forked child\n");
+
+		SCX_FAIL_IF(status != 0, "Pre-forked child %d exited with status %d\n", i,
+			    status);
+	}
+
 	/* SCHED_EXT children */
 	for (i = 0; i < num_children; i++) {
 		pids[i] = fork();
@@ -101,6 +125,14 @@ static enum scx_test_status run_test(bool global)
 		SCX_EQ(skel->bss->enable_cnt, num_children);
 		SCX_EQ(skel->bss->disable_cnt, num_children);
 	}
+	/*
+	 * We forked a ton of tasks before we attached the scheduler above, so
+	 * this should be fine. Technically it could be flaky if a ton of forks
+	 * are happening at the same time in other processes, but that should
+	 * be exceedingly unlikely.
+	 */
+	SCX_GT(skel->bss->init_transition_cnt, skel->bss->init_fork_cnt);
+	SCX_GE(skel->bss->init_fork_cnt, 2 * num_children);
 
 	bpf_link__destroy(link);
 	init_enable_count__destroy(skel);

tools/testing/selftests/scx/runner.c

@@ -21,11 +21,12 @@ const char help_fmt[] =
 "Usage: %s [-t TEST] [-h]\n"
 "\n"
 "  -t TEST       Only run tests whose name includes this string\n"
+"  -s            Include print output for skipped tests\n"
 "  -q            Don't print the test descriptions during run\n"
 "  -h            Display this help and exit\n";
 
 static volatile int exit_req;
-static bool quiet;
+static bool quiet, print_skipped;
 
 #define MAX_SCX_TESTS 2048
 
@@ -132,11 +133,14 @@ int main(int argc, char **argv)
 
 	libbpf_set_strict_mode(LIBBPF_STRICT_ALL);
 
-	while ((opt = getopt(argc, argv, "qt:h")) != -1) {
+	while ((opt = getopt(argc, argv, "qst:h")) != -1) {
 		switch (opt) {
 		case 'q':
 			quiet = true;
 			break;
+		case 's':
+			print_skipped = true;
+			break;
 		case 't':
 			filter = optarg;
 			break;
@@ -150,13 +154,21 @@ int main(int argc, char **argv)
 		enum scx_test_status status;
 		struct scx_test *test = &__scx_tests[i];
 
-		print_test_preamble(test, quiet);
-
 		if (filter && should_skip_test(test, filter)) {
-			print_test_result(test, SCX_TEST_SKIP, ++testnum);
+			/*
+			 * Printing the skipped tests and their preambles can
+			 * add a lot of noise to the runner output. Printing
+			 * this is only really useful for CI, so let's skip it
+			 * by default.
+			 */
+			if (print_skipped) {
+				print_test_preamble(test, quiet);
+				print_test_result(test, SCX_TEST_SKIP, ++testnum);
+			}
 			continue;
 		}
 
+		print_test_preamble(test, quiet);
 		status = run_test(test);
 		print_test_result(test, status, ++testnum);
 		switch (status) {