Allow setting custom backtrace hook

Existing backtrace implementations skip native stack frames from runtimes like
Python. The hook allows to augment the backtraces to attribute allocations to
native functions in heap profiles.

Makefile.in

@@ -247,6 +247,7 @@ TESTS_UNIT := \
 	$(srcroot)test/unit/prof_accum.c \
 	$(srcroot)test/unit/prof_active.c \
 	$(srcroot)test/unit/prof_gdump.c \
+	$(srcroot)test/unit/prof_hook.c \
 	$(srcroot)test/unit/prof_idump.c \
 	$(srcroot)test/unit/prof_log.c \
 	$(srcroot)test/unit/prof_mdump.c \

include/jemalloc/internal/prof_externs.h

@@ -2,6 +2,7 @@
 #define JEMALLOC_INTERNAL_PROF_EXTERNS_H
 
 #include "jemalloc/internal/mutex.h"
+#include "jemalloc/internal/prof_hook.h"
 
 extern bool opt_prof;
 extern bool opt_prof_active;
@@ -52,7 +53,8 @@ extern bool prof_booted;
  * otherwise difficult to guarantee that two allocations are reported as coming
  * from the exact same stack trace in the presence of an optimizing compiler.
  */
-extern void (* JET_MUTABLE prof_backtrace_hook)(prof_bt_t *bt);
+void prof_backtrace_hook_set(prof_backtrace_hook_t hook);
+prof_backtrace_hook_t prof_backtrace_hook_get();
 
 /* Functions only accessed in prof_inlines.h */
 prof_tdata_t *prof_tdata_init(tsd_t *tsd);

include/jemalloc/internal/prof_hook.h

@@ -0,0 +1,16 @@
+#ifndef JEMALLOC_INTERNAL_PROF_HOOK_H
+#define JEMALLOC_INTERNAL_PROF_HOOK_H
+
+/*
+ * The hooks types of which are declared in this file are experimental and
+ * undocumented, thus the typedefs are located in an 'internal' header.
+ */
+
+/*
+ * A hook to mock out backtrace functionality.  This can be handy, since it's
+ * otherwise difficult to guarantee that two allocations are reported as coming
+ * from the exact same stack trace in the presence of an optimizing compiler.
+ */
+typedef void (*prof_backtrace_hook_t)(void **, unsigned *, unsigned);
+
+#endif /* JEMALLOC_INTERNAL_PROF_HOOK_H */

include/jemalloc/internal/prof_structs.h

@@ -16,7 +16,8 @@ struct prof_bt_s {
 #ifdef JEMALLOC_PROF_LIBGCC
 /* Data structure passed to libgcc _Unwind_Backtrace() callback functions. */
 typedef struct {
-	prof_bt_t	*bt;
+	void 		**vec;
+	unsigned	*len;
 	unsigned	max;
 } prof_unwind_data_t;
 #endif

include/jemalloc/internal/prof_sys.h

@@ -6,6 +6,7 @@ extern base_t *prof_base;
 
 void bt_init(prof_bt_t *bt, void **vec);
 void prof_backtrace(tsd_t *tsd, prof_bt_t *bt);
+void prof_hooks_init();
 void prof_unwind_init();
 void prof_sys_thread_name_fetch(tsd_t *tsd);
 int prof_getpid(void);

src/ctl.c

@@ -305,6 +305,7 @@ CTL_PROTO(stats_retained)
 CTL_PROTO(stats_zero_reallocs)
 CTL_PROTO(experimental_hooks_install)
 CTL_PROTO(experimental_hooks_remove)
+CTL_PROTO(experimental_hooks_prof_backtrace)
 CTL_PROTO(experimental_thread_activity_callback)
 CTL_PROTO(experimental_utilization_query)
 CTL_PROTO(experimental_utilization_batch_query)
@@ -833,7 +834,8 @@ static const ctl_named_node_t stats_node[] = {
 
 static const ctl_named_node_t experimental_hooks_node[] = {
 	{NAME("install"),	CTL(experimental_hooks_install)},
-	{NAME("remove"),	CTL(experimental_hooks_remove)}
+	{NAME("remove"),	CTL(experimental_hooks_remove)},
+	{NAME("prof_backtrace"),	CTL(experimental_hooks_prof_backtrace)}
 };
 
 static const ctl_named_node_t experimental_thread_node[] = {
@@ -3328,6 +3330,38 @@ prof_log_stop_ctl(tsd_t *tsd, const size_t *mib, size_t miblen, void *oldp,
 	return 0;
 }
 
+static int
+experimental_hooks_prof_backtrace_ctl(tsd_t *tsd, const size_t *mib,
+    size_t miblen, void *oldp, size_t *oldlenp, void *newp, size_t newlen) {
+	int ret;
+
+	if (oldp == NULL && newp == NULL) {
+		ret = EINVAL;
+		goto label_return;
+	}
+	if (oldp != NULL) {
+		prof_backtrace_hook_t old_hook =
+		    prof_backtrace_hook_get();
+		READ(old_hook, prof_backtrace_hook_t);
+	}
+	if (newp != NULL) {
+		if (!opt_prof) {
+			ret = ENOENT;
+			goto label_return;
+		}
+		prof_backtrace_hook_t new_hook JEMALLOC_CC_SILENCE_INIT(NULL);
+		WRITE(new_hook, prof_backtrace_hook_t);
+		if (new_hook == NULL) {
+			ret = EINVAL;
+			goto label_return;
+		}
+		prof_backtrace_hook_set(new_hook);
+	}
+	ret = 0;
+label_return:
+	return ret;
+}
+
 /******************************************************************************/
 
 CTL_RO_CGEN(config_stats, stats_allocated, ctl_stats->allocated, size_t)

src/prof.c

@@ -10,6 +10,7 @@
 #include "jemalloc/internal/prof_recent.h"
 #include "jemalloc/internal/prof_stats.h"
 #include "jemalloc/internal/prof_sys.h"
+#include "jemalloc/internal/prof_hook.h"
 #include "jemalloc/internal/thread_event.h"
 
 /*
@@ -69,6 +70,9 @@ static malloc_mutex_t next_thr_uid_mtx;
 /* Do not dump any profiles until bootstrapping is complete. */
 bool prof_booted = false;
 
+/* Logically a prof_backtrace_hook_t. */
+atomic_p_t prof_backtrace_hook;
+
 /******************************************************************************/
 
 void
@@ -518,6 +522,17 @@ prof_gdump_set(tsdn_t *tsdn, bool gdump) {
 	return prof_gdump_old;
 }
 
+void
+prof_backtrace_hook_set(prof_backtrace_hook_t hook) {
+	atomic_store_p(&prof_backtrace_hook, hook, ATOMIC_RELEASE);
+}
+
+prof_backtrace_hook_t
+prof_backtrace_hook_get() {
+	return (prof_backtrace_hook_t)atomic_load_p(&prof_backtrace_hook,
+	    ATOMIC_ACQUIRE);
+}
+
 void
 prof_boot0(void) {
 	cassert(config_prof);
@@ -657,6 +672,7 @@ prof_boot2(tsd_t *tsd, base_t *base) {
 			}
 		}
 
+		prof_hooks_init();
 		prof_unwind_init();
 	}
 	prof_booted = true;

src/prof_sys.c

@@ -49,18 +49,18 @@ bt_init(prof_bt_t *bt, void **vec) {
 
 #ifdef JEMALLOC_PROF_LIBUNWIND
 static void
-prof_backtrace_impl(prof_bt_t *bt) {
+prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 	int nframes;
 
 	cassert(config_prof);
-	assert(bt->len == 0);
-	assert(bt->vec != NULL);
+	assert(*len == 0);
+	assert(vec != NULL);
 
-	nframes = unw_backtrace(bt->vec, PROF_BT_MAX);
+	nframes = unw_backtrace(vec, PROF_BT_MAX);
 	if (nframes <= 0) {
 		return;
 	}
-	bt->len = nframes;
+	*len = nframes;
 }
 #elif (defined(JEMALLOC_PROF_LIBGCC))
 static _Unwind_Reason_Code
@@ -81,9 +81,9 @@ prof_unwind_callback(struct _Unwind_Context *context, void *arg) {
 	if (ip == NULL) {
 		return _URC_END_OF_STACK;
 	}
-	data->bt->vec[data->bt->len] = ip;
-	data->bt->len++;
-	if (data->bt->len == data->max) {
+	data->vec[*data->len] = ip;
+	(*data->len)++;
+	if (*data->len == data->max) {
 		return _URC_END_OF_STACK;
 	}
 
@@ -91,8 +91,8 @@ prof_unwind_callback(struct _Unwind_Context *context, void *arg) {
 }
 
 static void
-prof_backtrace_impl(prof_bt_t *bt) {
-	prof_unwind_data_t data = {bt, PROF_BT_MAX};
+prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
+	prof_unwind_data_t data = {vec, len, max_len};
 
 	cassert(config_prof);
 
@@ -100,9 +100,9 @@ prof_backtrace_impl(prof_bt_t *bt) {
 }
 #elif (defined(JEMALLOC_PROF_GCC))
 static void
-prof_backtrace_impl(prof_bt_t *bt) {
+prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 #define BT_FRAME(i)							\
-	if ((i) < PROF_BT_MAX) {					\
+	if ((i) < max_len) {						\
 		void *p;						\
 		if (__builtin_frame_address(i) == 0) {			\
 			return;						\
@@ -111,8 +111,8 @@ prof_backtrace_impl(prof_bt_t *bt) {
 		if (p == NULL) {					\
 			return;						\
 		}							\
-		bt->vec[(i)] = p;					\
-		bt->len = (i) + 1;					\
+		vec[(i)] = p;						\
+		*len = (i) + 1;						\
 	} else {							\
 		return;							\
 	}
@@ -263,24 +263,28 @@ prof_backtrace_impl(prof_bt_t *bt) {
 }
 #else
 static void
-prof_backtrace_impl(prof_bt_t *bt) {
+prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 	cassert(config_prof);
 	not_reached();
 }
 #endif
 
-
-void (* JET_MUTABLE prof_backtrace_hook)(prof_bt_t *bt) = &prof_backtrace_impl;
-
 void
 prof_backtrace(tsd_t *tsd, prof_bt_t *bt) {
 	cassert(config_prof);
 	pre_reentrancy(tsd, NULL);
-	prof_backtrace_hook(bt);
+	prof_backtrace_hook_t prof_backtrace_hook = prof_backtrace_hook_get();
+	prof_backtrace_hook(bt->vec, &bt->len, PROF_BT_MAX);
 	post_reentrancy(tsd);
 }
 
-void prof_unwind_init() {
+void
+prof_hooks_init() {
+	prof_backtrace_hook_set(&prof_backtrace_impl);
+}
+
+void
+prof_unwind_init() {
 #ifdef JEMALLOC_PROF_LIBGCC
 	/*
 	 * Cause the backtracing machinery to allocate its internal

test/analyze/prof_bias.c

@@ -24,12 +24,12 @@
  */
 
 static void
-mock_backtrace(prof_bt_t *bt) {
-	bt->len = 4;
-	bt->vec[0] = (void *)0x111;
-	bt->vec[1] = (void *)0x222;
-	bt->vec[2] = (void *)0x333;
-	bt->vec[3] = (void *)0x444;
+mock_backtrace(void **vec, unsigned *len, unsigned max_len) {
+	*len = 4;
+	vec[0] = (void *)0x111;
+	vec[1] = (void *)0x222;
+	vec[2] = (void *)0x333;
+	vec[3] = (void *)0x444;
 }
 
 static void
@@ -50,7 +50,7 @@ main(void) {
 	    sizeof(lg_prof_sample));
 	assert(err == 0);
 
-	prof_backtrace_hook = &mock_backtrace;
+	prof_backtrace_hook_set(mock_backtrace);
 	do_allocs(16, 32 * 1024 * 1024, /* do_frees */ true);
 	do_allocs(32 * 1024* 1024, 16, /* do_frees */ true);
 	do_allocs(16, 32 * 1024 * 1024, /* do_frees */ false);

test/unit/prof_hook.c

@@ -0,0 +1,61 @@
+#include "test/jemalloc_test.h"
+
+bool mock_bt_hook_called = false;
+
+void
+mock_bt_hook(void **vec, unsigned *len, unsigned max_len) {
+	*len = max_len;
+	for (unsigned i = 0; i < max_len; ++i) {
+		vec[i] = (void *)((uintptr_t)i);
+	}
+	mock_bt_hook_called = true;
+}
+
+TEST_BEGIN(test_prof_backtrace_hook) {
+
+	test_skip_if(!config_prof);
+
+	mock_bt_hook_called = false;
+
+	void *p0 = mallocx(1, 0);
+	assert_ptr_not_null(p0, "Failed to allocate");
+
+	expect_false(mock_bt_hook_called, "Called mock hook before it's set");
+
+	prof_backtrace_hook_t null_hook = NULL;
+	expect_d_eq(mallctl("experimental.hooks.prof_backtrace",
+	    NULL, 0, (void *)&null_hook,  sizeof(null_hook)),
+		EINVAL, "Incorrectly allowed NULL backtrace hook");
+
+	prof_backtrace_hook_t default_hook;
+	size_t default_hook_sz = sizeof(prof_backtrace_hook_t);
+	prof_backtrace_hook_t hook = &mock_bt_hook;
+	expect_d_eq(mallctl("experimental.hooks.prof_backtrace",
+	    (void *)&default_hook, &default_hook_sz, (void *)&hook,
+	    sizeof(hook)), 0, "Unexpected mallctl failure setting hook");
+
+	void *p1 = mallocx(1, 0);
+	assert_ptr_not_null(p1, "Failed to allocate");
+
+	expect_true(mock_bt_hook_called, "Didn't call mock hook");
+
+	prof_backtrace_hook_t current_hook;
+	size_t current_hook_sz = sizeof(prof_backtrace_hook_t);
+	expect_d_eq(mallctl("experimental.hooks.prof_backtrace",
+	    (void *)&current_hook, &current_hook_sz, (void *)&default_hook,
+	    sizeof(default_hook)), 0,
+	    "Unexpected mallctl failure resetting hook to default");
+
+	expect_ptr_eq(current_hook, hook,
+	    "Hook returned by mallctl is not equal to mock hook");
+
+	dallocx(p1, 0);
+	dallocx(p0, 0);
+}
+TEST_END
+
+int
+main(void) {
+	return test(
+	    test_prof_backtrace_hook);
+}

test/unit/prof_hook.sh

@@ -0,0 +1,6 @@
+#!/bin/sh
+
+if [ "x${enable_prof}" = "x1" ] ; then
+  export MALLOC_CONF="prof:true,lg_prof_sample:0"
+fi
+