Allow setting a dump hook

If users want to be notified when a heap dump occurs, they can set this hook.

include/jemalloc/internal/prof_externs.h

@@ -48,14 +48,12 @@ extern size_t lg_prof_sample;
 
 extern bool prof_booted;
 
-/*
- * 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.
- */
 void prof_backtrace_hook_set(prof_backtrace_hook_t hook);
 prof_backtrace_hook_t prof_backtrace_hook_get();
 
+void prof_dump_hook_set(prof_dump_hook_t hook);
+prof_dump_hook_t prof_dump_hook_get();
+
 /* Functions only accessed in prof_inlines.h */
 prof_tdata_t *prof_tdata_init(tsd_t *tsd);
 prof_tdata_t *prof_tdata_reinit(tsd_t *tsd, prof_tdata_t *tdata);

include/jemalloc/internal/prof_hook.h

@@ -13,4 +13,9 @@
  */
 typedef void (*prof_backtrace_hook_t)(void **, unsigned *, unsigned);
 
+/*
+ * A callback hook that notifies about recently dumped heap profile.
+ */
+typedef void (*prof_dump_hook_t)(const char *filename);
+
 #endif /* JEMALLOC_INTERNAL_PROF_HOOK_H */

src/ctl.c

@@ -306,6 +306,7 @@ CTL_PROTO(stats_zero_reallocs)
 CTL_PROTO(experimental_hooks_install)
 CTL_PROTO(experimental_hooks_remove)
 CTL_PROTO(experimental_hooks_prof_backtrace)
+CTL_PROTO(experimental_hooks_prof_dump)
 CTL_PROTO(experimental_thread_activity_callback)
 CTL_PROTO(experimental_utilization_query)
 CTL_PROTO(experimental_utilization_batch_query)
@@ -835,7 +836,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("prof_backtrace"),	CTL(experimental_hooks_prof_backtrace)}
+	{NAME("prof_backtrace"),	CTL(experimental_hooks_prof_backtrace)},
+	{NAME("prof_dump"),	CTL(experimental_hooks_prof_dump)},
 };
 
 static const ctl_named_node_t experimental_thread_node[] = {
@@ -3362,6 +3364,34 @@ label_return:
 	return ret;
 }
 
+static int
+experimental_hooks_prof_dump_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_dump_hook_t old_hook =
+		    prof_dump_hook_get();
+		READ(old_hook, prof_dump_hook_t);
+	}
+	if (newp != NULL) {
+		if (!opt_prof) {
+			ret = ENOENT;
+			goto label_return;
+		}
+		prof_dump_hook_t new_hook JEMALLOC_CC_SILENCE_INIT(NULL);
+		WRITE(new_hook, prof_dump_hook_t);
+		prof_dump_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

@@ -73,6 +73,9 @@ bool prof_booted = false;
 /* Logically a prof_backtrace_hook_t. */
 atomic_p_t prof_backtrace_hook;
 
+/* Logically a prof_dump_hook_t. */
+atomic_p_t prof_dump_hook;
+
 /******************************************************************************/
 
 void
@@ -533,6 +536,17 @@ prof_backtrace_hook_get() {
 	    ATOMIC_ACQUIRE);
 }
 
+void
+prof_dump_hook_set(prof_dump_hook_t hook) {
+	atomic_store_p(&prof_dump_hook, hook, ATOMIC_RELEASE);
+}
+
+prof_dump_hook_t
+prof_dump_hook_get() {
+	return (prof_dump_hook_t)atomic_load_p(&prof_dump_hook,
+	    ATOMIC_ACQUIRE);
+}
+
 void
 prof_boot0(void) {
 	cassert(config_prof);
@@ -672,8 +686,8 @@ prof_boot2(tsd_t *tsd, base_t *base) {
 			}
 		}
 
-		prof_hooks_init();
 		prof_unwind_init();
+		prof_hooks_init();
 	}
 	prof_booted = true;
 

src/prof_sys.c

@@ -55,6 +55,7 @@ prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 	cassert(config_prof);
 	assert(*len == 0);
 	assert(vec != NULL);
+	assert(max_len == PROF_BT_MAX);
 
 	nframes = unw_backtrace(vec, PROF_BT_MAX);
 	if (nframes <= 0) {
@@ -95,6 +96,8 @@ prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 	prof_unwind_data_t data = {vec, len, max_len};
 
 	cassert(config_prof);
+	assert(vec != NULL);
+	assert(max_len == PROF_BT_MAX);
 
 	_Unwind_Backtrace(prof_unwind_callback, &data);
 }
@@ -118,6 +121,8 @@ prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 	}
 
 	cassert(config_prof);
+	assert(vec != NULL);
+	assert(max_len == PROF_BT_MAX);
 
 	BT_FRAME(0)
 	BT_FRAME(1)
@@ -272,8 +277,10 @@ prof_backtrace_impl(void **vec, unsigned *len, unsigned max_len) {
 void
 prof_backtrace(tsd_t *tsd, prof_bt_t *bt) {
 	cassert(config_prof);
-	pre_reentrancy(tsd, NULL);
 	prof_backtrace_hook_t prof_backtrace_hook = prof_backtrace_hook_get();
+	assert(prof_backtrace_hook != NULL);
+
+	pre_reentrancy(tsd, NULL);
 	prof_backtrace_hook(bt->vec, &bt->len, PROF_BT_MAX);
 	post_reentrancy(tsd);
 }
@@ -281,6 +288,7 @@ prof_backtrace(tsd_t *tsd, prof_bt_t *bt) {
 void
 prof_hooks_init() {
 	prof_backtrace_hook_set(&prof_backtrace_impl);
+	prof_dump_hook_set(NULL);
 }
 
 void
@@ -506,6 +514,10 @@ prof_dump(tsd_t *tsd, bool propagate_err, const char *filename,
 	buf_writer_terminate(tsd_tsdn(tsd), &buf_writer);
 	prof_dump_close(&arg);
 
+	prof_dump_hook_t dump_hook = prof_dump_hook_get();
+	if (dump_hook != NULL) {
+		dump_hook(filename);
+	}
 	malloc_mutex_unlock(tsd_tsdn(tsd), &prof_dump_mtx);
 	post_reentrancy(tsd);
 

test/unit/prof_hook.c

@@ -1,6 +1,11 @@
 #include "test/jemalloc_test.h"
 
+const char *dump_filename = "/dev/null";
+
+prof_backtrace_hook_t default_hook;
+
 bool mock_bt_hook_called = false;
+bool mock_dump_hook_called = false;
 
 void
 mock_bt_hook(void **vec, unsigned *len, unsigned max_len) {
@@ -11,7 +16,38 @@ mock_bt_hook(void **vec, unsigned *len, unsigned max_len) {
 	mock_bt_hook_called = true;
 }
 
-TEST_BEGIN(test_prof_backtrace_hook) {
+void
+mock_bt_augmenting_hook(void **vec, unsigned *len, unsigned max_len) {
+	default_hook(vec, len, max_len);
+	expect_u_gt(*len, 0, "Default backtrace hook returned empty backtrace");
+	expect_u_lt(*len, max_len,
+	    "Default backtrace hook returned too large backtrace");
+
+	/* Add a separator between default frames and augmented */
+	vec[*len] = (void *)0x030303030;
+	(*len)++;
+
+	/* Add more stack frames */
+	for (unsigned i = 0; i < 3; ++i) {
+		if (*len == max_len) {
+			break;
+		}
+		vec[*len] = (void *)((uintptr_t)i);
+		(*len)++;
+	}
+
+
+	mock_bt_hook_called = true;
+}
+
+void
+mock_dump_hook(const char *filename) {
+	mock_dump_hook_called = true;
+	expect_str_eq(filename, dump_filename,
+	    "Incorrect file name passed to the dump hook");
+}
+
+TEST_BEGIN(test_prof_backtrace_hook_replace) {
 
 	test_skip_if(!config_prof);
 
@@ -27,7 +63,6 @@ TEST_BEGIN(test_prof_backtrace_hook) {
 	    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",
@@ -54,8 +89,81 @@ TEST_BEGIN(test_prof_backtrace_hook) {
 }
 TEST_END
 
+TEST_BEGIN(test_prof_backtrace_hook_augment) {
+
+	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");
+
+	size_t default_hook_sz = sizeof(prof_backtrace_hook_t);
+	prof_backtrace_hook_t hook = &mock_bt_augmenting_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
+
+TEST_BEGIN(test_prof_dump_hook) {
+
+	test_skip_if(!config_prof);
+
+	mock_dump_hook_called = false;
+
+	expect_d_eq(mallctl("prof.dump", NULL, NULL, (void *)&dump_filename,
+	    sizeof(dump_filename)), 0, "Failed to dump heap profile");
+
+	expect_false(mock_dump_hook_called, "Called dump hook before it's set");
+
+	size_t default_hook_sz = sizeof(prof_dump_hook_t);
+	prof_dump_hook_t hook = &mock_dump_hook;
+	expect_d_eq(mallctl("experimental.hooks.prof_dump",
+	    (void *)&default_hook, &default_hook_sz, (void *)&hook,
+	    sizeof(hook)), 0, "Unexpected mallctl failure setting hook");
+
+	expect_d_eq(mallctl("prof.dump", NULL, NULL, (void *)&dump_filename,
+	    sizeof(dump_filename)), 0, "Failed to dump heap profile");
+
+	expect_true(mock_dump_hook_called, "Didn't call mock hook");
+
+	prof_dump_hook_t current_hook;
+	size_t current_hook_sz = sizeof(prof_dump_hook_t);
+	expect_d_eq(mallctl("experimental.hooks.prof_dump",
+	    (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");
+}
+TEST_END
+
 int
 main(void) {
 	return test(
-	    test_prof_backtrace_hook);
+	    test_prof_backtrace_hook_replace,
+	    test_prof_backtrace_hook_augment,
+	    test_prof_dump_hook);
 }