From c39b30fec6a92aa0a6d410903e02ee0f620c24de Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Wed, 6 Mar 2024 11:18:41 -1000
Subject: [PATCH 1/5] scx: Add sched_ext_ops.exit_dump_len and add userspace
 plumbing

Actual feature is not implemented yet.
---
 include/linux/sched/ext.h                |  6 ++++
 kernel/sched/ext.c                       |  3 ++
 tools/sched_ext/include/scx/compat.bpf.h | 40 ++++++++++++++++++++++++
 tools/sched_ext/include/scx/compat.h     | 33 ++++++++++++++++++-
 tools/sched_ext/scx_qmap.bpf.c           | 33 ++++++++++++-------
 tools/sched_ext/scx_qmap.c               | 26 +++++++++++++--
 6 files changed, 125 insertions(+), 16 deletions(-)

diff --git a/include/linux/sched/ext.h b/include/linux/sched/ext.h
index 048a688995c9d..a2de731361d21 100644
--- a/include/linux/sched/ext.h
+++ b/include/linux/sched/ext.h
@@ -598,6 +598,12 @@ struct sched_ext_ops {
 	 */
 	u32 timeout_ms;
 
+	/**
+	 * exit_dump_len - scx_exit_info.dump buffer length. If 0, the default
+	 * value of 32768 is used.
+	 */
+	u32 exit_dump_len;
+
 	/**
 	 * name - BPF scheduler's name
 	 *
diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 07e6d71a3c069..9ca195f539796 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -4086,6 +4086,9 @@ static int bpf_scx_init_member(const struct btf_type *t,
 			return -E2BIG;
 		ops->timeout_ms = *(u32 *)(udata + moff);
 		return 1;
+	case offsetof(struct sched_ext_ops, exit_dump_len):
+		ops->exit_dump_len = *(u32 *)(udata + moff);
+		return 1;
 	}
 
 	return 0;
diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index 01b804848cdd3..46a40ae37cb90 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -35,3 +35,43 @@ static inline void __COMPAT_scx_bpf_switch_all(void)
 }
 
 #endif
+
+/*
+ * sched_ext_ops.exit_dump_len is a recent addition. Use the following
+ * definition to support older kernels. See scx_qmap for usage example.
+ */
+struct sched_ext_ops___no_exit_dump_len {
+	s32 (*select_cpu)(struct task_struct *, s32, u64);
+	void (*enqueue)(struct task_struct *, u64);
+	void (*dequeue)(struct task_struct *, u64);
+	void (*dispatch)(s32, struct task_struct *);
+	void (*runnable)(struct task_struct *, u64);
+	void (*running)(struct task_struct *);
+	void (*stopping)(struct task_struct *, bool);
+	void (*quiescent)(struct task_struct *, u64);
+	bool (*yield)(struct task_struct *, struct task_struct *);
+	bool (*core_sched_before)(struct task_struct *, struct task_struct *);
+	void (*set_weight)(struct task_struct *, u32);
+	void (*set_cpumask)(struct task_struct *, const struct cpumask *);
+	void (*update_idle)(s32, bool);
+	void (*cpu_acquire)(s32, struct scx_cpu_acquire_args *);
+	void (*cpu_release)(s32, struct scx_cpu_release_args *);
+	s32 (*init_task)(struct task_struct *, struct scx_init_task_args *);
+	void (*exit_task)(struct task_struct *, struct scx_exit_task_args *);
+	void (*enable)(struct task_struct *);
+	void (*disable)(struct task_struct *);
+	s32 (*cgroup_init)(struct cgroup *, struct scx_cgroup_init_args *);
+	void (*cgroup_exit)(struct cgroup *);
+	s32 (*cgroup_prep_move)(struct task_struct *, struct cgroup *, struct cgroup *);
+	void (*cgroup_move)(struct task_struct *, struct cgroup *, struct cgroup *);
+	void (*cgroup_cancel_move)(struct task_struct *, struct cgroup *, struct cgroup *);
+	void (*cgroup_set_weight)(struct cgroup *, u32);
+	void (*cpu_online)(s32);
+	void (*cpu_offline)(s32);
+	s32 (*init)();
+	void (*exit)(struct scx_exit_info *);
+	u32 dispatch_max_batch;
+	u64 flags;
+	u32 timeout_ms;
+	char name[128];
+};
diff --git a/tools/sched_ext/include/scx/compat.h b/tools/sched_ext/include/scx/compat.h
index 1e6e3cbe96bb3..e962b18bf964c 100644
--- a/tools/sched_ext/include/scx/compat.h
+++ b/tools/sched_ext/include/scx/compat.h
@@ -69,6 +69,34 @@ static inline bool __COMPAT_read_enum(const char *type, const char *name, u64 *v
 	__val;									\
 })
 
+static inline bool __COMPAT_struct_has_field(const char *type, const char *field)
+{
+	const struct btf_type *t;
+	const struct btf_member *m;
+	const char *n;
+	s32 tid;
+	int i;
+
+	__COMPAT_load_vmlinux_btf();
+	tid = btf__find_by_name_kind(__COMPAT_vmlinux_btf, type, BTF_KIND_STRUCT);
+	if (tid < 0)
+		return false;
+
+	t = btf__type_by_id(__COMPAT_vmlinux_btf, tid);
+	SCX_BUG_ON(!t, "btf__type_by_id(%d)", tid);
+
+	m = btf_members(t);
+
+	for (i = 0; i < BTF_INFO_VLEN(t->info); i++) {
+		n = btf__name_by_offset(__COMPAT_vmlinux_btf, m[i].name_off);
+		SCX_BUG_ON(!n, "btf__name_by_offset()");
+			if (!strcmp(n, field))
+				return true;
+	}
+
+	return false;
+}
+
 /*
  * An ops flag, %SCX_OPS_SWITCH_PARTIAL, replaced scx_bpf_switch_all() which had
  * to be called from ops.init(). To support both before and after, use both
@@ -78,4 +106,7 @@ static inline bool __COMPAT_read_enum(const char *type, const char *name, u64 *v
 #define __COMPAT_SCX_OPS_SWITCH_PARTIAL						\
 	__COMPAT_ENUM_OR_ZERO("scx_ops_flags", "SCX_OPS_SWITCH_PARTIAL")
 
-#endif
+#define __COMPAT_KERNEL_HAS_OPS_EXIT_DUMP_LEN					\
+	__COMPAT_struct_has_field("sched_ext_ops", "exit_dump_len")
+
+#endif	/* __SCX_COMPAT_H */
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index a7b0b385c52a3..b621a5f7907d4 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -383,18 +383,27 @@ void BPF_STRUCT_OPS(qmap_exit, struct scx_exit_info *ei)
 	uei_record(&uei, ei);
 }
 
+#define QMAP_OPS_INIT_COMMON							\
+	.select_cpu		= (void *)qmap_select_cpu,			\
+	.enqueue		= (void *)qmap_enqueue,				\
+	.dequeue		= (void *)qmap_dequeue,				\
+	.dispatch		= (void *)qmap_dispatch,			\
+	.core_sched_before	= (void *)qmap_core_sched_before,		\
+	.cpu_release		= (void *)qmap_cpu_release,			\
+	.init_task		= (void *)qmap_init_task,			\
+	.init			= (void *)qmap_init,				\
+	.exit			= (void *)qmap_exit,				\
+	.flags			= SCX_OPS_ENQ_LAST,				\
+	.timeout_ms		= 5000U,					\
+	.name			= "qmap",
+
 SEC(".struct_ops.link")
 struct sched_ext_ops qmap_ops = {
-	.select_cpu		= (void *)qmap_select_cpu,
-	.enqueue		= (void *)qmap_enqueue,
-	.dequeue		= (void *)qmap_dequeue,
-	.dispatch		= (void *)qmap_dispatch,
-	.core_sched_before	= (void *)qmap_core_sched_before,
-	.cpu_release		= (void *)qmap_cpu_release,
-	.init_task		= (void *)qmap_init_task,
-	.init			= (void *)qmap_init,
-	.exit			= (void *)qmap_exit,
-	.flags			= SCX_OPS_ENQ_LAST,
-	.timeout_ms		= 5000U,
-	.name			= "qmap",
+	QMAP_OPS_INIT_COMMON
+	/* .exit_dump_len will be set while loading */
+};
+
+SEC(".struct_ops.link")
+struct sched_ext_ops___no_exit_dump_len qmap_ops___no_exit_dump_len = {
+	QMAP_OPS_INIT_COMMON
 };
diff --git a/tools/sched_ext/scx_qmap.c b/tools/sched_ext/scx_qmap.c
index 1a71bee00fbf4..af586ed73cf82 100644
--- a/tools/sched_ext/scx_qmap.c
+++ b/tools/sched_ext/scx_qmap.c
@@ -19,7 +19,8 @@ const char help_fmt[] =
 "\n"
 "See the top-level comment in .bpf.c for more details.\n"
 "\n"
-"Usage: %s [-s SLICE_US] [-e COUNT] [-t COUNT] [-T COUNT] [-l COUNT] [-d PID] [-p]\n"
+"Usage: %s [-s SLICE_US] [-e COUNT] [-t COUNT] [-T COUNT] [-l COUNT] [-d PID]\n"
+"       [-D LEN] [-p]\n"
 "\n"
 "  -s SLICE_US   Override slice duration\n"
 "  -e COUNT      Trigger scx_bpf_error() after COUNT enqueues\n"
@@ -27,6 +28,7 @@ const char help_fmt[] =
 "  -T COUNT      Stall every COUNT'th kernel thread\n"
 "  -l COUNT      Trigger dispatch infinite looping after COUNT dispatches\n"
 "  -d PID        Disallow a process from switching into SCHED_EXT (-1 for self)\n"
+"  -D LEN        Set scx_exit_info.dump buffer length\n"
 "  -p            Switch only tasks on SCHED_EXT policy intead of all\n"
 "  -h            Display this help and exit\n";
 
@@ -39,6 +41,7 @@ static void sigint_handler(int dummy)
 
 int main(int argc, char **argv)
 {
+	bool has_ops_exit_dump_len = __COMPAT_KERNEL_HAS_OPS_EXIT_DUMP_LEN;
 	struct scx_qmap *skel;
 	struct bpf_link *link;
 	int opt;
@@ -51,7 +54,15 @@ int main(int argc, char **argv)
 	skel = scx_qmap__open();
 	SCX_BUG_ON(!skel, "Failed to open skel");
 
-	while ((opt = getopt(argc, argv, "s:e:t:T:l:d:ph")) != -1) {
+	if (has_ops_exit_dump_len) {
+		bpf_map__set_autocreate(skel->maps.qmap_ops, true);
+		bpf_map__set_autocreate(skel->maps.qmap_ops___no_exit_dump_len, false);
+	} else {
+		bpf_map__set_autocreate(skel->maps.qmap_ops, false);
+		bpf_map__set_autocreate(skel->maps.qmap_ops___no_exit_dump_len, true);
+	}
+
+	while ((opt = getopt(argc, argv, "s:e:t:T:l:d:D:ph")) != -1) {
 		switch (opt) {
 		case 's':
 			skel->rodata->slice_ns = strtoull(optarg, NULL, 0) * 1000;
@@ -73,6 +84,11 @@ int main(int argc, char **argv)
 			if (skel->rodata->disallow_tgid < 0)
 				skel->rodata->disallow_tgid = getpid();
 			break;
+		case 'D':
+			if (!has_ops_exit_dump_len)
+				fprintf(stderr, "WARNING: kernel doesn't support setting exit dump len\n");
+			skel->struct_ops.qmap_ops->exit_dump_len = strtoul(optarg, NULL, 0);
+			break;
 		case 'p':
 			skel->rodata->switch_partial = true;
 			skel->struct_ops.qmap_ops->flags |= __COMPAT_SCX_OPS_SWITCH_PARTIAL;
@@ -85,7 +101,11 @@ int main(int argc, char **argv)
 
 	SCX_BUG_ON(scx_qmap__load(skel), "Failed to load skel");
 
-	link = bpf_map__attach_struct_ops(skel->maps.qmap_ops);
+	if (has_ops_exit_dump_len)
+		link = bpf_map__attach_struct_ops(skel->maps.qmap_ops);
+	else
+		link = bpf_map__attach_struct_ops(skel->maps.qmap_ops___no_exit_dump_len);
+
 	SCX_BUG_ON(!link, "Failed to attach struct_ops");
 
 	while (!exit_req && !uei_exited(&skel->bss->uei)) {

From 910e0890bc44bea2f2068c592b4c8aa1fe570a74 Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Wed, 6 Mar 2024 16:11:18 -1000
Subject: [PATCH 2/5] compat.bpf.h: Add DEFINE_SCX_OPS()

---
 tools/sched_ext/include/scx/compat.bpf.h | 10 +++++++
 tools/sched_ext/scx_qmap.bpf.c           | 37 +++++++++---------------
 2 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index 46a40ae37cb90..9616b42562cbd 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -75,3 +75,13 @@ struct sched_ext_ops___no_exit_dump_len {
 	u32 timeout_ms;
 	char name[128];
 };
+
+#define DEFINE_SCX_OPS(__name, ...)						\
+	SEC(".struct_ops.link")							\
+	struct sched_ext_ops __name = {						\
+		__VA_ARGS__,							\
+	};									\
+	SEC(".struct_ops.link")							\
+	struct sched_ext_ops___no_exit_dump_len __name##___no_exit_dump_len = {	\
+		__VA_ARGS__							\
+	};									\
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index b621a5f7907d4..fd07e581d0515 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -383,27 +383,16 @@ void BPF_STRUCT_OPS(qmap_exit, struct scx_exit_info *ei)
 	uei_record(&uei, ei);
 }
 
-#define QMAP_OPS_INIT_COMMON							\
-	.select_cpu		= (void *)qmap_select_cpu,			\
-	.enqueue		= (void *)qmap_enqueue,				\
-	.dequeue		= (void *)qmap_dequeue,				\
-	.dispatch		= (void *)qmap_dispatch,			\
-	.core_sched_before	= (void *)qmap_core_sched_before,		\
-	.cpu_release		= (void *)qmap_cpu_release,			\
-	.init_task		= (void *)qmap_init_task,			\
-	.init			= (void *)qmap_init,				\
-	.exit			= (void *)qmap_exit,				\
-	.flags			= SCX_OPS_ENQ_LAST,				\
-	.timeout_ms		= 5000U,					\
-	.name			= "qmap",
-
-SEC(".struct_ops.link")
-struct sched_ext_ops qmap_ops = {
-	QMAP_OPS_INIT_COMMON
-	/* .exit_dump_len will be set while loading */
-};
-
-SEC(".struct_ops.link")
-struct sched_ext_ops___no_exit_dump_len qmap_ops___no_exit_dump_len = {
-	QMAP_OPS_INIT_COMMON
-};
+DEFINE_SCX_OPS(qmap_ops,
+	       .select_cpu		= (void *)qmap_select_cpu,
+	       .enqueue			= (void *)qmap_enqueue,
+	       .dequeue			= (void *)qmap_dequeue,
+	       .dispatch		= (void *)qmap_dispatch,
+	       .core_sched_before	= (void *)qmap_core_sched_before,
+	       .cpu_release		= (void *)qmap_cpu_release,
+	       .init_task		= (void *)qmap_init_task,
+	       .init			= (void *)qmap_init,
+	       .exit			= (void *)qmap_exit,
+	       .flags			= SCX_OPS_ENQ_LAST,
+	       .timeout_ms		= 5000U,
+	       .name			= "qmap");

From 37402bbba24263c8facd3a0586d38b34844ccb8f Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Wed, 6 Mar 2024 21:50:40 -1000
Subject: [PATCH 3/5] compat: Introduce SCX_LOAD/ATTACH() and apply to all
 example schedulers

---
 tools/sched_ext/include/scx/compat.bpf.h |  3 ++-
 tools/sched_ext/include/scx/compat.h     | 34 ++++++++++++++++++++++--
 tools/sched_ext/scx_central.bpf.c        | 34 +++++++++++-------------
 tools/sched_ext/scx_central.c            |  5 ++--
 tools/sched_ext/scx_flatcg.bpf.c         | 34 +++++++++++-------------
 tools/sched_ext/scx_flatcg.c             |  6 ++---
 tools/sched_ext/scx_qmap.bpf.c           |  2 +-
 tools/sched_ext/scx_qmap.c               | 21 ++-------------
 tools/sched_ext/scx_simple.bpf.c         | 22 +++++++--------
 tools/sched_ext/scx_simple.c             |  6 ++---
 10 files changed, 85 insertions(+), 82 deletions(-)

diff --git a/tools/sched_ext/include/scx/compat.bpf.h b/tools/sched_ext/include/scx/compat.bpf.h
index 9616b42562cbd..b86ec08787a00 100644
--- a/tools/sched_ext/include/scx/compat.bpf.h
+++ b/tools/sched_ext/include/scx/compat.bpf.h
@@ -76,7 +76,8 @@ struct sched_ext_ops___no_exit_dump_len {
 	char name[128];
 };
 
-#define DEFINE_SCX_OPS(__name, ...)						\
+/* define sched_ext_ops, see compat.h::SCX_OPS_LOAD/ATTACH() */
+#define SCX_OPS_DEFINE(__name, ...)						\
 	SEC(".struct_ops.link")							\
 	struct sched_ext_ops __name = {						\
 		__VA_ARGS__,							\
diff --git a/tools/sched_ext/include/scx/compat.h b/tools/sched_ext/include/scx/compat.h
index e962b18bf964c..b66bea556196e 100644
--- a/tools/sched_ext/include/scx/compat.h
+++ b/tools/sched_ext/include/scx/compat.h
@@ -106,7 +106,37 @@ static inline bool __COMPAT_struct_has_field(const char *type, const char *field
 #define __COMPAT_SCX_OPS_SWITCH_PARTIAL						\
 	__COMPAT_ENUM_OR_ZERO("scx_ops_flags", "SCX_OPS_SWITCH_PARTIAL")
 
-#define __COMPAT_KERNEL_HAS_OPS_EXIT_DUMP_LEN					\
-	__COMPAT_struct_has_field("sched_ext_ops", "exit_dump_len")
+/*
+ * struct sched_ext_ops can change over time. If compat.bpf.h::SCX_OPS_DEFINE()
+ * is used to define ops and compat.h::SCX_OPS_LOAD/ATTACH() are used to load
+ * and attach it, backward compatibility is automatically maintained where
+ * reasonable.
+ *
+ * - sched_ext_ops.exit_dump_len was added later. On kernels which don't support
+ *   it, the value is ignored and a warning is triggered if the value is
+ *   requested to be non-zero.
+ */
+#define SCX_OPS_LOAD(__skel, __ops_name, __scx_name) ({				\
+	if (__COMPAT_struct_has_field("sched_ext_ops", "exit_dump_len")) {	\
+		bpf_map__set_autocreate((__skel)->maps.__ops_name, true);	\
+		bpf_map__set_autocreate((__skel)->maps.__ops_name##___no_exit_dump_len, false); \
+	} else {								\
+		if ((__skel)->struct_ops.__ops_name->exit_dump_len)		\
+			fprintf(stderr, "WARNING: kernel doesn't support setting exit dump len\n"); \
+		bpf_map__set_autocreate((__skel)->maps.__ops_name, false);	\
+		bpf_map__set_autocreate((__skel)->maps.__ops_name##___no_exit_dump_len, true); \
+	}									\
+	SCX_BUG_ON(__scx_name##__load((__skel)), "Failed to load skel");	\
+})
+
+#define SCX_OPS_ATTACH(__skel, __ops_name) ({					\
+	struct bpf_link *__link;						\
+	if (__COMPAT_struct_has_field("sched_ext_ops", "exit_dump_len"))	\
+		__link = bpf_map__attach_struct_ops((__skel)->maps.__ops_name);	\
+	else									\
+		__link = bpf_map__attach_struct_ops((__skel)->maps.__ops_name##___no_exit_dump_len); \
+	SCX_BUG_ON(!__link, "Failed to attach struct_ops");			\
+	__link;									\
+})
 
 #endif	/* __SCX_COMPAT_H */
diff --git a/tools/sched_ext/scx_central.bpf.c b/tools/sched_ext/scx_central.bpf.c
index 1b2ed7e4dd4e7..3d9911f812d04 100644
--- a/tools/sched_ext/scx_central.bpf.c
+++ b/tools/sched_ext/scx_central.bpf.c
@@ -343,21 +343,19 @@ void BPF_STRUCT_OPS(central_exit, struct scx_exit_info *ei)
 	uei_record(&uei, ei);
 }
 
-SEC(".struct_ops.link")
-struct sched_ext_ops central_ops = {
-	/*
-	 * We are offloading all scheduling decisions to the central CPU and
-	 * thus being the last task on a given CPU doesn't mean anything
-	 * special. Enqueue the last tasks like any other tasks.
-	 */
-	.flags			= SCX_OPS_ENQ_LAST,
-
-	.select_cpu		= (void *)central_select_cpu,
-	.enqueue		= (void *)central_enqueue,
-	.dispatch		= (void *)central_dispatch,
-	.running		= (void *)central_running,
-	.stopping		= (void *)central_stopping,
-	.init			= (void *)central_init,
-	.exit			= (void *)central_exit,
-	.name			= "central",
-};
+SCX_OPS_DEFINE(central_ops,
+	       /*
+		* We are offloading all scheduling decisions to the central CPU
+		* and thus being the last task on a given CPU doesn't mean
+		* anything special. Enqueue the last tasks like any other tasks.
+		*/
+	       .flags			= SCX_OPS_ENQ_LAST,
+
+	       .select_cpu		= (void *)central_select_cpu,
+	       .enqueue			= (void *)central_enqueue,
+	       .dispatch		= (void *)central_dispatch,
+	       .running			= (void *)central_running,
+	       .stopping		= (void *)central_stopping,
+	       .init			= (void *)central_init,
+	       .exit			= (void *)central_exit,
+	       .name			= "central");
diff --git a/tools/sched_ext/scx_central.c b/tools/sched_ext/scx_central.c
index e9669f215a6b7..fd186ce285183 100644
--- a/tools/sched_ext/scx_central.c
+++ b/tools/sched_ext/scx_central.c
@@ -70,7 +70,7 @@ int main(int argc, char **argv)
 	RESIZE_ARRAY(data, cpu_gimme_task, skel->rodata->nr_cpu_ids);
 	RESIZE_ARRAY(data, cpu_started_at, skel->rodata->nr_cpu_ids);
 
-	SCX_BUG_ON(scx_central__load(skel), "Failed to load skel");
+	SCX_OPS_LOAD(skel, central_ops, scx_central);
 
 	/*
 	 * Affinitize the loading thread to the central CPU, as:
@@ -92,8 +92,7 @@ int main(int argc, char **argv)
 		   skel->rodata->central_cpu, skel->rodata->nr_cpu_ids - 1);
 	CPU_FREE(cpuset);
 
-	link = bpf_map__attach_struct_ops(skel->maps.central_ops);
-	SCX_BUG_ON(!link, "Failed to attach struct_ops");
+	link = SCX_OPS_ATTACH(skel, central_ops);
 
 	if (!skel->data->timer_pinned)
 		printf("WARNING : BPF_F_TIMER_CPU_PIN not available, timer not pinned to central\n");
diff --git a/tools/sched_ext/scx_flatcg.bpf.c b/tools/sched_ext/scx_flatcg.bpf.c
index e1a5c6cb538e2..9c1c3c4e6afb6 100644
--- a/tools/sched_ext/scx_flatcg.bpf.c
+++ b/tools/sched_ext/scx_flatcg.bpf.c
@@ -931,21 +931,19 @@ void BPF_STRUCT_OPS(fcg_exit, struct scx_exit_info *ei)
 	uei_record(&uei, ei);
 }
 
-SEC(".struct_ops.link")
-struct sched_ext_ops flatcg_ops = {
-	.select_cpu		= (void *)fcg_select_cpu,
-	.enqueue		= (void *)fcg_enqueue,
-	.dispatch		= (void *)fcg_dispatch,
-	.runnable		= (void *)fcg_runnable,
-	.running		= (void *)fcg_running,
-	.stopping		= (void *)fcg_stopping,
-	.quiescent		= (void *)fcg_quiescent,
-	.init_task		= (void *)fcg_init_task,
-	.cgroup_set_weight	= (void *)fcg_cgroup_set_weight,
-	.cgroup_init		= (void *)fcg_cgroup_init,
-	.cgroup_exit		= (void *)fcg_cgroup_exit,
-	.cgroup_move		= (void *)fcg_cgroup_move,
-	.exit			= (void *)fcg_exit,
-	.flags			= SCX_OPS_CGROUP_KNOB_WEIGHT | SCX_OPS_ENQ_EXITING,
-	.name			= "flatcg",
-};
+SCX_OPS_DEFINE(flatcg_ops,
+	       .select_cpu		= (void *)fcg_select_cpu,
+	       .enqueue			= (void *)fcg_enqueue,
+	       .dispatch		= (void *)fcg_dispatch,
+	       .runnable		= (void *)fcg_runnable,
+	       .running			= (void *)fcg_running,
+	       .stopping		= (void *)fcg_stopping,
+	       .quiescent		= (void *)fcg_quiescent,
+	       .init_task		= (void *)fcg_init_task,
+	       .cgroup_set_weight	= (void *)fcg_cgroup_set_weight,
+	       .cgroup_init		= (void *)fcg_cgroup_init,
+	       .cgroup_exit		= (void *)fcg_cgroup_exit,
+	       .cgroup_move		= (void *)fcg_cgroup_move,
+	       .exit			= (void *)fcg_exit,
+	       .flags			= SCX_OPS_CGROUP_KNOB_WEIGHT | SCX_OPS_ENQ_EXITING,
+	       .name			= "flatcg");
diff --git a/tools/sched_ext/scx_flatcg.c b/tools/sched_ext/scx_flatcg.c
index 8841ef89ae47f..80988014cd210 100644
--- a/tools/sched_ext/scx_flatcg.c
+++ b/tools/sched_ext/scx_flatcg.c
@@ -161,10 +161,8 @@ int main(int argc, char **argv)
 	       (double)intv_ts.tv_sec + (double)intv_ts.tv_nsec / 1000000000.0,
 	       dump_cgrps);
 
-	SCX_BUG_ON(scx_flatcg__load(skel), "Failed to load skel");
-
-	link = bpf_map__attach_struct_ops(skel->maps.flatcg_ops);
-	SCX_BUG_ON(!link, "Failed to attach struct_ops");
+	SCX_OPS_LOAD(skel, flatcg_ops, scx_flatcg);
+	link = SCX_OPS_ATTACH(skel, flatcg_ops);
 
 	while (!exit_req && !uei_exited(&skel->bss->uei)) {
 		__u64 acc_stats[FCG_NR_STATS];
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index fd07e581d0515..6c0ccae5f1933 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -383,7 +383,7 @@ void BPF_STRUCT_OPS(qmap_exit, struct scx_exit_info *ei)
 	uei_record(&uei, ei);
 }
 
-DEFINE_SCX_OPS(qmap_ops,
+SCX_OPS_DEFINE(qmap_ops,
 	       .select_cpu		= (void *)qmap_select_cpu,
 	       .enqueue			= (void *)qmap_enqueue,
 	       .dequeue			= (void *)qmap_dequeue,
diff --git a/tools/sched_ext/scx_qmap.c b/tools/sched_ext/scx_qmap.c
index af586ed73cf82..d6674864312bd 100644
--- a/tools/sched_ext/scx_qmap.c
+++ b/tools/sched_ext/scx_qmap.c
@@ -41,7 +41,6 @@ static void sigint_handler(int dummy)
 
 int main(int argc, char **argv)
 {
-	bool has_ops_exit_dump_len = __COMPAT_KERNEL_HAS_OPS_EXIT_DUMP_LEN;
 	struct scx_qmap *skel;
 	struct bpf_link *link;
 	int opt;
@@ -54,14 +53,6 @@ int main(int argc, char **argv)
 	skel = scx_qmap__open();
 	SCX_BUG_ON(!skel, "Failed to open skel");
 
-	if (has_ops_exit_dump_len) {
-		bpf_map__set_autocreate(skel->maps.qmap_ops, true);
-		bpf_map__set_autocreate(skel->maps.qmap_ops___no_exit_dump_len, false);
-	} else {
-		bpf_map__set_autocreate(skel->maps.qmap_ops, false);
-		bpf_map__set_autocreate(skel->maps.qmap_ops___no_exit_dump_len, true);
-	}
-
 	while ((opt = getopt(argc, argv, "s:e:t:T:l:d:D:ph")) != -1) {
 		switch (opt) {
 		case 's':
@@ -85,8 +76,6 @@ int main(int argc, char **argv)
 				skel->rodata->disallow_tgid = getpid();
 			break;
 		case 'D':
-			if (!has_ops_exit_dump_len)
-				fprintf(stderr, "WARNING: kernel doesn't support setting exit dump len\n");
 			skel->struct_ops.qmap_ops->exit_dump_len = strtoul(optarg, NULL, 0);
 			break;
 		case 'p':
@@ -99,14 +88,8 @@ int main(int argc, char **argv)
 		}
 	}
 
-	SCX_BUG_ON(scx_qmap__load(skel), "Failed to load skel");
-
-	if (has_ops_exit_dump_len)
-		link = bpf_map__attach_struct_ops(skel->maps.qmap_ops);
-	else
-		link = bpf_map__attach_struct_ops(skel->maps.qmap_ops___no_exit_dump_len);
-
-	SCX_BUG_ON(!link, "Failed to attach struct_ops");
+	SCX_OPS_LOAD(skel, qmap_ops, scx_qmap);
+	link = SCX_OPS_ATTACH(skel, qmap_ops);
 
 	while (!exit_req && !uei_exited(&skel->bss->uei)) {
 		long nr_enqueued = skel->bss->nr_enqueued;
diff --git a/tools/sched_ext/scx_simple.bpf.c b/tools/sched_ext/scx_simple.bpf.c
index 26b30ffd3043d..b1aba9d49b007 100644
--- a/tools/sched_ext/scx_simple.bpf.c
+++ b/tools/sched_ext/scx_simple.bpf.c
@@ -137,15 +137,13 @@ void BPF_STRUCT_OPS(simple_exit, struct scx_exit_info *ei)
 	uei_record(&uei, ei);
 }
 
-SEC(".struct_ops.link")
-struct sched_ext_ops simple_ops = {
-	.select_cpu		= (void *)simple_select_cpu,
-	.enqueue		= (void *)simple_enqueue,
-	.dispatch		= (void *)simple_dispatch,
-	.running		= (void *)simple_running,
-	.stopping		= (void *)simple_stopping,
-	.enable			= (void *)simple_enable,
-	.init			= (void *)simple_init,
-	.exit			= (void *)simple_exit,
-	.name			= "simple",
-};
+SCX_OPS_DEFINE(simple_ops,
+	       .select_cpu		= (void *)simple_select_cpu,
+	       .enqueue			= (void *)simple_enqueue,
+	       .dispatch		= (void *)simple_dispatch,
+	       .running			= (void *)simple_running,
+	       .stopping		= (void *)simple_stopping,
+	       .enable			= (void *)simple_enable,
+	       .init			= (void *)simple_init,
+	       .exit			= (void *)simple_exit,
+	       .name			= "simple");
diff --git a/tools/sched_ext/scx_simple.c b/tools/sched_ext/scx_simple.c
index 02f03bff454ba..63e030d08a464 100644
--- a/tools/sched_ext/scx_simple.c
+++ b/tools/sched_ext/scx_simple.c
@@ -74,10 +74,8 @@ int main(int argc, char **argv)
 		}
 	}
 
-	SCX_BUG_ON(scx_simple__load(skel), "Failed to load skel");
-
-	link = bpf_map__attach_struct_ops(skel->maps.simple_ops);
-	SCX_BUG_ON(!link, "Failed to attach struct_ops");
+	SCX_OPS_LOAD(skel, simple_ops, scx_simple);
+	link = SCX_OPS_ATTACH(skel, simple_ops);
 
 	while (!exit_req && !uei_exited(&skel->bss->uei)) {
 		__u64 stats[2];

From 1e7f99f9042022ce84540b14f8a0ece2a7eabce2 Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Wed, 6 Mar 2024 23:46:36 -1000
Subject: [PATCH 4/5] scx: Implement configurable user_exit_info dump buffer
 size

The default dump buffer size of 32k is okay on smaller and mostly idle
systems but it's not difficult to run over. Userspace can already
communicate the desired size through sched_ext_ops.exit_dump_len. This
commit actually implements dynamic buffer sizing. This unfortunately makes
user_exit_info interface macro-fest but the usage in schedulers doesn't get
more complicated at least.
---
 kernel/sched/ext.c                           | 27 ++++---
 tools/sched_ext/include/scx/common.h         |  3 +-
 tools/sched_ext/include/scx/compat.h         |  3 +-
 tools/sched_ext/include/scx/user_exit_info.h | 79 ++++++++++++--------
 tools/sched_ext/scx_central.bpf.c            |  4 +-
 tools/sched_ext/scx_central.c                |  6 +-
 tools/sched_ext/scx_flatcg.bpf.c             |  4 +-
 tools/sched_ext/scx_flatcg.c                 |  6 +-
 tools/sched_ext/scx_qmap.bpf.c               |  4 +-
 tools/sched_ext/scx_qmap.c                   |  6 +-
 tools/sched_ext/scx_simple.bpf.c             |  4 +-
 tools/sched_ext/scx_simple.c                 |  6 +-
 12 files changed, 84 insertions(+), 68 deletions(-)

diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 9ca195f539796..a784bec9c7f8c 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -21,7 +21,7 @@ enum scx_internal_consts {
 
 	SCX_EXIT_BT_LEN			= 64,
 	SCX_EXIT_MSG_LEN		= 1024,
-	SCX_EXIT_DUMP_LEN		= 32768,
+	SCX_EXIT_DUMP_DFL_LEN		= 32768,
 };
 
 enum scx_ops_enable_state {
@@ -3270,7 +3270,7 @@ static void free_exit_info(struct scx_exit_info *ei)
 	kfree(ei);
 }
 
-static struct scx_exit_info *alloc_exit_info(void)
+static struct scx_exit_info *alloc_exit_info(size_t exit_dump_len)
 {
 	struct scx_exit_info *ei;
 
@@ -3280,7 +3280,7 @@ static struct scx_exit_info *alloc_exit_info(void)
 
 	ei->bt = kcalloc(sizeof(ei->bt[0]), SCX_EXIT_BT_LEN, GFP_KERNEL);
 	ei->msg = kzalloc(SCX_EXIT_MSG_LEN, GFP_KERNEL);
-	ei->dump = kzalloc(SCX_EXIT_DUMP_LEN, GFP_KERNEL);
+	ei->dump = kzalloc(exit_dump_len, GFP_KERNEL);
 
 	if (!ei->bt || !ei->msg || !ei->dump) {
 		free_exit_info(ei);
@@ -3519,7 +3519,7 @@ static void scx_dump_task(struct seq_buf *s, struct task_struct *p, char marker,
 	seq_buf_commit(s, used < avail ? used : -1);
 }
 
-static void scx_dump_state(struct scx_exit_info *ei)
+static void scx_dump_state(struct scx_exit_info *ei, size_t dump_len)
 {
 	const char trunc_marker[] = "\n\n~~~~ TRUNCATED ~~~~\n";
 	unsigned long now = jiffies;
@@ -3528,7 +3528,10 @@ static void scx_dump_state(struct scx_exit_info *ei)
 	char *buf;
 	int cpu;
 
-	seq_buf_init(&s, ei->dump, SCX_EXIT_DUMP_LEN - sizeof(trunc_marker));
+	if (dump_len <= sizeof(trunc_marker))
+		return;
+
+	seq_buf_init(&s, ei->dump, dump_len - sizeof(trunc_marker));
 
 	seq_buf_printf(&s, "%s[%d] triggered exit kind %d:\n  %s (%s)\n\n",
 		       current->comm, current->pid, ei->kind, ei->reason, ei->msg);
@@ -3580,15 +3583,14 @@ static void scx_dump_state(struct scx_exit_info *ei)
 		rq_unlock(rq, &rf);
 	}
 
-	if (seq_buf_has_overflowed(&s)) {
-		used = strlen(seq_buf_str(&s));
-		memcpy(ei->dump + used, trunc_marker, sizeof(trunc_marker));
-	}
+	if (seq_buf_has_overflowed(&s))
+		memcpy(ei->dump + seq_buf_used(&s) - 1, trunc_marker,
+		       sizeof(trunc_marker));
 }
 
 static void scx_ops_error_irq_workfn(struct irq_work *irq_work)
 {
-	scx_dump_state(scx_exit_info);
+	scx_dump_state(scx_exit_info, scx_ops.exit_dump_len);
 	schedule_scx_ops_disable_work();
 }
 
@@ -3678,7 +3680,7 @@ static int scx_ops_enable(struct sched_ext_ops *ops)
 	if (ret < 0)
 		goto err;
 
-	scx_exit_info = alloc_exit_info();
+	scx_exit_info = alloc_exit_info(ops->exit_dump_len);
 	if (!scx_exit_info) {
 		ret = -ENOMEM;
 		goto err_del;
@@ -4087,7 +4089,8 @@ static int bpf_scx_init_member(const struct btf_type *t,
 		ops->timeout_ms = *(u32 *)(udata + moff);
 		return 1;
 	case offsetof(struct sched_ext_ops, exit_dump_len):
-		ops->exit_dump_len = *(u32 *)(udata + moff);
+		ops->exit_dump_len =
+			*(u32 *)(udata + moff) ?: SCX_EXIT_DUMP_DFL_LEN;
 		return 1;
 	}
 
diff --git a/tools/sched_ext/include/scx/common.h b/tools/sched_ext/include/scx/common.h
index 2a8820108f980..8d5a6775f64d9 100644
--- a/tools/sched_ext/include/scx/common.h
+++ b/tools/sched_ext/include/scx/common.h
@@ -17,8 +17,6 @@
 #include <stdint.h>
 #include <errno.h>
 
-#include "user_exit_info.h"
-
 typedef uint8_t u8;
 typedef uint16_t u16;
 typedef uint32_t u32;
@@ -67,6 +65,7 @@ typedef int64_t s64;
 			bpf_map__initial_value(skel->maps.elfsec##_##arr, &__sz); \
 	} while (0)
 
+#include "user_exit_info.h"
 #include "compat.h"
 
 #endif	/* __SCHED_EXT_COMMON_H */
diff --git a/tools/sched_ext/include/scx/compat.h b/tools/sched_ext/include/scx/compat.h
index b66bea556196e..a9f0c2b955cd9 100644
--- a/tools/sched_ext/include/scx/compat.h
+++ b/tools/sched_ext/include/scx/compat.h
@@ -116,7 +116,8 @@ static inline bool __COMPAT_struct_has_field(const char *type, const char *field
  *   it, the value is ignored and a warning is triggered if the value is
  *   requested to be non-zero.
  */
-#define SCX_OPS_LOAD(__skel, __ops_name, __scx_name) ({				\
+#define SCX_OPS_LOAD(__skel, __ops_name, __scx_name, __uei_name) ({		\
+	UEI_SET_SIZE(__skel, __ops_name, __uei_name);				\
 	if (__COMPAT_struct_has_field("sched_ext_ops", "exit_dump_len")) {	\
 		bpf_map__set_autocreate((__skel)->maps.__ops_name, true);	\
 		bpf_map__set_autocreate((__skel)->maps.__ops_name##___no_exit_dump_len, false); \
diff --git a/tools/sched_ext/include/scx/user_exit_info.h b/tools/sched_ext/include/scx/user_exit_info.h
index a51b0baf1cbf4..4ccf0d822ae89 100644
--- a/tools/sched_ext/include/scx/user_exit_info.h
+++ b/tools/sched_ext/include/scx/user_exit_info.h
@@ -11,16 +11,15 @@
 #define __USER_EXIT_INFO_H
 
 enum uei_sizes {
-	UEI_REASON_SIZE	= 128,
-	UEI_MSG_SIZE	= 1024,
-	UEI_DUMP_SIZE	= 32768,
+	UEI_REASON_LEN		= 128,
+	UEI_MSG_LEN		= 1024,
+	UEI_DUMP_DFL_LEN	= 32768,
 };
 
 struct user_exit_info {
 	int		kind;
-	char		reason[UEI_REASON_SIZE];
-	char		msg[UEI_MSG_SIZE];
-	char		dump[UEI_DUMP_SIZE];
+	char		reason[UEI_REASON_LEN];
+	char		msg[UEI_MSG_LEN];
 };
 
 #ifdef __bpf__
@@ -28,40 +27,54 @@ struct user_exit_info {
 #include "vmlinux.h"
 #include <bpf/bpf_core_read.h>
 
-static inline void uei_record(struct user_exit_info *uei,
-			      const struct scx_exit_info *ei)
-{
-	bpf_probe_read_kernel_str(uei->reason, sizeof(uei->reason), ei->reason);
-	bpf_probe_read_kernel_str(uei->msg, sizeof(uei->msg), ei->msg);
-	bpf_probe_read_kernel_str(uei->dump, sizeof(uei->dump), ei->dump);
-	/* use __sync to force memory barrier */
-	__sync_val_compare_and_swap(&uei->kind, uei->kind, ei->kind);
-}
+#define UEI_DEFINE(__name)							\
+	char RESIZABLE_ARRAY(data, __name##_dump);				\
+	const volatile u32 __name##_dump_len;					\
+	struct user_exit_info __name SEC(".data")
+
+#define UEI_RECORD(__uei_name, __ei) ({						\
+	bpf_probe_read_kernel_str(__uei_name.reason,				\
+				  sizeof(__uei_name.reason), (__ei)->reason);	\
+	bpf_probe_read_kernel_str(__uei_name.msg,				\
+				  sizeof(__uei_name.msg), (__ei)->msg);		\
+	bpf_probe_read_kernel_str(__uei_name##_dump,				\
+				  __uei_name##_dump_len, (__ei)->dump);		\
+	/* use __sync to force memory barrier */				\
+	__sync_val_compare_and_swap(&__uei_name.kind, __uei_name.kind,		\
+				    (__ei)->kind);				\
+})
 
 #else	/* !__bpf__ */
 
 #include <stdio.h>
 #include <stdbool.h>
 
-static inline bool uei_exited(struct user_exit_info *uei)
-{
-	/* use __sync to force memory barrier */
-	return __sync_val_compare_and_swap(&uei->kind, -1, -1);
-}
+/* no need to call the following explicitly if SCX_OPS_LOAD() is used */
+#define UEI_SET_SIZE(__skel, __ops_name, __uei_name) ({				\
+	u32 __len = (__skel)->struct_ops.__ops_name->exit_dump_len ?: UEI_DUMP_DFL_LEN; \
+	(__skel)->rodata->__uei_name##_dump_len = __len;			\
+	RESIZE_ARRAY(data, __uei_name##_dump, __len);				\
+})
+
+#define UEI_EXITED(__skel, __uei_name) ({					\
+	/* use __sync to force memory barrier */				\
+	__sync_val_compare_and_swap(&(__skel)->data->__uei_name.kind, -1, -1);	\
+})
 
-static inline void uei_print(const struct user_exit_info *uei)
-{
-	if (uei->dump[0] != '\0') {
-		fputs("\nDEBUG DUMP\n", stderr);
-		fputs("================================================================================\n\n", stderr);
-		fputs(uei->dump, stderr);
-		fputs("\n================================================================================\n\n", stderr);
-	}
-	fprintf(stderr, "EXIT: %s", uei->reason);
-	if (uei->msg[0] != '\0')
-		fprintf(stderr, " (%s)", uei->msg);
-	fputs("\n", stderr);
-}
+#define UEI_REPORT(__skel, __uei_name) ({					\
+	struct user_exit_info *__uei = &(__skel)->data->__uei_name;		\
+	char *__uei_dump = (__skel)->data_##__uei_name##_dump->__uei_name##_dump; \
+	if (__uei_dump[0] != '\0') {						\
+		fputs("\nDEBUG DUMP\n", stderr);				\
+		fputs("================================================================================\n\n", stderr); \
+		fputs(__uei_dump, stderr);					\
+		fputs("\n================================================================================\n\n", stderr); \
+	}									\
+	fprintf(stderr, "EXIT: %s", __uei->reason);				\
+	if (__uei->msg[0] != '\0')						\
+		fprintf(stderr, " (%s)", __uei->msg);				\
+	fputs("\n", stderr);							\
+})
 
 #endif	/* __bpf__ */
 #endif	/* __USER_EXIT_INFO_H */
diff --git a/tools/sched_ext/scx_central.bpf.c b/tools/sched_ext/scx_central.bpf.c
index 3d9911f812d04..1ab8a42edbe79 100644
--- a/tools/sched_ext/scx_central.bpf.c
+++ b/tools/sched_ext/scx_central.bpf.c
@@ -64,7 +64,7 @@ u64 nr_total, nr_locals, nr_queued, nr_lost_pids;
 u64 nr_timers, nr_dispatches, nr_mismatches, nr_retries;
 u64 nr_overflows;
 
-struct user_exit_info uei;
+UEI_DEFINE(uei);
 
 struct {
 	__uint(type, BPF_MAP_TYPE_QUEUE);
@@ -340,7 +340,7 @@ int BPF_STRUCT_OPS_SLEEPABLE(central_init)
 
 void BPF_STRUCT_OPS(central_exit, struct scx_exit_info *ei)
 {
-	uei_record(&uei, ei);
+	UEI_RECORD(uei, ei);
 }
 
 SCX_OPS_DEFINE(central_ops,
diff --git a/tools/sched_ext/scx_central.c b/tools/sched_ext/scx_central.c
index fd186ce285183..237a2c6add1bf 100644
--- a/tools/sched_ext/scx_central.c
+++ b/tools/sched_ext/scx_central.c
@@ -70,7 +70,7 @@ int main(int argc, char **argv)
 	RESIZE_ARRAY(data, cpu_gimme_task, skel->rodata->nr_cpu_ids);
 	RESIZE_ARRAY(data, cpu_started_at, skel->rodata->nr_cpu_ids);
 
-	SCX_OPS_LOAD(skel, central_ops, scx_central);
+	SCX_OPS_LOAD(skel, central_ops, scx_central, uei);
 
 	/*
 	 * Affinitize the loading thread to the central CPU, as:
@@ -97,7 +97,7 @@ int main(int argc, char **argv)
 	if (!skel->data->timer_pinned)
 		printf("WARNING : BPF_F_TIMER_CPU_PIN not available, timer not pinned to central\n");
 
-	while (!exit_req && !uei_exited(&skel->bss->uei)) {
+	while (!exit_req && !UEI_EXITED(skel, uei)) {
 		printf("[SEQ %llu]\n", seq++);
 		printf("total   :%10" PRIu64 "    local:%10" PRIu64 "   queued:%10" PRIu64 "  lost:%10" PRIu64 "\n",
 		       skel->bss->nr_total,
@@ -116,7 +116,7 @@ int main(int argc, char **argv)
 	}
 
 	bpf_link__destroy(link);
-	uei_print(&skel->bss->uei);
+	UEI_REPORT(skel, uei);
 	scx_central__destroy(skel);
 	return 0;
 }
diff --git a/tools/sched_ext/scx_flatcg.bpf.c b/tools/sched_ext/scx_flatcg.bpf.c
index 9c1c3c4e6afb6..9946b8949b7a2 100644
--- a/tools/sched_ext/scx_flatcg.bpf.c
+++ b/tools/sched_ext/scx_flatcg.bpf.c
@@ -58,7 +58,7 @@ const volatile u64 cgrp_slice_ns = SCX_SLICE_DFL;
 const volatile bool fifo_sched;
 
 u64 cvtime_now;
-struct user_exit_info uei;
+UEI_DEFINE(uei);
 
 struct {
 	__uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
@@ -928,7 +928,7 @@ void BPF_STRUCT_OPS(fcg_cgroup_move, struct task_struct *p,
 
 void BPF_STRUCT_OPS(fcg_exit, struct scx_exit_info *ei)
 {
-	uei_record(&uei, ei);
+	UEI_RECORD(uei, ei);
 }
 
 SCX_OPS_DEFINE(flatcg_ops,
diff --git a/tools/sched_ext/scx_flatcg.c b/tools/sched_ext/scx_flatcg.c
index 80988014cd210..a2b440d0a7a28 100644
--- a/tools/sched_ext/scx_flatcg.c
+++ b/tools/sched_ext/scx_flatcg.c
@@ -161,10 +161,10 @@ int main(int argc, char **argv)
 	       (double)intv_ts.tv_sec + (double)intv_ts.tv_nsec / 1000000000.0,
 	       dump_cgrps);
 
-	SCX_OPS_LOAD(skel, flatcg_ops, scx_flatcg);
+	SCX_OPS_LOAD(skel, flatcg_ops, scx_flatcg, uei);
 	link = SCX_OPS_ATTACH(skel, flatcg_ops);
 
-	while (!exit_req && !uei_exited(&skel->bss->uei)) {
+	while (!exit_req && !UEI_EXITED(skel, uei)) {
 		__u64 acc_stats[FCG_NR_STATS];
 		__u64 stats[FCG_NR_STATS];
 		float cpu_util;
@@ -213,7 +213,7 @@ int main(int argc, char **argv)
 	}
 
 	bpf_link__destroy(link);
-	uei_print(&skel->bss->uei);
+	UEI_REPORT(skel, uei);
 	scx_flatcg__destroy(skel);
 	return 0;
 }
diff --git a/tools/sched_ext/scx_qmap.bpf.c b/tools/sched_ext/scx_qmap.bpf.c
index 6c0ccae5f1933..89ac8d1260678 100644
--- a/tools/sched_ext/scx_qmap.bpf.c
+++ b/tools/sched_ext/scx_qmap.bpf.c
@@ -35,7 +35,7 @@ const volatile s32 disallow_tgid;
 
 u32 test_error_cnt;
 
-struct user_exit_info uei;
+UEI_DEFINE(uei);
 
 struct qmap {
 	__uint(type, BPF_MAP_TYPE_QUEUE);
@@ -380,7 +380,7 @@ s32 BPF_STRUCT_OPS(qmap_init)
 
 void BPF_STRUCT_OPS(qmap_exit, struct scx_exit_info *ei)
 {
-	uei_record(&uei, ei);
+	UEI_RECORD(uei, ei);
 }
 
 SCX_OPS_DEFINE(qmap_ops,
diff --git a/tools/sched_ext/scx_qmap.c b/tools/sched_ext/scx_qmap.c
index d6674864312bd..0171a5f7dcb8d 100644
--- a/tools/sched_ext/scx_qmap.c
+++ b/tools/sched_ext/scx_qmap.c
@@ -88,10 +88,10 @@ int main(int argc, char **argv)
 		}
 	}
 
-	SCX_OPS_LOAD(skel, qmap_ops, scx_qmap);
+	SCX_OPS_LOAD(skel, qmap_ops, scx_qmap, uei);
 	link = SCX_OPS_ATTACH(skel, qmap_ops);
 
-	while (!exit_req && !uei_exited(&skel->bss->uei)) {
+	while (!exit_req && !UEI_EXITED(skel, uei)) {
 		long nr_enqueued = skel->bss->nr_enqueued;
 		long nr_dispatched = skel->bss->nr_dispatched;
 
@@ -104,7 +104,7 @@ int main(int argc, char **argv)
 	}
 
 	bpf_link__destroy(link);
-	uei_print(&skel->bss->uei);
+	UEI_REPORT(skel, uei);
 	scx_qmap__destroy(skel);
 	return 0;
 }
diff --git a/tools/sched_ext/scx_simple.bpf.c b/tools/sched_ext/scx_simple.bpf.c
index b1aba9d49b007..535fe52f5b739 100644
--- a/tools/sched_ext/scx_simple.bpf.c
+++ b/tools/sched_ext/scx_simple.bpf.c
@@ -27,7 +27,7 @@ char _license[] SEC("license") = "GPL";
 const volatile bool fifo_sched;
 
 static u64 vtime_now;
-struct user_exit_info uei;
+UEI_DEFINE(uei);
 
 #define SHARED_DSQ 0
 
@@ -134,7 +134,7 @@ s32 BPF_STRUCT_OPS_SLEEPABLE(simple_init)
 
 void BPF_STRUCT_OPS(simple_exit, struct scx_exit_info *ei)
 {
-	uei_record(&uei, ei);
+	UEI_RECORD(uei, ei);
 }
 
 SCX_OPS_DEFINE(simple_ops,
diff --git a/tools/sched_ext/scx_simple.c b/tools/sched_ext/scx_simple.c
index 63e030d08a464..57abb7b4dadf1 100644
--- a/tools/sched_ext/scx_simple.c
+++ b/tools/sched_ext/scx_simple.c
@@ -74,10 +74,10 @@ int main(int argc, char **argv)
 		}
 	}
 
-	SCX_OPS_LOAD(skel, simple_ops, scx_simple);
+	SCX_OPS_LOAD(skel, simple_ops, scx_simple, uei);
 	link = SCX_OPS_ATTACH(skel, simple_ops);
 
-	while (!exit_req && !uei_exited(&skel->bss->uei)) {
+	while (!exit_req && !UEI_EXITED(skel, uei)) {
 		__u64 stats[2];
 
 		read_stats(skel, stats);
@@ -87,7 +87,7 @@ int main(int argc, char **argv)
 	}
 
 	bpf_link__destroy(link);
-	uei_print(&skel->bss->uei);
+	UEI_REPORT(skel, uei);
 	scx_simple__destroy(skel);
 	return 0;
 }

From 541f061156417689acf05cc5ec3a74ab6573eaaa Mon Sep 17 00:00:00 2001
From: Tejun Heo <tj@kernel.org>
Date: Thu, 7 Mar 2024 10:45:22 -1000
Subject: [PATCH 5/5] scx examples: Drop stray "[-p]" from usage help messages

---
 tools/sched_ext/scx_central.c | 2 +-
 tools/sched_ext/scx_flatcg.c  | 2 +-
 tools/sched_ext/scx_simple.c  | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tools/sched_ext/scx_central.c b/tools/sched_ext/scx_central.c
index 237a2c6add1bf..bc209ad6a6235 100644
--- a/tools/sched_ext/scx_central.c
+++ b/tools/sched_ext/scx_central.c
@@ -20,7 +20,7 @@ const char help_fmt[] =
 "\n"
 "See the top-level comment in .bpf.c for more details.\n"
 "\n"
-"Usage: %s [-s SLICE_US] [-c CPU] [-p]\n"
+"Usage: %s [-s SLICE_US] [-c CPU]\n"
 "\n"
 "  -s SLICE_US   Override slice duration\n"
 "  -c CPU        Override the central CPU (default: 0)\n"
diff --git a/tools/sched_ext/scx_flatcg.c b/tools/sched_ext/scx_flatcg.c
index a2b440d0a7a28..b6f0184f106d2 100644
--- a/tools/sched_ext/scx_flatcg.c
+++ b/tools/sched_ext/scx_flatcg.c
@@ -26,7 +26,7 @@ const char help_fmt[] =
 "\n"
 "See the top-level comment in .bpf.c for more details.\n"
 "\n"
-"Usage: %s [-s SLICE_US] [-i INTERVAL] [-f] [-p]\n"
+"Usage: %s [-s SLICE_US] [-i INTERVAL] [-f]\n"
 "\n"
 "  -s SLICE_US   Override slice duration\n"
 "  -i INTERVAL   Report interval\n"
diff --git a/tools/sched_ext/scx_simple.c b/tools/sched_ext/scx_simple.c
index 57abb7b4dadf1..5fdee03cfbdcd 100644
--- a/tools/sched_ext/scx_simple.c
+++ b/tools/sched_ext/scx_simple.c
@@ -17,7 +17,7 @@ const char help_fmt[] =
 "\n"
 "See the top-level comment in .bpf.c for more details.\n"
 "\n"
-"Usage: %s [-f] [-p]\n"
+"Usage: %s [-f]\n"
 "\n"
 "  -f            Use FIFO scheduling instead of weighted vtime scheduling\n"
 "  -h            Display this help and exit\n";