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.

3 files changed, 74 insertions, 7 deletions

diff --git a/test/analyze/prof_bias.c b/test/analyze/prof_bias.c
index 0aae766b..4b960a66 100644
--- a/test/analyze/prof_bias.c
+++ b/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);
diff --git a/test/unit/prof_hook.c b/test/unit/prof_hook.c
new file mode 100644
index 00000000..32d0e9ea
--- /dev/null
+++ b/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);
+}
diff --git a/test/unit/prof_hook.sh b/test/unit/prof_hook.sh
new file mode 100644
index 00000000..d14cb8c5
--- /dev/null
+++ b/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
+