scx: Reattach scheduler on hotplug in scx_qmap

Now that we have the ability to detect restart events, let's update scx_qmap to leverage it. Signed-off-by: David Vernet <[email protected]>
sched-ext · Apr 9, 2024 · eb03eb3 · eb03eb3
1 parent 0d2385d
commit eb03eb3
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/tools/sched_ext/include/scx/user_exit_info.h b/tools/sched_ext/include/scx/user_exit_info.h
@@ -52,6 +52,9 @@ struct user_exit_info {
 #include <stdio.h>
 #include <stdbool.h>
 
+#define ECODE_USER_MASK ((1ULL << 32) - 1)
+#define ECODE_SYS_MASK ~ECODE_USER_MASK
+
 /* 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; \
@@ -64,6 +67,12 @@ struct user_exit_info {
 	__sync_val_compare_and_swap(&(__skel)->data->__uei_name.kind, -1, -1);	\
 })
 
+#define UEI_KIND(__skel, __uei_name) ((__skel)->data->__uei_name.kind)
+
+#define UEI_ECODE(__skel, __uei_name) (__skel)->data->__uei_name.exit_code
+#define UEI_ECODE_SYS(__skel, __uei_name) (UEI_ECODE(__skel, __uei_name) & ECODE_SYS_MASK)
+#define UEI_ECODE_USER(__skel, __uei_name) (UEI_ECODE(__skel, __uei_name) & ECODE_USER_MASK)
+
 #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; \
@@ -79,5 +88,14 @@ struct user_exit_info {
 	fputs("\n", stderr);							\
 })
 
+#define UEI_RESET(__skel, __uei_name) ({					\
+	struct user_exit_info *__uei = &(__skel)->data->__uei_name;		\
+	char *__uei_dump = (__skel)->data_##__uei_name##_dump->__uei_name##_dump; \
+	size_t __uei_dump_len = (__skel)->rodata->__uei_name##_dump_len;	\
+										\
+	memset(__uei, 0, sizeof(struct user_exit_info));			\
+	memset(__uei_dump, 0, __uei_dump_len);					\
+})
+
 #endif	/* __bpf__ */
 #endif	/* __USER_EXIT_INFO_H */
diff --git a/tools/sched_ext/scx_qmap.c b/tools/sched_ext/scx_qmap.c
@@ -39,6 +39,23 @@ static void sigint_handler(int dummy)
 	exit_req = 1;
 }
 
+static bool should_restart(const struct scx_qmap *skel)
+{
+	u64 kind = 0, code = 0;
+
+        if (!__COMPAT_read_enum("scx_exit_kind", "SCX_EXIT_UNREG_KERN", &kind) ||
+            kind != UEI_KIND(skel, uei)) {
+                return false;
+	}
+
+        if (!__COMPAT_read_enum("scx_exit_code", "SCX_ECODE_RESTART", &code) ||
+            (s64)code != UEI_ECODE_SYS(skel, uei)) {
+                return false;
+	}
+
+        return true;
+}
+
 int main(int argc, char **argv)
 {
 	struct scx_qmap *skel;
@@ -89,7 +106,10 @@ int main(int argc, char **argv)
 	}
 
 	SCX_OPS_LOAD(skel, qmap_ops, scx_qmap, uei);
+
+reattach:
 	link = SCX_OPS_ATTACH(skel, qmap_ops);
+	SCX_BUG_ON(!link, "Failed to attach link");
 
 	while (!exit_req && !UEI_EXITED(skel, uei)) {
 		long nr_enqueued = skel->bss->nr_enqueued;
@@ -105,6 +125,12 @@ int main(int argc, char **argv)
 
 	bpf_link__destroy(link);
 	UEI_REPORT(skel, uei);
+
+	if (should_restart(skel)) {
+		UEI_RESET(skel, uei);
+		goto reattach;
+	}
+
 	scx_qmap__destroy(skel);
 	return 0;
 }