Implement prof sample hooks "experimental.hooks.prof_sample(_free)".

The added hooks hooks.prof_sample and hooks.prof_sample_free are intended to
allow advanced users to track additional information, to enable new ways of
profiling on top of the jemalloc heap profile and sample features.

The sample hook is invoked after the allocation and backtracing, and forwards
the both the allocation and backtrace to the user hook; the sample_free hook
happens before the actual deallocation, and forwards only the ptr and usz to the
hook.

7 files changed, 307 insertions, 19 deletions

diff --git a/include/jemalloc/internal/prof_externs.h b/include/jemalloc/internal/prof_externs.h
index d1101561..412378a2 100644
--- a/include/jemalloc/internal/prof_externs.h
+++ b/include/jemalloc/internal/prof_externs.h
@@ -56,6 +56,12 @@ 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();
 
+void prof_sample_hook_set(prof_sample_hook_t hook);
+prof_sample_hook_t prof_sample_hook_get();
+
+void prof_sample_free_hook_set(prof_sample_free_hook_t hook);
+prof_sample_free_hook_t prof_sample_free_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);
@@ -63,7 +69,8 @@ prof_tdata_t *prof_tdata_reinit(tsd_t *tsd, prof_tdata_t *tdata);
 void prof_alloc_rollback(tsd_t *tsd, prof_tctx_t *tctx);
 void prof_malloc_sample_object(tsd_t *tsd, const void *ptr, size_t size,
     size_t usize, prof_tctx_t *tctx);
-void prof_free_sampled_object(tsd_t *tsd, size_t usize, prof_info_t *prof_info);
+void prof_free_sampled_object(tsd_t *tsd, const void *ptr, size_t usize,
+    prof_info_t *prof_info);
 prof_tctx_t *prof_tctx_create(tsd_t *tsd);
 void prof_idump(tsdn_t *tsdn);
 bool prof_mdump(tsd_t *tsd, const char *filename);
diff --git a/include/jemalloc/internal/prof_hook.h b/include/jemalloc/internal/prof_hook.h
index 150d19d3..8615dc53 100644
--- a/include/jemalloc/internal/prof_hook.h
+++ b/include/jemalloc/internal/prof_hook.h
@@ -18,4 +18,10 @@ typedef void (*prof_backtrace_hook_t)(void **, unsigned *, unsigned);
  */
 typedef void (*prof_dump_hook_t)(const char *filename);
 
+/* ptr, size, backtrace vector, backtrace vector length */
+typedef void (*prof_sample_hook_t)(const void *, size_t, void **, unsigned);
+
+/* ptr, size */
+typedef void (*prof_sample_free_hook_t)(const void *, size_t);
+
 #endif /* JEMALLOC_INTERNAL_PROF_HOOK_H */
diff --git a/include/jemalloc/internal/prof_inlines.h b/include/jemalloc/internal/prof_inlines.h
index 7d9608b5..ab3e01f6 100644
--- a/include/jemalloc/internal/prof_inlines.h
+++ b/include/jemalloc/internal/prof_inlines.h
@@ -213,7 +213,8 @@ prof_realloc(tsd_t *tsd, const void *ptr, size_t size, size_t usize,
 	 * counters.
 	 */
 	if (unlikely(old_sampled)) {
-		prof_free_sampled_object(tsd, old_usize, old_prof_info);
+		prof_free_sampled_object(tsd, old_ptr, old_usize,
+		    old_prof_info);
 	}
 }
 
@@ -250,7 +251,7 @@ prof_free(tsd_t *tsd, const void *ptr, size_t usize,
 
 	if (unlikely((uintptr_t)prof_info.alloc_tctx > (uintptr_t)1U)) {
 		assert(prof_sample_aligned(ptr));
-		prof_free_sampled_object(tsd, usize, &prof_info);
+		prof_free_sampled_object(tsd, ptr, usize, &prof_info);
 	}
 }
 
diff --git a/src/ctl.c b/src/ctl.c
index acf5d366..eafbdc61 100644
--- a/src/ctl.c
+++ b/src/ctl.c
@@ -315,6 +315,8 @@ 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_hooks_prof_sample)
+CTL_PROTO(experimental_hooks_prof_sample_free)
 CTL_PROTO(experimental_hooks_safety_check_abort)
 CTL_PROTO(experimental_thread_activity_callback)
 CTL_PROTO(experimental_utilization_query)
@@ -858,6 +860,8 @@ static const ctl_named_node_t experimental_hooks_node[] = {
 	{NAME("remove"),	CTL(experimental_hooks_remove)},
 	{NAME("prof_backtrace"),	CTL(experimental_hooks_prof_backtrace)},
 	{NAME("prof_dump"),	CTL(experimental_hooks_prof_dump)},
+	{NAME("prof_sample"),	CTL(experimental_hooks_prof_sample)},
+	{NAME("prof_sample_free"),	CTL(experimental_hooks_prof_sample_free)},
 	{NAME("safety_check_abort"),	CTL(experimental_hooks_safety_check_abort)},
 };
 
@@ -3505,6 +3509,62 @@ label_return:
 	return ret;
 }
 
+static int
+experimental_hooks_prof_sample_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_sample_hook_t old_hook =
+		    prof_sample_hook_get();
+		READ(old_hook, prof_sample_hook_t);
+	}
+	if (newp != NULL) {
+		if (!opt_prof) {
+			ret = ENOENT;
+			goto label_return;
+		}
+		prof_sample_hook_t new_hook JEMALLOC_CC_SILENCE_INIT(NULL);
+		WRITE(new_hook, prof_sample_hook_t);
+		prof_sample_hook_set(new_hook);
+	}
+	ret = 0;
+label_return:
+	return ret;
+}
+
+static int
+experimental_hooks_prof_sample_free_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_sample_free_hook_t old_hook =
+		    prof_sample_free_hook_get();
+		READ(old_hook, prof_sample_free_hook_t);
+	}
+	if (newp != NULL) {
+		if (!opt_prof) {
+			ret = ENOENT;
+			goto label_return;
+		}
+		prof_sample_free_hook_t new_hook JEMALLOC_CC_SILENCE_INIT(NULL);
+		WRITE(new_hook, prof_sample_free_hook_t);
+		prof_sample_free_hook_set(new_hook);
+	}
+	ret = 0;
+label_return:
+	return ret;
+}
+
 /* For integration test purpose only.  No plan to move out of experimental. */
 static int
 experimental_hooks_safety_check_abort_ctl(tsd_t *tsd, const size_t *mib,
diff --git a/src/prof.c b/src/prof.c
index 3deac0b5..91425371 100644
--- a/src/prof.c
+++ b/src/prof.c
@@ -78,6 +78,12 @@ atomic_p_t prof_backtrace_hook;
 /* Logically a prof_dump_hook_t. */
 atomic_p_t prof_dump_hook;
 
+/* Logically a prof_sample_hook_t. */
+atomic_p_t prof_sample_hook;
+
+/* Logically a prof_sample_free_hook_t. */
+atomic_p_t prof_sample_free_hook;
+
 /******************************************************************************/
 
 void
@@ -145,10 +151,20 @@ prof_malloc_sample_object(tsd_t *tsd, const void *ptr, size_t size,
 	if (opt_prof_stats) {
 		prof_stats_inc(tsd, szind, size);
 	}
+
+	/* Sample hook. */
+	prof_sample_hook_t prof_sample_hook = prof_sample_hook_get();
+	if (prof_sample_hook != NULL) {
+		prof_bt_t *bt = &tctx->gctx->bt;
+		pre_reentrancy(tsd, NULL);
+		prof_sample_hook(ptr, size, bt->vec, bt->len);
+		post_reentrancy(tsd);
+	}
 }
 
 void
-prof_free_sampled_object(tsd_t *tsd, size_t usize, prof_info_t *prof_info) {
+prof_free_sampled_object(tsd_t *tsd, const void *ptr, size_t usize,
+    prof_info_t *prof_info) {
 	cassert(config_prof);
 
 	assert(prof_info != NULL);
@@ -156,6 +172,16 @@ prof_free_sampled_object(tsd_t *tsd, size_t usize, prof_info_t *prof_info) {
 	assert((uintptr_t)tctx > (uintptr_t)1U);
 
 	szind_t szind = sz_size2index(usize);
+
+	/* Unsample hook. */
+	prof_sample_free_hook_t prof_sample_free_hook =
+	    prof_sample_free_hook_get();
+	if (prof_sample_free_hook != NULL) {
+		pre_reentrancy(tsd, NULL);
+		prof_sample_free_hook(ptr, usize);
+		post_reentrancy(tsd);
+	}
+
 	malloc_mutex_lock(tsd_tsdn(tsd), tctx->tdata->lock);
 
 	assert(tctx->cnts.curobjs > 0);
@@ -550,6 +576,28 @@ prof_dump_hook_get() {
 }
 
 void
+prof_sample_hook_set(prof_sample_hook_t hook) {
+	atomic_store_p(&prof_sample_hook, hook, ATOMIC_RELEASE);
+}
+
+prof_sample_hook_t
+prof_sample_hook_get() {
+	return (prof_sample_hook_t)atomic_load_p(&prof_sample_hook,
+	    ATOMIC_ACQUIRE);
+}
+
+void
+prof_sample_free_hook_set(prof_sample_free_hook_t hook) {
+	atomic_store_p(&prof_sample_free_hook, hook, ATOMIC_RELEASE);
+}
+
+prof_sample_free_hook_t
+prof_sample_free_hook_get() {
+	return (prof_sample_free_hook_t)atomic_load_p(&prof_sample_free_hook,
+	    ATOMIC_ACQUIRE);
+}
+
+void
 prof_boot0(void) {
 	cassert(config_prof);
 
diff --git a/src/prof_sys.c b/src/prof_sys.c
index 99fa3a77..d2487fd6 100644
--- a/src/prof_sys.c
+++ b/src/prof_sys.c
@@ -431,6 +431,8 @@ void
 prof_hooks_init() {
 	prof_backtrace_hook_set(&prof_backtrace_impl);
 	prof_dump_hook_set(NULL);
+	prof_sample_hook_set(NULL);
+	prof_sample_free_hook_set(NULL);
 }
 
 void
diff --git a/test/unit/prof_hook.c b/test/unit/prof_hook.c
index fc06d84e..a48b237b 100644
--- a/test/unit/prof_hook.c
+++ b/test/unit/prof_hook.c
@@ -1,11 +1,23 @@
 #include "test/jemalloc_test.h"
 
+/*
+ * The MALLOC_CONF of this test has lg_prof_sample:0, meaning that every single
+ * allocation will be sampled (and trigger relevant hooks).
+ */
+
 const char *dump_filename = "/dev/null";
 
-prof_backtrace_hook_t default_hook;
+prof_backtrace_hook_t default_bt_hook;
 
 bool mock_bt_hook_called = false;
 bool mock_dump_hook_called = false;
+bool mock_prof_sample_hook_called = false;
+bool mock_prof_sample_free_hook_called = false;
+
+void *sampled_ptr = NULL;
+size_t sampled_ptr_sz = 0;
+void *free_sampled_ptr = NULL;
+size_t free_sampled_ptr_sz = 0;
 
 void
 mock_bt_hook(void **vec, unsigned *len, unsigned max_len) {
@@ -18,7 +30,7 @@ mock_bt_hook(void **vec, unsigned *len, unsigned max_len) {
 
 void
 mock_bt_augmenting_hook(void **vec, unsigned *len, unsigned max_len) {
-	default_hook(vec, len, max_len);
+	default_bt_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");
@@ -47,6 +59,24 @@ mock_dump_hook(const char *filename) {
 	    "Incorrect file name passed to the dump hook");
 }
 
+void
+mock_prof_sample_hook(const void *ptr, size_t sz, void **vec, unsigned len) {
+	mock_prof_sample_hook_called = true;
+	sampled_ptr = (void *)ptr;
+	sampled_ptr_sz = sz;
+	for (unsigned i = 0; i < len; i++) {
+		expect_ptr_not_null((void **)vec[i],
+		    "Backtrace should not contain NULL");
+	}
+}
+
+void
+mock_prof_sample_free_hook(const void *ptr, size_t sz) {
+	mock_prof_sample_free_hook_called = true;
+	free_sampled_ptr = (void *)ptr;
+	free_sampled_ptr_sz = sz;
+}
+
 TEST_BEGIN(test_prof_backtrace_hook_replace) {
 
 	test_skip_if(!config_prof);
@@ -63,10 +93,10 @@ TEST_BEGIN(test_prof_backtrace_hook_replace) {
 	    NULL, 0, (void *)&null_hook,  sizeof(null_hook)),
 		EINVAL, "Incorrectly allowed NULL backtrace hook");
 
-	size_t default_hook_sz = sizeof(prof_backtrace_hook_t);
+	size_t default_bt_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,
+	    (void *)&default_bt_hook, &default_bt_hook_sz, (void *)&hook,
 	    sizeof(hook)), 0, "Unexpected mallctl failure setting hook");
 
 	void *p1 = mallocx(1, 0);
@@ -77,8 +107,8 @@ TEST_BEGIN(test_prof_backtrace_hook_replace) {
 	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,
+	    (void *)&current_hook, &current_hook_sz, (void *)&default_bt_hook,
+	    sizeof(default_bt_hook)), 0,
 	    "Unexpected mallctl failure resetting hook to default");
 
 	expect_ptr_eq(current_hook, hook,
@@ -100,10 +130,10 @@ TEST_BEGIN(test_prof_backtrace_hook_augment) {
 
 	expect_false(mock_bt_hook_called, "Called mock hook before it's set");
 
-	size_t default_hook_sz = sizeof(prof_backtrace_hook_t);
+	size_t default_bt_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,
+	    (void *)&default_bt_hook, &default_bt_hook_sz, (void *)&hook,
 	    sizeof(hook)), 0, "Unexpected mallctl failure setting hook");
 
 	void *p1 = mallocx(1, 0);
@@ -114,8 +144,8 @@ TEST_BEGIN(test_prof_backtrace_hook_augment) {
 	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,
+	    (void *)&current_hook, &current_hook_sz, (void *)&default_bt_hook,
+	    sizeof(default_bt_hook)), 0,
 	    "Unexpected mallctl failure resetting hook to default");
 
 	expect_ptr_eq(current_hook, hook,
@@ -138,10 +168,10 @@ TEST_BEGIN(test_prof_dump_hook) {
 
 	expect_false(mock_dump_hook_called, "Called dump hook before it's set");
 
-	size_t default_hook_sz = sizeof(prof_dump_hook_t);
+	size_t default_bt_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,
+	    (void *)&default_bt_hook, &default_bt_hook_sz, (void *)&hook,
 	    sizeof(hook)), 0, "Unexpected mallctl failure setting hook");
 
 	expect_d_eq(mallctl("prof.dump", NULL, NULL, (void *)&dump_filename,
@@ -152,8 +182,8 @@ TEST_BEGIN(test_prof_dump_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,
+	    (void *)&current_hook, &current_hook_sz, (void *)&default_bt_hook,
+	    sizeof(default_bt_hook)), 0,
 	    "Unexpected mallctl failure resetting hook to default");
 
 	expect_ptr_eq(current_hook, hook,
@@ -161,10 +191,144 @@ TEST_BEGIN(test_prof_dump_hook) {
 }
 TEST_END
 
+/* Need the do_write flag because NULL is a valid to_write value. */
+static void
+read_write_prof_sample_hook(prof_sample_hook_t *to_read, bool do_write,
+    prof_sample_hook_t to_write) {
+	size_t hook_sz = sizeof(prof_sample_hook_t);
+	expect_d_eq(mallctl("experimental.hooks.prof_sample",
+	    (void *)to_read, &hook_sz, do_write ? &to_write : NULL, hook_sz), 0,
+	    "Unexpected prof_sample_hook mallctl failure");
+}
+
+static void
+write_prof_sample_hook(prof_sample_hook_t new_hook) {
+	read_write_prof_sample_hook(NULL, true, new_hook);
+}
+
+static prof_sample_hook_t
+read_prof_sample_hook(void) {
+	prof_sample_hook_t curr_hook;
+	read_write_prof_sample_hook(&curr_hook, false, NULL);
+
+	return curr_hook;
+}
+
+static void
+read_write_prof_sample_free_hook(prof_sample_free_hook_t *to_read,
+    bool do_write, prof_sample_free_hook_t to_write) {
+	size_t hook_sz = sizeof(prof_sample_free_hook_t);
+	expect_d_eq(mallctl("experimental.hooks.prof_sample_free",
+	    (void *)to_read, &hook_sz, do_write ? &to_write : NULL, hook_sz), 0,
+	    "Unexpected prof_sample_free_hook mallctl failure");
+}
+
+static void
+write_prof_sample_free_hook(prof_sample_free_hook_t new_hook) {
+	read_write_prof_sample_free_hook(NULL, true, new_hook);
+}
+
+static prof_sample_free_hook_t
+read_prof_sample_free_hook(void) {
+	prof_sample_free_hook_t curr_hook;
+	read_write_prof_sample_free_hook(&curr_hook, false, NULL);
+
+	return curr_hook;
+}
+
+static void
+check_prof_sample_hooks(bool sample_hook_set, bool sample_free_hook_set) {
+	expect_false(mock_prof_sample_hook_called,
+	    "Should not have called prof_sample hook");
+	expect_false(mock_prof_sample_free_hook_called,
+	    "Should not have called prof_sample_free hook");
+	expect_ptr_null(sampled_ptr, "Unexpected sampled ptr");
+	expect_zu_eq(sampled_ptr_sz, 0, "Unexpected sampled ptr size");
+	expect_ptr_null(free_sampled_ptr, "Unexpected free sampled ptr");
+	expect_zu_eq(free_sampled_ptr_sz, 0,
+	    "Unexpected free sampled ptr size");
+
+	prof_sample_hook_t curr_hook = read_prof_sample_hook();
+	expect_ptr_eq(curr_hook, sample_hook_set ? mock_prof_sample_hook : NULL,
+	    "Unexpected non NULL default hook");
+
+	prof_sample_free_hook_t curr_free_hook = read_prof_sample_free_hook();
+	expect_ptr_eq(curr_free_hook, sample_free_hook_set ?
+	    mock_prof_sample_free_hook : NULL,
+	    "Unexpected non NULL default hook");
+
+	size_t alloc_sz = 10;
+	void *p = mallocx(alloc_sz, 0);
+	expect_ptr_not_null(p, "Failed to allocate");
+	expect_true(mock_prof_sample_hook_called == sample_hook_set,
+	   "Incorrect prof_sample hook usage");
+	if (sample_hook_set) {
+		expect_ptr_eq(p, sampled_ptr, "Unexpected sampled ptr");
+		expect_zu_eq(alloc_sz, sampled_ptr_sz,
+		    "Unexpected sampled usize");
+	}
+
+	dallocx(p, 0);
+	expect_true(mock_prof_sample_free_hook_called == sample_free_hook_set,
+	   "Incorrect prof_sample_free hook usage");
+	if (sample_free_hook_set) {
+		size_t usz = sz_s2u(alloc_sz);
+		expect_ptr_eq(p, free_sampled_ptr, "Unexpected sampled ptr");
+		expect_zu_eq(usz, free_sampled_ptr_sz, "Unexpected sampled usize");
+	}
+
+	sampled_ptr = free_sampled_ptr = NULL;
+	sampled_ptr_sz = free_sampled_ptr_sz = 0;
+	mock_prof_sample_hook_called = false;
+	mock_prof_sample_free_hook_called = false;
+}
+
+TEST_BEGIN(test_prof_sample_hooks) {
+	test_skip_if(!config_prof);
+
+	check_prof_sample_hooks(false, false);
+
+	write_prof_sample_hook(mock_prof_sample_hook);
+	check_prof_sample_hooks(true, false);
+
+	write_prof_sample_free_hook(mock_prof_sample_free_hook);
+	check_prof_sample_hooks(true, true);
+
+	write_prof_sample_hook(NULL);
+	check_prof_sample_hooks(false, true);
+
+	write_prof_sample_free_hook(NULL);
+	check_prof_sample_hooks(false, false);
+
+	/* Test read+write together. */
+	prof_sample_hook_t sample_hook;
+	read_write_prof_sample_hook(&sample_hook, true, mock_prof_sample_hook);
+	expect_ptr_null(sample_hook, "Unexpected non NULL default hook");
+	check_prof_sample_hooks(true, false);
+
+	prof_sample_free_hook_t sample_free_hook;
+	read_write_prof_sample_free_hook(&sample_free_hook, true,
+	    mock_prof_sample_free_hook);
+	expect_ptr_null(sample_free_hook, "Unexpected non NULL default hook");
+	check_prof_sample_hooks(true, true);
+
+	read_write_prof_sample_hook(&sample_hook, true, NULL);
+	expect_ptr_eq(sample_hook, mock_prof_sample_hook,
+	    "Unexpected prof_sample hook");
+	check_prof_sample_hooks(false, true);
+
+	read_write_prof_sample_free_hook(&sample_free_hook, true, NULL);
+	expect_ptr_eq(sample_free_hook, mock_prof_sample_free_hook,
+	    "Unexpected prof_sample_free hook");
+	check_prof_sample_hooks(false, false);
+}
+TEST_END
+
 int
 main(void) {
 	return test(
 	    test_prof_backtrace_hook_replace,
 	    test_prof_backtrace_hook_augment,
-	    test_prof_dump_hook);
+	    test_prof_dump_hook,
+	    test_prof_sample_hooks);
 }