From 91dd776cd20dbe6b0e6c367896afa1623dd9fb56 Mon Sep 17 00:00:00 2001
From: John Birrell <jb@FreeBSD.org>
Date: Fri, 23 May 2008 00:43:36 +0000
Subject: [PATCH] Add support for the DTrace malloc provider which can enable
 probes on a per-malloc type basis.

---
 sys/kern/kern_malloc.c | 63 ++++++++++++++++++++++++++++++++++++++++++
 sys/sys/malloc.h       | 15 ++++++++++
 2 files changed, 78 insertions(+)

diff --git a/sys/kern/kern_malloc.c b/sys/kern/kern_malloc.c
index 8f2b5099cb23..3b169d53b21c 100644
--- a/sys/kern/kern_malloc.c
+++ b/sys/kern/kern_malloc.c
@@ -46,6 +46,7 @@
 __FBSDID("$FreeBSD$");
 
 #include "opt_ddb.h"
+#include "opt_kdtrace.h"
 #include "opt_vm.h"
 
 #include <sys/param.h>
@@ -86,6 +87,12 @@ __FBSDID("$FreeBSD$");
 
 #include <ddb/ddb.h>
 
+#ifdef KDTRACE_HOOKS
+#include <sys/dtrace_bsd.h>
+
+dtrace_malloc_probe_func_t	dtrace_malloc_probe;
+#endif
+
 /*
  * When realloc() is called, if the new size is sufficiently smaller than
  * the old size, realloc() will allocate a new, smaller block to avoid
@@ -255,6 +262,17 @@ malloc_type_zone_allocated(struct malloc_type *mtp, unsigned long size,
 	}
 	if (zindx != -1)
 		mtsp->mts_size |= 1 << zindx;
+
+#ifdef KDTRACE_HOOKS
+	if (dtrace_malloc_probe != NULL) {
+		uint32_t probe_id = mtip->mti_probes[DTMALLOC_PROBE_MALLOC];
+		if (probe_id != 0)
+			(dtrace_malloc_probe)(probe_id,
+			    (uintptr_t) mtp, (uintptr_t) mtip,
+			    (uintptr_t) mtsp, size, zindx);
+	}
+#endif
+
 	critical_exit();
 }
 
@@ -283,6 +301,17 @@ malloc_type_freed(struct malloc_type *mtp, unsigned long size)
 	mtsp = &mtip->mti_stats[curcpu];
 	mtsp->mts_memfreed += size;
 	mtsp->mts_numfrees++;
+
+#ifdef KDTRACE_HOOKS
+	if (dtrace_malloc_probe != NULL) {
+		uint32_t probe_id = mtip->mti_probes[DTMALLOC_PROBE_FREE];
+		if (probe_id != 0)
+			(dtrace_malloc_probe)(probe_id,
+			    (uintptr_t) mtp, (uintptr_t) mtip,
+			    (uintptr_t) mtsp, size, 0);
+	}
+#endif
+
 	critical_exit();
 }
 
@@ -804,6 +833,40 @@ SYSCTL_PROC(_kern, OID_AUTO, malloc_stats, CTLFLAG_RD|CTLTYPE_STRUCT,
 SYSCTL_INT(_kern, OID_AUTO, malloc_count, CTLFLAG_RD, &kmemcount, 0,
     "Count of kernel malloc types");
 
+void
+malloc_type_list(malloc_type_list_func_t *func, void *arg)
+{
+	struct malloc_type *mtp, **bufmtp;
+	int count, i;
+	size_t buflen;
+
+	mtx_lock(&malloc_mtx);
+restart:
+	mtx_assert(&malloc_mtx, MA_OWNED);
+	count = kmemcount;
+	mtx_unlock(&malloc_mtx);
+
+	buflen = sizeof(struct malloc_type *) * count;
+	bufmtp = malloc(buflen, M_TEMP, M_WAITOK);
+
+	mtx_lock(&malloc_mtx);
+
+	if (count < kmemcount) {
+		free(bufmtp, M_TEMP);
+		goto restart;
+	}
+
+	for (mtp = kmemstatistics, i = 0; mtp != NULL; mtp = mtp->ks_next, i++)
+		bufmtp[i] = mtp;
+
+	mtx_unlock(&malloc_mtx);
+
+	for (i = 0; i < count; i++)
+		(func)(bufmtp[i], arg);
+
+	free(bufmtp, M_TEMP);
+}
+
 #ifdef DDB
 DB_SHOW_COMMAND(malloc, db_show_malloc)
 {
diff --git a/sys/sys/malloc.h b/sys/sys/malloc.h
index 8ea18875972a..33b287f83d92 100644
--- a/sys/sys/malloc.h
+++ b/sys/sys/malloc.h
@@ -80,7 +80,16 @@ struct malloc_type_stats {
 	uint64_t	_mts_reserved3;	/* Reserved field. */
 };
 
+/*
+ * Index definitions for the mti_probes[] array.
+ */
+#define DTMALLOC_PROBE_MALLOC		0
+#define DTMALLOC_PROBE_FREE		1
+#define DTMALLOC_PROBE_MAX		2
+
 struct malloc_type_internal {
+	uint32_t	mti_probes[DTMALLOC_PROBE_MAX];
+					/* DTrace probe ID array. */
 	struct malloc_type_stats	mti_stats[MAXCPU];
 };
 
@@ -173,6 +182,11 @@ MALLOC_DECLARE(M_IOV);
 
 extern struct mtx malloc_mtx;
 
+/*
+ * Function type used when iterating over the list of malloc types.
+ */
+typedef void malloc_type_list_func_t(struct malloc_type *, void *);
+
 void	contigfree(void *addr, unsigned long size, struct malloc_type *type);
 void	*contigmalloc(unsigned long size, struct malloc_type *type, int flags,
 	    vm_paddr_t low, vm_paddr_t high, unsigned long alignment,
@@ -183,6 +197,7 @@ void	malloc_init(void *);
 int	malloc_last_fail(void);
 void	malloc_type_allocated(struct malloc_type *type, unsigned long size);
 void	malloc_type_freed(struct malloc_type *type, unsigned long size);
+void	malloc_type_list(malloc_type_list_func_t *, void *);
 void	malloc_uninit(void *);
 void	*realloc(void *addr, unsigned long size, struct malloc_type *type,
 	    int flags);