From 220faa4d7473db67e7357a2e90f3a7cb102f5bdc Mon Sep 17 00:00:00 2001
From: Taylor R Campbell <riastradh@NetBSD.org>
Date: Sun, 7 Aug 2022 16:08:49 +0000
Subject: [PATCH] dtrace_sdt: Register sdt providers and probes in loaded
 modules too.

---
 external/cddl/osnet/dev/sdt/sdt.c | 58 +++++++++++++++++++++++--------
 sys/modules/dtrace/sdt/Makefile   |  1 +
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/external/cddl/osnet/dev/sdt/sdt.c b/external/cddl/osnet/dev/sdt/sdt.c
index b18ba093f3cd..377deb7e817c 100644
--- a/external/cddl/osnet/dev/sdt/sdt.c
+++ b/external/cddl/osnet/dev/sdt/sdt.c
@@ -52,6 +52,9 @@ __KERNEL_RCSID(0, "$NetBSD: sdt.c,v 1.21 2022/03/28 12:33:20 riastradh Exp $");
 #endif
 #include <sys/kernel.h>
 #include <sys/syslimits.h>
+#ifdef __NetBSD__
+#include <sys/kobj.h>
+#endif
 #ifdef __FreeBSD__
 #include <sys/linker.h>
 #include <sys/linker_set.h>
@@ -142,6 +145,8 @@ static eventhandler_tag	sdt_kld_unload_try_tag;
 #endif
 
 #ifdef __NetBSD__
+struct module_callbacks *sdt_link_set_callbacks;
+
 static char *
 strdup(const char *s, const struct malloc_type *m)
 {
@@ -410,30 +415,47 @@ __link_set_decl(sdt_providers_set, struct sdt_provider);
 __link_set_decl(sdt_probes_set, struct sdt_probe);
 __link_set_decl(sdt_argtypes_set, struct sdt_argtype);
 
-/*
- * Unfortunately we don't have linker set functions and event handlers
- * to support loading and unloading probes in modules... Currently if
- * modules have probes, if the modules are loaded when sdt is loaded
- * they will work, but they will crash unloading.
- */
 static void
-sdt_link_set_load(void)
+sdt_link_set_load(struct module *mod)
 {
 	struct sdt_provider * const *provider;
 	struct sdt_probe * const *probe;
 	struct sdt_argtype * const *argtype;
+	void *p;
+	size_t size, n;
+
+	/*
+	 * Skip builtin modules other than the main kernel.  The link
+	 * sets of all builtin modules are aggregated into one for the
+	 * kernel.
+	 */
+	if (module_source(mod) == MODULE_SOURCE_KERNEL &&
+	    strcmp(module_name(mod), "netbsd") != 0)
+		return;
 
-	__link_set_foreach(provider, sdt_providers_set) {
+	if (kobj_find_section(mod->mod_kobj, "link_set_sdt_providers_set",
+		&p, &size) != 0)
+		return;
+	n = size/sizeof(*provider);
+	for (provider = p; n --> 0; provider++) {
 		sdt_create_provider(*provider);
 	}
 
-	__link_set_foreach(probe, sdt_probes_set) {
+	if (kobj_find_section(mod->mod_kobj, "link_set_sdt_probes_set",
+		&p, &size) != 0)
+		return;
+	n = size/sizeof(*probe);
+	for (probe = p; n --> 0; probe++) {
 		(*probe)->sdtp_lf = NULL;	// XXX: we don't support it
 		sdt_create_probe(*probe);
 		TAILQ_INIT(&(*probe)->argtype_list);
 	}
 
-	__link_set_foreach(argtype, sdt_argtypes_set) {
+	if (kobj_find_section(mod->mod_kobj, "link_set_sdt_argtypes_set",
+		&p, &size) != 0)
+		return;
+	n = size/sizeof(*argtype);
+	for (argtype = p; n --> 0; argtype++) {
 		(*argtype)->probe->n_args++;
 		TAILQ_INSERT_TAIL(&(*argtype)->probe->argtype_list,
 		    *argtype, argtype_entry);
@@ -441,15 +463,21 @@ sdt_link_set_load(void)
 }
 
 static void
-sdt_link_set_unload(void)
+sdt_link_set_unload(struct module *mod)
 {
 	struct sdt_provider * const *curr, *prov, *tmp;
+	void *p;
+	size_t size, n;
 
 	/*
 	 * Go through all the providers declared in this linker file and
 	 * unregister any that aren't declared in another loaded file.
 	 */
-	__link_set_foreach(curr, sdt_providers_set) {
+	if (kobj_find_section(mod->mod_kobj, "link_set_sdt_providers_set",
+		&p, &size) != 0)
+		return;
+	n = size/sizeof(*curr);
+	for (curr = p; n --> 0; curr++) {
 		TAILQ_FOREACH_SAFE(prov, &sdt_prov_list, prov_entry, tmp) {
 			if (strcmp(prov->name, (*curr)->name) != 0)
 				continue;
@@ -488,7 +516,8 @@ sdt_load(void)
 #endif
 #ifdef __NetBSD__
 	sdt_init(dtrace_probe);
-	sdt_link_set_load();
+	sdt_link_set_callbacks = module_register_callbacks(sdt_link_set_load,
+	    sdt_link_set_unload);
 #endif
 }
 
@@ -507,8 +536,7 @@ sdt_unload(void)
 
 #ifdef __NetBSD__
 	sdt_exit();
-
-	sdt_link_set_unload();
+	module_unregister_callbacks(sdt_link_set_callbacks);
 #endif
 	TAILQ_FOREACH_SAFE(prov, &sdt_prov_list, prov_entry, tmp) {
 		ret = dtrace_unregister(prov->id);
diff --git a/sys/modules/dtrace/sdt/Makefile b/sys/modules/dtrace/sdt/Makefile
index 9f96f71a9064..5129401efa69 100644
--- a/sys/modules/dtrace/sdt/Makefile
+++ b/sys/modules/dtrace/sdt/Makefile
@@ -11,5 +11,6 @@ CPPFLAGS+=	-I${NETBSDSRCDIR}/external/cddl/osnet/sys \
 		-I${NETBSDSRCDIR}/external/cddl/osnet/dist/uts/common
 
 CPPFLAGS+=	-Wno-unknown-pragmas
+CPPFLAGS+=	-Wno-sign-compare
 
 .include <bsd.kmodule.mk>