[PATCH 4/5] selftests/sched_ext: Add SCX_ASSERT_ALIVE and adopt it

[Tao Cui <cui.tao@linux.dev>]

From: Tao Cui <cuitao@kylinos.cn>

A scheduler that dies mid-test freezes the state that the test is
about to assert on, so the assertions can pass vacuously. Add
SCX_ASSERT_ALIVE(), which fails the test with the recorded exit kind
and code, and use it in maximal, create_dsq and init_enable_count.
Existing open-coded checks keep their goto-based cleanup and are not
converted. init_enable_count also clears the uei between its two
attach cycles, since the clean unregister of the first cycle would
otherwise trip the check of the second.

Signed-off-by: Tao Cui <cuitao@kylinos.cn>
---
 tools/testing/selftests/sched_ext/create_dsq.bpf.c |  8 ++++++++
 tools/testing/selftests/sched_ext/create_dsq.c     |  2 ++
 .../selftests/sched_ext/init_enable_count.bpf.c    |  8 ++++++++
 .../selftests/sched_ext/init_enable_count.c        | 13 +++++++++++++
 tools/testing/selftests/sched_ext/maximal.bpf.c    |  6 +++++-
 tools/testing/selftests/sched_ext/maximal.c        |  2 ++
 tools/testing/selftests/sched_ext/scx_test.h       | 14 ++++++++++++++
 7 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/tools/testing/selftests/sched_ext/create_dsq.bpf.c b/tools/testing/selftests/sched_ext/create_dsq.bpf.c
index 680cc4b6d8c7..24dacd1aab0c 100644
--- a/tools/testing/selftests/sched_ext/create_dsq.bpf.c
+++ b/tools/testing/selftests/sched_ext/create_dsq.bpf.c
@@ -12,6 +12,8 @@ char _license[] SEC("license") = "GPL";
 
 u32 nr_lifecycle_tests;
 
+UEI_DEFINE(uei);
+
 void BPF_STRUCT_OPS(create_dsq_exit_task, struct task_struct *p,
 		    struct scx_exit_task_args *args)
 {
@@ -84,10 +86,16 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(create_dsq_init)
 	return 0;
 }
 
+void BPF_STRUCT_OPS(create_dsq_exit, struct scx_exit_info *ei)
+{
+	UEI_RECORD(uei, ei);
+}
+
 SEC(".struct_ops.link")
 struct sched_ext_ops create_dsq_ops = {
 	.init_task		= (void *) create_dsq_init_task,
 	.exit_task		= (void *) create_dsq_exit_task,
 	.init			= (void *) create_dsq_init,
+	.exit			= (void *) create_dsq_exit,
 	.name			= "create_dsq",
 };
diff --git a/tools/testing/selftests/sched_ext/create_dsq.c b/tools/testing/selftests/sched_ext/create_dsq.c
index 422e6532ee71..84143d5710f4 100644
--- a/tools/testing/selftests/sched_ext/create_dsq.c
+++ b/tools/testing/selftests/sched_ext/create_dsq.c
@@ -35,6 +35,8 @@ static enum scx_test_status run(void *ctx)
 		return SCX_TEST_FAIL;
 	}
 
+	SCX_ASSERT_ALIVE(skel, uei);
+
 	bpf_link__destroy(link);
 
 	SCX_EQ(skel->bss->nr_lifecycle_tests, 1024);
diff --git a/tools/testing/selftests/sched_ext/init_enable_count.bpf.c b/tools/testing/selftests/sched_ext/init_enable_count.bpf.c
index 5eb9edb1837d..c0a5dd8fe183 100644
--- a/tools/testing/selftests/sched_ext/init_enable_count.bpf.c
+++ b/tools/testing/selftests/sched_ext/init_enable_count.bpf.c
@@ -15,6 +15,8 @@ char _license[] SEC("license") = "GPL";
 u64 init_task_cnt, exit_task_cnt, enable_cnt, disable_cnt;
 u64 init_fork_cnt, init_transition_cnt;
 
+UEI_DEFINE(uei);
+
 s32 BPF_STRUCT_OPS_SLEEPABLE(cnt_init_task, struct task_struct *p,
 			     struct scx_init_task_args *args)
 {
@@ -43,11 +45,17 @@ void BPF_STRUCT_OPS(cnt_disable, struct task_struct *p)
 	__sync_fetch_and_add(&disable_cnt, 1);
 }
 
+void BPF_STRUCT_OPS(cnt_exit, struct scx_exit_info *ei)
+{
+	UEI_RECORD(uei, ei);
+}
+
 SEC(".struct_ops.link")
 struct sched_ext_ops init_enable_count_ops = {
 	.init_task	= (void *) cnt_init_task,
 	.exit_task	= (void *) cnt_exit_task,
 	.enable		= (void *) cnt_enable,
 	.disable	= (void *) cnt_disable,
+	.exit		= (void *) cnt_exit,
 	.name		= "init_enable_count",
 };
diff --git a/tools/testing/selftests/sched_ext/init_enable_count.c b/tools/testing/selftests/sched_ext/init_enable_count.c
index 44577e30e764..c4a57038a1d6 100644
--- a/tools/testing/selftests/sched_ext/init_enable_count.c
+++ b/tools/testing/selftests/sched_ext/init_enable_count.c
@@ -5,6 +5,7 @@
  * Copyright (c) 2023 Tejun Heo <tj@kernel.org>
  */
 #include <signal.h>
+#include <string.h>
 #include <stdio.h>
 #include <unistd.h>
 #include <sched.h>
@@ -72,7 +73,17 @@ static enum scx_test_status run_test(bool global)
 	close(pipe_fds[1]);
 	signal(SIGCHLD, SIG_DFL);
 
+	SCX_ASSERT_ALIVE(skel, uei);
+
 	bpf_link__destroy(link);
+
+	/*
+	 * The clean unregister above is recorded in uei and would
+	 * otherwise trip the liveness check of the second cycle, which
+	 * reuses this skeleton.
+	 */
+	memset(&skel->data->uei, 0, sizeof(skel->data->uei));
+
 	SCX_GE(skel->bss->init_task_cnt, num_pre_forks);
 	SCX_GE(skel->bss->exit_task_cnt, num_pre_forks);
 
@@ -124,6 +135,8 @@ static enum scx_test_status run_test(bool global)
 			    status);
 	}
 
+	SCX_ASSERT_ALIVE(skel, uei);
+
 	bpf_link__destroy(link);
 
 	SCX_GE(skel->bss->init_task_cnt, 2 * num_children);
diff --git a/tools/testing/selftests/sched_ext/maximal.bpf.c b/tools/testing/selftests/sched_ext/maximal.bpf.c
index 04a369078aac..9ea664e2ff70 100644
--- a/tools/testing/selftests/sched_ext/maximal.bpf.c
+++ b/tools/testing/selftests/sched_ext/maximal.bpf.c
@@ -14,6 +14,8 @@ char _license[] SEC("license") = "GPL";
 
 #define DSQ_ID 0
 
+UEI_DEFINE(uei);
+
 s32 BPF_STRUCT_OPS(maximal_select_cpu, struct task_struct *p, s32 prev_cpu,
 		   u64 wake_flags)
 {
@@ -132,7 +134,9 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(maximal_init)
 }
 
 void BPF_STRUCT_OPS(maximal_exit, struct scx_exit_info *info)
-{}
+{
+	UEI_RECORD(uei, info);
+}
 
 SEC(".struct_ops.link")
 struct sched_ext_ops maximal_ops = {
diff --git a/tools/testing/selftests/sched_ext/maximal.c b/tools/testing/selftests/sched_ext/maximal.c
index 1dc369224670..6c6b80e0446e 100644
--- a/tools/testing/selftests/sched_ext/maximal.c
+++ b/tools/testing/selftests/sched_ext/maximal.c
@@ -35,6 +35,8 @@ static enum scx_test_status run(void *ctx)
 	link = bpf_map__attach_struct_ops(skel->maps.maximal_ops);
 	SCX_FAIL_IF(!link, "Failed to attach scheduler");
 
+	SCX_ASSERT_ALIVE(skel, uei);
+
 	bpf_link__destroy(link);
 
 	return SCX_TEST_PASS;
diff --git a/tools/testing/selftests/sched_ext/scx_test.h b/tools/testing/selftests/sched_ext/scx_test.h
index fe0e6dabe7a4..aa676cd5bb8d 100644
--- a/tools/testing/selftests/sched_ext/scx_test.h
+++ b/tools/testing/selftests/sched_ext/scx_test.h
@@ -125,6 +125,20 @@ void scx_test_register(struct scx_test *test);
 #define SCX_ASSERT(_x) SCX_FAIL_IF(!(_x), "Expected %s to be true (%lu)",		\
 				   #_x, (u64)(_x))
 
+/*
+ * SCX_ASSERT_ALIVE - Fail the test if the scheduler under test exited
+ * on its own between attach and this check, e.g. from a runtime error.
+ * A scheduler that dies mid-test freezes whatever state the test is
+ * about to assert on, which can make the assertions pass vacuously.
+ */
+#define SCX_ASSERT_ALIVE(__skel, __uei)						\
+	do {									\
+		if (UEI_EXITED(__skel, __uei))					\
+			SCX_FAIL("Scheduler exited unexpectedly (kind=%llu code=%lld)", \
+				 (unsigned long long)(__skel)->data->__uei.kind, \
+				 (long long)(__skel)->data->__uei.exit_code);	\
+	} while (0)
+
 #define SCX_ECODE_VAL(__ecode) ({						\
         u64 __val = 0;								\
 	bool __found = false;							\
-- 
2.53.0