From 9cac3fa8f588c828a0a94bdc911383d2952b40e0 Mon Sep 17 00:00:00 2001 From: Yinan Zhang Date: Mon, 3 Feb 2020 15:56:13 -0800 Subject: [PATCH] Encapsulate buffer allocation in buffered writer --- include/jemalloc/internal/buf_writer.h | 27 +++--- src/buf_writer.c | 118 +++++++++++++++++++++++-- src/jemalloc.c | 19 ++-- src/prof_log.c | 22 ++--- src/prof_recent.c | 24 ++--- test/unit/buf_writer.c | 85 +++++++++++++++--- 6 files changed, 215 insertions(+), 80 deletions(-) diff --git a/include/jemalloc/internal/buf_writer.h b/include/jemalloc/internal/buf_writer.h index b2644a86..c1e2a827 100644 --- a/include/jemalloc/internal/buf_writer.h +++ b/include/jemalloc/internal/buf_writer.h @@ -10,27 +10,24 @@ * some "option like" content for the write_cb, so it doesn't matter. */ +typedef void (write_cb_t)(void *, const char *); + typedef struct { - void (*write_cb)(void *, const char *); - void *cbopaque; + write_cb_t *public_write_cb; + void *public_cbopaque; + write_cb_t *private_write_cb; + void *private_cbopaque; char *buf; size_t buf_size; size_t buf_end; + bool internal_buf; } buf_writer_t; -JEMALLOC_ALWAYS_INLINE void -buf_writer_init(buf_writer_t *buf_writer, void (*write_cb)(void *, - const char *), void *cbopaque, char *buf, size_t buf_len) { - buf_writer->write_cb = write_cb; - buf_writer->cbopaque = cbopaque; - assert(buf != NULL); - buf_writer->buf = buf; - assert(buf_len >= 2); - buf_writer->buf_size = buf_len - 1; /* Allowing for '\0' at the end. */ - buf_writer->buf_end = 0; -} - +bool buf_writer_init(tsdn_t *tsdn, buf_writer_t *buf_writer, + write_cb_t *write_cb, void *cbopaque, char *buf, size_t buf_len); +write_cb_t *buf_writer_get_write_cb(buf_writer_t *buf_writer); +void *buf_writer_get_cbopaque(buf_writer_t *buf_writer); void buf_writer_flush(buf_writer_t *buf_writer); -void buf_writer_cb(void *buf_writer_arg, const char *s); +void buf_writer_terminate(tsdn_t *tsdn, buf_writer_t *buf_writer); #endif /* JEMALLOC_INTERNAL_BUF_WRITER_H */ diff --git a/src/buf_writer.c b/src/buf_writer.c index aed7d4a8..bb8763b9 100644 --- a/src/buf_writer.c +++ b/src/buf_writer.c @@ -5,23 +5,114 @@ #include "jemalloc/internal/buf_writer.h" #include "jemalloc/internal/malloc_io.h" -void -buf_writer_flush(buf_writer_t *buf_writer) { - assert(buf_writer->buf_end <= buf_writer->buf_size); - buf_writer->buf[buf_writer->buf_end] = '\0'; - if (buf_writer->write_cb == NULL) { - buf_writer->write_cb = je_malloc_message != NULL ? - je_malloc_message : wrtmessage; +static void * +buf_writer_allocate_internal_buf(tsdn_t *tsdn, size_t buf_len) { +#ifdef JEMALLOC_JET + if (buf_len > SC_LARGE_MAXCLASS) { + return NULL; + } +#else + assert(buf_len <= SC_LARGE_MAXCLASS); +#endif + return iallocztm(tsdn, buf_len, sz_size2index(buf_len), false, NULL, + true, arena_get(tsdn, 0, false), true); +} + +static void +buf_writer_free_internal_buf(tsdn_t *tsdn, void *buf) { + if (buf != NULL) { + idalloctm(tsdn, buf, NULL, NULL, true, true); + } +} + +static write_cb_t buf_writer_cb; + +static void +buf_writer_assert(buf_writer_t *buf_writer) { + if (buf_writer->buf != NULL) { + assert(buf_writer->public_write_cb == buf_writer_cb); + assert(buf_writer->public_cbopaque == buf_writer); + assert(buf_writer->private_write_cb != buf_writer_cb); + assert(buf_writer->private_cbopaque != buf_writer); + assert(buf_writer->buf_size > 0); + } else { + assert(buf_writer->public_write_cb != buf_writer_cb); + assert(buf_writer->public_cbopaque != buf_writer); + assert(buf_writer->private_write_cb == NULL); + assert(buf_writer->private_cbopaque == NULL); + assert(buf_writer->buf_size == 0); + } +} + +bool +buf_writer_init(tsdn_t *tsdn, buf_writer_t *buf_writer, write_cb_t *write_cb, + void *cbopaque, char *buf, size_t buf_len) { + assert(buf_len >= 2); + if (buf != NULL) { + buf_writer->buf = buf; + buf_writer->internal_buf = false; + } else { + buf_writer->buf = buf_writer_allocate_internal_buf(tsdn, + buf_len); + buf_writer->internal_buf = true; } - buf_writer->write_cb(buf_writer->cbopaque, buf_writer->buf); buf_writer->buf_end = 0; + if (buf_writer->buf != NULL) { + buf_writer->public_write_cb = buf_writer_cb; + buf_writer->public_cbopaque = buf_writer; + buf_writer->private_write_cb = write_cb; + buf_writer->private_cbopaque = cbopaque; + buf_writer->buf_size = buf_len - 1; /* Allowing for '\0'. */ + buf_writer_assert(buf_writer); + return false; + } else { + buf_writer->public_write_cb = write_cb; + buf_writer->public_cbopaque = cbopaque; + buf_writer->private_write_cb = NULL; + buf_writer->private_cbopaque = NULL; + buf_writer->buf_size = 0; + buf_writer_assert(buf_writer); + return true; + } +} + +write_cb_t * +buf_writer_get_write_cb(buf_writer_t *buf_writer) { + buf_writer_assert(buf_writer); + return buf_writer->public_write_cb; +} + +void * +buf_writer_get_cbopaque(buf_writer_t *buf_writer) { + buf_writer_assert(buf_writer); + return buf_writer->public_cbopaque; } void +buf_writer_flush(buf_writer_t *buf_writer) { + buf_writer_assert(buf_writer); + if (buf_writer->buf == NULL) { + return; + } + assert(buf_writer->buf_end <= buf_writer->buf_size); + buf_writer->buf[buf_writer->buf_end] = '\0'; + if (buf_writer->private_write_cb == NULL) { + buf_writer->private_write_cb = je_malloc_message != NULL ? + je_malloc_message : wrtmessage; + } + assert(buf_writer->private_write_cb != NULL); + buf_writer->private_write_cb(buf_writer->private_cbopaque, + buf_writer->buf); + buf_writer->buf_end = 0; +} + +static void buf_writer_cb(void *buf_writer_arg, const char *s) { buf_writer_t *buf_writer = (buf_writer_t *)buf_writer_arg; - size_t i, slen, n, s_remain, buf_remain; + buf_writer_assert(buf_writer); + assert(buf_writer->buf != NULL); assert(buf_writer->buf_end <= buf_writer->buf_size); + size_t i, slen, n, s_remain, buf_remain; for (i = 0, slen = strlen(s); i < slen; i += n) { if (buf_writer->buf_end == buf_writer->buf_size) { buf_writer_flush(buf_writer); @@ -34,3 +125,12 @@ buf_writer_cb(void *buf_writer_arg, const char *s) { } assert(i == slen); } + +void +buf_writer_terminate(tsdn_t *tsdn, buf_writer_t *buf_writer) { + buf_writer_assert(buf_writer); + buf_writer_flush(buf_writer); + if (buf_writer->internal_buf) { + buf_writer_free_internal_buf(tsdn, buf_writer->buf); + } +} diff --git a/src/jemalloc.c b/src/jemalloc.c index 35c490be..ddb29e38 100644 --- a/src/jemalloc.c +++ b/src/jemalloc.c @@ -3740,19 +3740,12 @@ je_malloc_stats_print(void (*write_cb)(void *, const char *), void *cbopaque, if (config_debug) { stats_print(write_cb, cbopaque, opts); } else { - char *buf = (char *)iallocztm(tsdn, STATS_PRINT_BUFSIZE, - sz_size2index(STATS_PRINT_BUFSIZE), false, NULL, true, - arena_get(TSDN_NULL, 0, true), true); - if (buf == NULL) { - stats_print(write_cb, cbopaque, opts); - } else { - buf_writer_t buf_writer; - buf_writer_init(&buf_writer, write_cb, cbopaque, buf, - STATS_PRINT_BUFSIZE); - stats_print(buf_writer_cb, &buf_writer, opts); - buf_writer_flush(&buf_writer); - idalloctm(tsdn, buf, NULL, NULL, true, true); - } + buf_writer_t buf_writer; + buf_writer_init(tsdn, &buf_writer, write_cb, cbopaque, NULL, + STATS_PRINT_BUFSIZE); + stats_print(buf_writer_get_write_cb(&buf_writer), + buf_writer_get_cbopaque(&buf_writer), opts); + buf_writer_terminate(tsdn, &buf_writer); } check_entry_exit_locking(tsdn); diff --git a/src/prof_log.c b/src/prof_log.c index 95cf246d..c29fa350 100644 --- a/src/prof_log.c +++ b/src/prof_log.c @@ -629,19 +629,12 @@ prof_log_stop(tsdn_t *tsdn) { struct prof_emitter_cb_arg_s arg; arg.fd = fd; - char *buf = (char *)iallocztm(tsdn, PROF_LOG_STOP_BUFSIZE, - sz_size2index(PROF_LOG_STOP_BUFSIZE), false, NULL, true, - arena_get(TSDN_NULL, 0, true), true); buf_writer_t buf_writer; - if (buf == NULL) { - emitter_init(&emitter, emitter_output_json_compact, - prof_emitter_write_cb, &arg); - } else { - buf_writer_init(&buf_writer, prof_emitter_write_cb, &arg, buf, - PROF_LOG_STOP_BUFSIZE); - emitter_init(&emitter, emitter_output_json_compact, - buf_writer_cb, &buf_writer); - } + buf_writer_init(tsdn, &buf_writer, prof_emitter_write_cb, &arg, NULL, + PROF_LOG_STOP_BUFSIZE); + emitter_init(&emitter, emitter_output_json_compact, + buf_writer_get_write_cb(&buf_writer), + buf_writer_get_cbopaque(&buf_writer)); emitter_begin(&emitter); prof_log_emit_metadata(&emitter); @@ -650,10 +643,7 @@ prof_log_stop(tsdn_t *tsdn) { prof_log_emit_allocs(tsd, &emitter); emitter_end(&emitter); - if (buf != NULL) { - buf_writer_flush(&buf_writer); - idalloctm(tsdn, buf, NULL, NULL, true, true); - } + buf_writer_terminate(tsdn, &buf_writer); /* Reset global state. */ if (log_tables_initialized) { diff --git a/src/prof_recent.c b/src/prof_recent.c index dde029ce..7a98cc58 100644 --- a/src/prof_recent.c +++ b/src/prof_recent.c @@ -462,20 +462,13 @@ dump_bt(emitter_t *emitter, prof_tctx_t *tctx) { void prof_recent_alloc_dump(tsd_t *tsd, void (*write_cb)(void *, const char *), void *cbopaque) { - char *buf = (char *)iallocztm(tsd_tsdn(tsd), PROF_RECENT_PRINT_BUFSIZE, - sz_size2index(PROF_RECENT_PRINT_BUFSIZE), false, NULL, true, - arena_get(tsd_tsdn(tsd), 0, false), true); - emitter_t emitter; buf_writer_t buf_writer; - if (buf == NULL) { - emitter_init(&emitter, emitter_output_json_compact, write_cb, - cbopaque); - } else { - buf_writer_init(&buf_writer, write_cb, cbopaque, buf, - PROF_RECENT_PRINT_BUFSIZE); - emitter_init(&emitter, emitter_output_json_compact, - buf_writer_cb, &buf_writer); - } + buf_writer_init(tsd_tsdn(tsd), &buf_writer, write_cb, cbopaque, NULL, + PROF_RECENT_PRINT_BUFSIZE); + emitter_t emitter; + emitter_init(&emitter, emitter_output_json_compact, + buf_writer_get_write_cb(&buf_writer), + buf_writer_get_cbopaque(&buf_writer)); emitter_begin(&emitter); malloc_mutex_lock(tsd_tsdn(tsd), &prof_recent_alloc_mtx); @@ -535,10 +528,7 @@ prof_recent_alloc_dump(tsd_t *tsd, void (*write_cb)(void *, const char *), malloc_mutex_unlock(tsd_tsdn(tsd), &prof_recent_alloc_mtx); emitter_end(&emitter); - if (buf != NULL) { - buf_writer_flush(&buf_writer); - idalloctm(tsd_tsdn(tsd), buf, NULL, NULL, true, true); - } + buf_writer_terminate(tsd_tsdn(tsd), &buf_writer); } #undef PROF_RECENT_PRINT_BUFSIZE diff --git a/test/unit/buf_writer.c b/test/unit/buf_writer.c index 63fd0c67..5171d618 100644 --- a/test/unit/buf_writer.c +++ b/test/unit/buf_writer.c @@ -7,6 +7,7 @@ static size_t test_write_len; static char test_buf[TEST_BUF_SIZE]; +static uint64_t arg; static uint64_t arg_store; static void test_write_cb(void *cbopaque, const char *s) { @@ -17,16 +18,16 @@ static void test_write_cb(void *cbopaque, const char *s) { "Test write overflowed"); } -TEST_BEGIN(test_buf_write) { +static void test_buf_writer_body(tsdn_t *tsdn, buf_writer_t *buf_writer) { char s[UNIT_MAX + 1]; size_t n_unit, remain, i; ssize_t unit; - uint64_t arg = 4; /* Starting value of random argument. */ - buf_writer_t buf_writer; - buf_writer_init(&buf_writer, test_write_cb, &arg, test_buf, - TEST_BUF_SIZE); + assert_ptr_not_null(buf_writer->buf, "Buffer is null"); + write_cb_t *write_cb = buf_writer_get_write_cb(buf_writer); + void *cbopaque = buf_writer_get_cbopaque(buf_writer); memset(s, 'a', UNIT_MAX); + arg = 4; /* Starting value of random argument. */ arg_store = arg; for (unit = UNIT_MAX; unit >= 0; --unit) { /* unit keeps decreasing, so strlen(s) is always unit. */ @@ -36,32 +37,96 @@ TEST_BEGIN(test_buf_write) { remain = 0; for (i = 1; i <= n_unit; ++i) { arg = prng_lg_range_u64(&arg, 64); - buf_writer_cb(&buf_writer, s); + write_cb(cbopaque, s); remain += unit; - if (remain > buf_writer.buf_size) { + if (remain > buf_writer->buf_size) { /* Flushes should have happened. */ assert_u64_eq(arg_store, arg, "Call " "back argument didn't get through"); - remain %= buf_writer.buf_size; + remain %= buf_writer->buf_size; if (remain == 0) { /* Last flush should be lazy. */ - remain += buf_writer.buf_size; + remain += buf_writer->buf_size; } } assert_zu_eq(test_write_len + remain, i * unit, "Incorrect length after writing %zu strings" " of length %zu", i, unit); } + buf_writer_flush(buf_writer); + assert_zu_eq(test_write_len, n_unit * unit, + "Incorrect length after flushing at the end of" + " writing %zu strings of length %zu", n_unit, unit); + } + } + buf_writer_terminate(tsdn, buf_writer); +} + +TEST_BEGIN(test_buf_write_static) { + buf_writer_t buf_writer; + tsdn_t *tsdn = tsdn_fetch(); + assert_false(buf_writer_init(tsdn, &buf_writer, test_write_cb, &arg, + test_buf, TEST_BUF_SIZE), + "buf_writer_init() should not encounter error on static buffer"); + test_buf_writer_body(tsdn, &buf_writer); +} +TEST_END + +TEST_BEGIN(test_buf_write_dynamic) { + buf_writer_t buf_writer; + tsdn_t *tsdn = tsdn_fetch(); + assert_false(buf_writer_init(tsdn, &buf_writer, test_write_cb, &arg, + NULL, TEST_BUF_SIZE), "buf_writer_init() should not OOM"); + test_buf_writer_body(tsdn, &buf_writer); +} +TEST_END + +TEST_BEGIN(test_buf_write_oom) { + buf_writer_t buf_writer; + tsdn_t *tsdn = tsdn_fetch(); + assert_true(buf_writer_init(tsdn, &buf_writer, test_write_cb, &arg, + NULL, SC_LARGE_MAXCLASS + 1), "buf_writer_init() should OOM"); + assert_ptr_null(buf_writer.buf, "Buffer should be null"); + write_cb_t *write_cb = buf_writer_get_write_cb(&buf_writer); + assert_ptr_eq(write_cb, test_write_cb, "Should use test_write_cb"); + void *cbopaque = buf_writer_get_cbopaque(&buf_writer); + assert_ptr_eq(cbopaque, &arg, "Should use arg"); + + char s[UNIT_MAX + 1]; + size_t n_unit, i; + ssize_t unit; + + memset(s, 'a', UNIT_MAX); + arg = 4; /* Starting value of random argument. */ + arg_store = arg; + for (unit = UNIT_MAX; unit >= 0; unit -= UNIT_MAX / 4) { + /* unit keeps decreasing, so strlen(s) is always unit. */ + s[unit] = '\0'; + for (n_unit = 1; n_unit <= 3; ++n_unit) { + test_write_len = 0; + for (i = 1; i <= n_unit; ++i) { + arg = prng_lg_range_u64(&arg, 64); + write_cb(cbopaque, s); + assert_u64_eq(arg_store, arg, + "Call back argument didn't get through"); + assert_zu_eq(test_write_len, i * unit, + "Incorrect length after writing %zu strings" + " of length %zu", i, unit); + } buf_writer_flush(&buf_writer); assert_zu_eq(test_write_len, n_unit * unit, "Incorrect length after flushing at the end of" " writing %zu strings of length %zu", n_unit, unit); } } + buf_writer_terminate(tsdn, &buf_writer); } TEST_END int main(void) { - return test(test_buf_write); + return test( + test_buf_write_static, + test_buf_write_dynamic, + test_buf_write_oom); }