From a9031a0970df9c999873617423f789bd46bfe619 Mon Sep 17 00:00:00 2001 From: Alex Lapenkou Date: Wed, 1 Sep 2021 13:00:01 -0700 Subject: [PATCH] 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 | 8 +- include/jemalloc/internal/prof_hook.h | 5 + src/ctl.c | 32 ++++++- src/prof.c | 16 +++- src/prof_sys.c | 14 ++- test/unit/prof_hook.c | 114 ++++++++++++++++++++++- 6 files changed, 178 insertions(+), 11 deletions(-) diff --git a/include/jemalloc/internal/prof_externs.h b/include/jemalloc/internal/prof_externs.h index 75d1d7a0..75dd90bf 100644 --- a/include/jemalloc/internal/prof_externs.h +++ b/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); diff --git a/include/jemalloc/internal/prof_hook.h b/include/jemalloc/internal/prof_hook.h index 277cd992..150d19d3 100644 --- a/include/jemalloc/internal/prof_hook.h +++ b/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 */ diff --git a/src/ctl.c b/src/ctl.c index 6bf1c946..3aaa5a74 100644 --- a/src/ctl.c +++ b/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) diff --git a/src/prof.c b/src/prof.c index d0cae0e9..625bcd73 100644 --- a/src/prof.c +++ b/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; diff --git a/src/prof_sys.c b/src/prof_sys.c index 1485e8b2..fd41e86c 100644 --- a/src/prof_sys.c +++ b/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); diff --git a/test/unit/prof_hook.c b/test/unit/prof_hook.c index 32d0e9ea..6480d930 100644 --- a/test/unit/prof_hook.c +++ b/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 *)¤t_hook, ¤t_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 *)¤t_hook, ¤t_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); }