[PATCH] selftests/sched_ext: Fix false positives of init_enable_count test

[Ze Gao]

Tests run in VM might be slow, so that children may exit before bpf
programs are loaded. SCX_GE(skel->bss->init_task_cnt, num_pre_forks)
would fail in this case.

For tests working in any env, use signals to control the lifetime of
children beyond bpf prog loading deterministically to get expected
results.

Signed-off-by: Ze Gao <zegao@xxxxxxxxxxx>
---
 .../selftests/sched_ext/init_enable_count.c   | 27 ++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/tools/testing/selftests/sched_ext/init_enable_count.c b/tools/testing/selftests/sched_ext/init_enable_count.c
index 97d45f1e5597..3b2c8ab8464f 100644
--- a/tools/testing/selftests/sched_ext/init_enable_count.c
+++ b/tools/testing/selftests/sched_ext/init_enable_count.c
@@ -31,6 +31,11 @@ open_load_prog(bool global)
 	return skel;
 }
 
+/* Signal handler for children */
+void sigusr1_handler(int sig)
+{
+}
+
 static enum scx_test_status run_test(bool global)
 {
 	struct init_enable_count *skel;
@@ -39,9 +44,15 @@ static enum scx_test_status run_test(bool global)
 	int ret, i, status;
 	struct sched_param param = {};
 	pid_t pids[num_pre_forks];
+	sigset_t blocked_set;
 
 	skel = open_load_prog(global);
 
+	/* Block SIGUSR1 in parent, children will inherit this*/
+	sigemptyset(&blocked_set);
+	sigaddset(&blocked_set, SIGUSR1);
+	sigprocmask(SIG_BLOCK, &blocked_set, NULL);
+
 	/*
 	 * 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
@@ -52,7 +63,13 @@ static enum scx_test_status run_test(bool global)
 		pids[i] = fork();
 		SCX_FAIL_IF(pids[i] < 0, "Failed to fork child");
 		if (pids[i] == 0) {
-			sleep(1);
+			signal(SIGUSR1, sigusr1_handler);
+			sigprocmask(SIG_UNBLOCK, &blocked_set, NULL);
+			/*
+			 * Wait indefinitely for signal, will be interrupted
+			 * by signal handler.
+			 */
+			pause();
 			exit(0);
 		}
 	}
@@ -60,6 +77,13 @@ static enum scx_test_status run_test(bool global)
 	link = bpf_map__attach_struct_ops(skel->maps.init_enable_count_ops);
 	SCX_FAIL_IF(!link, "Failed to attach struct_ops");
 
+	/* Give children time to set up handlers */
+	sleep(1);
+
+	/* Send SIGUSR1 to all children */
+	for (int i = 0; i < num_pre_forks; i++)
+		kill(pids[i], SIGUSR1);
+
 	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");
@@ -69,6 +93,7 @@ static enum scx_test_status run_test(bool global)
 	}
 
 	bpf_link__destroy(link);
+	SCX_EQ(skel->bss->init_task_cnt, skel->bss->exit_task_cnt);
 	SCX_GE(skel->bss->init_task_cnt, num_pre_forks);
 	SCX_GE(skel->bss->exit_task_cnt, num_pre_forks);
 
-- 
2.41.1