FreeBSD's DTrace implementation has a few problems with respect to handling

probes declared in a kernel module when that module is unloaded. In
particular,

* Unloading a module with active SDT probes will cause a panic. [1]
* A module's (FBT/SDT) probes aren't destroyed when the module is unloaded;
  trying to use them after the fact will generally cause a panic.

This change fixes both problems by porting the DTrace module load/unload
handlers from illumos and registering them with the corresponding
EVENTHANDLER(9) handlers. This allows the DTrace framework to destroy all
probes defined in a module when that module is unloaded, and to prevent a
module unload from proceeding if some of its probes are active. The latter
problem has already been fixed for FBT probes by checking lf->nenabled in
kern_kldunload(), but moving the check into the DTrace framework generalizes
it to all kernel providers and also fixes a race in the current
implementation (since a probe may be activated between the check and the
call to linker_file_unload()).

Additionally, the SDT implementation has been reworked to define SDT
providers/probes/argtypes in linker sets rather than using SYSINIT/SYSUNINIT
to create and destroy SDT probes when a module is loaded or unloaded. This
simplifies things quite a bit since it means that pretty much all of the SDT
code can live in sdt.ko, and since it becomes easier to integrate SDT with
the DTrace framework. Furthermore, this allows FreeBSD to be quite flexible
in that SDT providers spanning multiple modules can be created on the fly
when a module is loaded; at the moment it looks like illumos' SDT
implementation requires all SDT probes to be statically defined in a single
kernel table.

PR:		166927, 166926, 166928
Reported by:	davide [1]
Reviewed by:	avg, trociny (earlier version)
MFC after:	1 month

4 files changed, 227 insertions, 79 deletions

diff --git a/sys/cddl/dev/dtrace/dtrace_load.c b/sys/cddl/dev/dtrace/dtrace_load.c
index 9c5681a..a153e7e 100644
--- a/sys/cddl/dev/dtrace/dtrace_load.c
+++ b/sys/cddl/dev/dtrace/dtrace_load.c
@@ -56,6 +56,12 @@ dtrace_load(void *dummy)
 	/* Hang our hook for exceptions. */
 	dtrace_invop_init();
 
+	/* Register callbacks for module load and unload events. */
+	dtrace_modload_tag = EVENTHANDLER_REGISTER(mod_load,
+	    dtrace_mod_load, NULL, EVENTHANDLER_PRI_ANY);
+	dtrace_modunload_tag = EVENTHANDLER_REGISTER(mod_unload,
+	    dtrace_mod_unload, NULL, EVENTHANDLER_PRI_ANY);
+
 	/*
 	 * Initialise the mutexes without 'witness' because the dtrace
 	 * code is mostly written to wait for memory. To have the
diff --git a/sys/cddl/dev/dtrace/dtrace_unload.c b/sys/cddl/dev/dtrace/dtrace_unload.c
index 33d7c40..f8d2ec6 100644
--- a/sys/cddl/dev/dtrace/dtrace_unload.c
+++ b/sys/cddl/dev/dtrace/dtrace_unload.c
@@ -67,6 +67,8 @@ dtrace_unload()
 	}
 
 	dtrace_provider = NULL;
+	EVENTHANDLER_DEREGISTER(mod_load, dtrace_modload_tag);
+	EVENTHANDLER_DEREGISTER(mod_unload, dtrace_modunload_tag);
 
 	if ((state = dtrace_anon_grab()) != NULL) {
 		/*
diff --git a/sys/cddl/dev/fbt/fbt.c b/sys/cddl/dev/fbt/fbt.c
index b828163..7b278f7 100644
--- a/sys/cddl/dev/fbt/fbt.c
+++ b/sys/cddl/dev/fbt/fbt.c
@@ -1335,6 +1335,15 @@ fbt_getargdesc(void *arg __unused, dtrace_id_t id __unused, void *parg, dtrace_a
 	return;
 }
 
+static int
+fbt_linker_file_cb(linker_file_t lf, void *arg)
+{
+
+	fbt_provide_module(arg, lf);
+
+	return (0);
+}
+
 static void
 fbt_load(void *dummy)
 {
@@ -1359,8 +1368,10 @@ fbt_load(void *dummy)
 	if (dtrace_register("fbt", &fbt_attr, DTRACE_PRIV_USER,
 	    NULL, &fbt_pops, NULL, &fbt_id) != 0)
 		return;
-}
 
+	/* Create probes for the kernel and already-loaded modules. */
+	linker_file_foreach(fbt_linker_file_cb, NULL);
+}
 
 static int
 fbt_unload()
diff --git a/sys/cddl/dev/sdt/sdt.c b/sys/cddl/dev/sdt/sdt.c
index 96ab550..493ab9b 100644
--- a/sys/cddl/dev/sdt/sdt.c
+++ b/sys/cddl/dev/sdt/sdt.c
@@ -24,36 +24,44 @@
  *
  */
 
-#ifndef KDTRACE_HOOKS
-#define KDTRACE_HOOKS
-#endif
+#include "opt_kdtrace.h"
 
 #include <sys/cdefs.h>
 #include <sys/param.h>
 #include <sys/systm.h>
+
 #include <sys/conf.h>
+#include <sys/eventhandler.h>
 #include <sys/kernel.h>
 #include <sys/limits.h>
-#include <sys/lock.h>
 #include <sys/linker.h>
+#include <sys/linker_set.h>
+#include <sys/lock.h>
+#include <sys/malloc.h>
 #include <sys/module.h>
 #include <sys/mutex.h>
-
-#include <sys/dtrace.h>
+#include <sys/queue.h>
 #include <sys/sdt.h>
 
-#define SDT_ADDR2NDX(addr) (((uintptr_t)(addr)) >> 4)
+#include <sys/dtrace.h>
+#include <sys/dtrace_bsd.h>
 
-static d_open_t sdt_open;
-static int	sdt_unload(void);
+/* DTrace methods. */
 static void	sdt_getargdesc(void *, dtrace_id_t, void *, dtrace_argdesc_t *);
 static void	sdt_provide_probes(void *, dtrace_probedesc_t *);
 static void	sdt_destroy(void *, dtrace_id_t, void *);
 static void	sdt_enable(void *, dtrace_id_t, void *);
 static void	sdt_disable(void *, dtrace_id_t, void *);
+
+static d_open_t	sdt_open;
 static void	sdt_load(void *);
-static int	sdt_provider_unreg_callback(struct sdt_provider *prov, 
-		    void *arg);
+static int	sdt_unload(void *);
+static void	sdt_create_provider(struct sdt_provider *);
+static void	sdt_create_probe(struct sdt_probe *);
+static void	sdt_modload(void *, struct linker_file *);
+static void	sdt_modunload(void *, struct linker_file *, int *);
+
+static MALLOC_DEFINE(M_SDT, "SDT", "DTrace SDT providers");
 
 static struct cdevsw sdt_cdevsw = {
 	.d_version	= D_VERSION,
@@ -79,141 +87,261 @@ static dtrace_pops_t sdt_pops = {
 	sdt_getargdesc,
 	NULL,
 	NULL,
-	sdt_destroy
+	sdt_destroy,
 };
 
-static struct cdev		*sdt_cdev;
+static struct cdev	*sdt_cdev;
 
-static int
-sdt_argtype_callback(struct sdt_argtype *argtype, void *arg)
-{
-	dtrace_argdesc_t *desc = arg;
+static TAILQ_HEAD(, sdt_provider) sdt_prov_list;
 
-	if (desc->dtargd_ndx == argtype->ndx) {
-		desc->dtargd_mapping = desc->dtargd_ndx; /* XXX */
-		strlcpy(desc->dtargd_native, argtype->type,
-		    sizeof(desc->dtargd_native));
-		desc->dtargd_xlate[0] = '\0'; /* XXX */
-	}
-
-	return (0);
-}
+eventhandler_tag	modload_tag;
+eventhandler_tag	modunload_tag;
 
 static void
-sdt_getargdesc(void *arg, dtrace_id_t id, void *parg, dtrace_argdesc_t *desc)
+sdt_create_provider(struct sdt_provider *prov)
 {
-	struct sdt_probe *probe = parg;
+	struct sdt_provider *curr, *newprov;
 
-	if (desc->dtargd_ndx < probe->n_args)
-		(void) (sdt_argtype_listall(probe, sdt_argtype_callback, desc));
-	else
-		desc->dtargd_ndx = DTRACE_ARGNONE;
+	TAILQ_FOREACH(curr, &sdt_prov_list, prov_entry)
+		if (strcmp(prov->name, curr->name) == 0) {
+			/* The provider has already been defined. */
+			curr->sdt_refs++;
+			return;
+		}
+
+	/*
+	 * Make a copy of prov so that we don't lose fields if its module is
+	 * unloaded but the provider isn't destroyed. This could happen with
+	 * a provider that spans multiple modules.
+	 */
+	newprov = malloc(sizeof(*newprov), M_SDT, M_WAITOK | M_ZERO);
+	newprov->name = strdup(prov->name, M_SDT);
+	prov->sdt_refs = newprov->sdt_refs = 1;
+	TAILQ_INIT(&newprov->probe_list);
+
+	TAILQ_INSERT_TAIL(&sdt_prov_list, newprov, prov_entry);
 
-	return;
+	(void)dtrace_register(newprov->name, &sdt_attr, DTRACE_PRIV_USER, NULL,
+	    &sdt_pops, NULL, (dtrace_provider_id_t *)&newprov->id);
+	prov->id = newprov->id;
 }
 
-static int
-sdt_probe_callback(struct sdt_probe *probe, void *arg __unused)
+static void
+sdt_create_probe(struct sdt_probe *probe)
 {
-	struct sdt_provider *prov = probe->prov;
-	char mod[64];
-	char func[64];
-	char name[64];
+	struct sdt_provider *prov;
+	char mod[DTRACE_MODNAMELEN];
+	char func[DTRACE_FUNCNAMELEN];
+	char name[DTRACE_NAMELEN];
+	size_t len;
+
+	TAILQ_FOREACH(prov, &sdt_prov_list, prov_entry)
+		if (strcmp(prov->name, probe->prov->name) == 0)
+			break;
+
+	KASSERT(prov != NULL, ("probe defined without a provider"));
+
+	/* If no module name was specified, use the module filename. */
+	if (*probe->mod == 0) {
+		len = strlcpy(mod, probe->sdtp_lf->filename, sizeof(mod));
+		if (len > 3 && strcmp(mod + len - 3, ".ko") == 0)
+			mod[len - 3] = '\0';
+	} else
+		strlcpy(mod, probe->mod, sizeof(mod));
 
 	/*
 	 * Unfortunately this is necessary because the Solaris DTrace
 	 * code mixes consts and non-consts with casts to override
 	 * the incompatibilies. On FreeBSD, we use strict warnings
-	 * in gcc, so we have to respect const vs non-const.
+	 * in the C compiler, so we have to respect const vs non-const.
 	 */
-	strlcpy(mod, probe->mod, sizeof(mod));
 	strlcpy(func, probe->func, sizeof(func));
 	strlcpy(name, probe->name, sizeof(name));
 
-	if (dtrace_probe_lookup(prov->id, mod, func, name) != 0)
-		return (0);
+	if (dtrace_probe_lookup(prov->id, mod, func, name) != DTRACE_IDNONE)
+		return;
 
-	(void) dtrace_probe_create(prov->id, probe->mod, probe->func,
-	    probe->name, 1, probe);
+	TAILQ_INSERT_TAIL(&prov->probe_list, probe, probe_entry);
 
-	return (0);
+	(void)dtrace_probe_create(prov->id, mod, func, name, 1, probe);
 }
 
-static int
-sdt_provider_entry(struct sdt_provider *prov, void *arg)
+/* Probes are created through the SDT module load/unload hook. */
+static void
+sdt_provide_probes(void *arg, dtrace_probedesc_t *desc)
 {
-	return (sdt_probe_listall(prov, sdt_probe_callback, NULL));
 }
 
 static void
-sdt_provide_probes(void *arg, dtrace_probedesc_t *desc)
+sdt_enable(void *arg __unused, dtrace_id_t id, void *parg)
 {
-	if (desc != NULL)
-		return;
+	struct sdt_probe *probe = parg;
 
-	(void) sdt_provider_listall(sdt_provider_entry, NULL);
+	probe->id = id;
+	probe->sdtp_lf->nenabled++;
 }
 
 static void
-sdt_destroy(void *arg, dtrace_id_t id, void *parg)
+sdt_disable(void *arg __unused, dtrace_id_t id, void *parg)
 {
-	/* Nothing to do here. */
+	struct sdt_probe *probe = parg;
+
+	KASSERT(probe->sdtp_lf->nenabled > 0, ("no probes enabled"));
+
+	probe->id = 0;
+	probe->sdtp_lf->nenabled--;
 }
 
 static void
-sdt_enable(void *arg, dtrace_id_t id, void *parg)
+sdt_getargdesc(void *arg, dtrace_id_t id, void *parg, dtrace_argdesc_t *desc)
 {
+	struct sdt_argtype *argtype;
 	struct sdt_probe *probe = parg;
 
-	probe->id = id;
+	if (desc->dtargd_ndx < probe->n_args) {
+		TAILQ_FOREACH(argtype, &probe->argtype_list, argtype_entry) {
+			if (desc->dtargd_ndx == argtype->ndx) {
+				/* XXX */
+				desc->dtargd_mapping = desc->dtargd_ndx;
+				strlcpy(desc->dtargd_native, argtype->type,
+				    sizeof(desc->dtargd_native));
+				desc->dtargd_xlate[0] = '\0'; /* XXX */
+			}
+		}
+	} else
+		desc->dtargd_ndx = DTRACE_ARGNONE;
 }
 
 static void
-sdt_disable(void *arg, dtrace_id_t id, void *parg)
+sdt_destroy(void *arg, dtrace_id_t id, void *parg)
 {
-	struct sdt_probe *probe = parg;
+	struct sdt_probe *probe;
 
-	probe->id = 0;
+	probe = parg;
+	TAILQ_REMOVE(&probe->prov->probe_list, probe, probe_entry);
+}
+
+/*
+ * Called from the kernel linker when a module is loaded, before
+ * dtrace_module_loaded() is called. This is done so that it's possible to
+ * register new providers when modules are loaded. We cannot do this in the
+ * provide_module method since it's called with the provider lock held
+ * and dtrace_register() will try to acquire it again.
+ */
+static void
+sdt_modload(void *arg __unused, struct linker_file *lf)
+{
+	struct sdt_provider **prov, **begin, **end;
+	struct sdt_probe **probe, **p_begin, **p_end;
+	struct sdt_argtype **argtype, **a_begin, **a_end;
+
+	if (linker_file_lookup_set(lf, "sdt_providers_set", &begin, &end, NULL))
+		return;
+	for (prov = begin; prov < end; prov++)
+		sdt_create_provider(*prov);
+
+	if (linker_file_lookup_set(lf, "sdt_probes_set", &p_begin, &p_end,
+	    NULL))
+		return;
+	for (probe = p_begin; probe < p_end; probe++) {
+		(*probe)->sdtp_lf = lf;
+		sdt_create_probe(*probe);
+		TAILQ_INIT(&(*probe)->argtype_list);
+	}
+
+	if (linker_file_lookup_set(lf, "sdt_argtypes_set", &a_begin, &a_end,
+	    NULL))
+		return;
+	for (argtype = a_begin; argtype < a_end; argtype++) {
+		(*argtype)->probe->n_args++;
+		TAILQ_INSERT_TAIL(&(*argtype)->probe->argtype_list, *argtype,
+		    argtype_entry);
+	}
+}
+
+static void
+sdt_modunload(void *arg __unused, struct linker_file *lf, int *error __unused)
+{
+	struct sdt_provider *prov, **curr, **begin, **end, *tmp;
+
+	if (*error != 0)
+		/* We already have an error, so don't do anything. */
+		return;
+	else if (linker_file_lookup_set(lf, "sdt_providers_set", &begin, &end, NULL))
+		/* No DTrace providers are declared in this module. */
+		return;
+
+	/*
+	 * Go through all the providers declared in this module and unregister
+	 * any that aren't declared in another loaded module.
+	 */
+	for (curr = begin; curr < end; curr++) {
+		TAILQ_FOREACH_SAFE(prov, &sdt_prov_list, prov_entry, tmp) {
+			if (strcmp(prov->name, (*curr)->name) == 0) {
+				if (prov->sdt_refs == 1) {
+					TAILQ_REMOVE(&sdt_prov_list, prov,
+					    prov_entry);
+					dtrace_unregister(prov->id);
+					free(prov->name, M_SDT);
+					free(prov, M_SDT);
+				} else
+					prov->sdt_refs--;
+				break;
+			}
+		}
+	}
 }
 
 static int
-sdt_provider_reg_callback(struct sdt_provider *prov, void *arg __unused)
+sdt_linker_file_cb(linker_file_t lf, void *arg __unused)
 {
-	return (dtrace_register(prov->name, &sdt_attr, DTRACE_PRIV_USER,
-	    NULL, &sdt_pops, NULL, (dtrace_provider_id_t *) &prov->id));
+
+	sdt_modload(NULL, lf);
+
+	return (0);
 }
 
 static void
-sdt_load(void *dummy)
+sdt_load(void *arg __unused)
 {
+
+	TAILQ_INIT(&sdt_prov_list);
+
 	/* Create the /dev/dtrace/sdt entry. */
 	sdt_cdev = make_dev(&sdt_cdevsw, 0, UID_ROOT, GID_WHEEL, 0600,
 	    "dtrace/sdt");
 
 	sdt_probe_func = dtrace_probe;
 
-	sdt_register_callbacks(sdt_provider_reg_callback, NULL,
-	    sdt_provider_unreg_callback, NULL, sdt_probe_callback, NULL);
-}
+	modload_tag = EVENTHANDLER_REGISTER(mod_load, sdt_modload, NULL,
+	    EVENTHANDLER_PRI_ANY);
+	modunload_tag = EVENTHANDLER_REGISTER(mod_unload, sdt_modunload, NULL,
+	    EVENTHANDLER_PRI_ANY);
 
-static int
-sdt_provider_unreg_callback(struct sdt_provider *prov, void *arg __unused)
-{
-	return (dtrace_unregister(prov->id));
+	/* Pick up probes from the kernel and already-loaded modules. */
+	linker_file_foreach(sdt_linker_file_cb, NULL);
 }
 
 static int
-sdt_unload()
+sdt_unload(void *arg __unused)
 {
-	int error = 0;
+	struct sdt_provider *prov, *tmp;
+
+	EVENTHANDLER_DEREGISTER(mod_load, modload_tag);
+	EVENTHANDLER_DEREGISTER(mod_unload, modunload_tag);
 
 	sdt_probe_func = sdt_probe_stub;
 
-	sdt_deregister_callbacks();
-	
+	TAILQ_FOREACH_SAFE(prov, &sdt_prov_list, prov_entry, tmp) {
+		TAILQ_REMOVE(&sdt_prov_list, prov, prov_entry);
+		dtrace_unregister(prov->id);
+		free(prov->name, M_SDT);
+		free(prov, M_SDT);
+	}
+
 	destroy_dev(sdt_cdev);
 
-	return (error);
+	return (0);
 }
 
 /* ARGSUSED */
@@ -235,7 +363,6 @@ sdt_modevent(module_t mod __unused, int type, void *data __unused)
 	default:
 		error = EOPNOTSUPP;
 		break;
-
 	}
 
 	return (error);
@@ -243,8 +370,10 @@ sdt_modevent(module_t mod __unused, int type, void *data __unused)
 
 /* ARGSUSED */
 static int
-sdt_open(struct cdev *dev __unused, int oflags __unused, int devtype __unused, struct thread *td __unused)
+sdt_open(struct cdev *dev __unused, int oflags __unused, int devtype __unused,
+    struct thread *td __unused)
 {
+
 	return (0);
 }