From 7fc6b1b259fd1c38a59341ad555a47790da6f773 Mon Sep 17 00:00:00 2001 From: Yinan Zhang Date: Fri, 7 Jun 2019 14:04:59 -0700 Subject: [PATCH] Add buffered writer The buffered writer adopts a signature identical to `write_cb`, so that it can be plugged into anywhere `write_cb` appears. --- Makefile.in | 1 + include/jemalloc/internal/malloc_io.h | 25 +++++++++++ src/malloc_io.c | 30 +++++++++++++ test/unit/buf_writer.c | 64 +++++++++++++++++++++++++++ 4 files changed, 120 insertions(+) create mode 100644 test/unit/buf_writer.c diff --git a/Makefile.in b/Makefile.in index 40daf115..ef75d8ac 100644 --- a/Makefile.in +++ b/Makefile.in @@ -177,6 +177,7 @@ TESTS_UNIT := \ $(srcroot)test/unit/bitmap.c \ $(srcroot)test/unit/bit_util.c \ $(srcroot)test/unit/binshard.c \ + $(srcroot)test/unit/buf_writer.c \ $(srcroot)test/unit/ckh.c \ $(srcroot)test/unit/decay.c \ $(srcroot)test/unit/div.c \ diff --git a/include/jemalloc/internal/malloc_io.h b/include/jemalloc/internal/malloc_io.h index 1d1a414e..f5d16a5e 100644 --- a/include/jemalloc/internal/malloc_io.h +++ b/include/jemalloc/internal/malloc_io.h @@ -99,4 +99,29 @@ malloc_read_fd(int fd, void *buf, size_t count) { return (ssize_t)result; } +/******************************************************************************/ + +/* + * The rest is buffered writing utility. + * + * The only difference when using the buffered writer is that cbopaque is + * passed to write_cb only when the buffer is flushed. It would make a + * difference if cbopaque points to something that's changing for each write_cb + * call, or something that affects write_cb in a way dependent on the content + * of the output string. However, the most typical usage case in practice is + * that cbopaque points to some "option like" content for the write_cb, so it + * doesn't matter. + */ + +typedef struct { + void (*write_cb)(void *, const char *); + void *cbopaque; + char *buf; + size_t buf_size; /* must be one less than the capacity of buf array */ + size_t buf_end; +} buf_writer_arg_t; + +void buf_writer_flush(buf_writer_arg_t *arg); +void buffered_write_cb(void *buf_writer_arg, const char *s); + #endif /* JEMALLOC_INTERNAL_MALLOC_IO_H */ diff --git a/src/malloc_io.c b/src/malloc_io.c index d7cb0f52..2fae7570 100644 --- a/src/malloc_io.c +++ b/src/malloc_io.c @@ -664,6 +664,36 @@ malloc_printf(const char *format, ...) { va_end(ap); } +void +buf_writer_flush(buf_writer_arg_t *arg) { + assert(arg->buf_end <= arg->buf_size); + arg->buf[arg->buf_end] = '\0'; + if (arg->write_cb == NULL) { + arg->write_cb = je_malloc_message != NULL ? + je_malloc_message : wrtmessage; + } + arg->write_cb(arg->cbopaque, arg->buf); + arg->buf_end = 0; +} + +void +buffered_write_cb(void *buf_writer_arg, const char *s) { + buf_writer_arg_t *arg = (buf_writer_arg_t *)buf_writer_arg; + size_t i, slen, n, s_remain, buf_remain; + assert(arg->buf_end <= arg->buf_size); + for (i = 0, slen = strlen(s); i < slen; i += n) { + if (arg->buf_end == arg->buf_size) { + buf_writer_flush(arg); + } + s_remain = slen - i; + buf_remain = arg->buf_size - arg->buf_end; + n = s_remain < buf_remain ? s_remain : buf_remain; + memcpy(arg->buf + arg->buf_end, s + i, n); + arg->buf_end += n; + } + assert(i == slen); +} + /* * Restore normal assertion macros, in order to make it possible to compile all * C files as a single concatenation. diff --git a/test/unit/buf_writer.c b/test/unit/buf_writer.c new file mode 100644 index 00000000..4d8ae99b --- /dev/null +++ b/test/unit/buf_writer.c @@ -0,0 +1,64 @@ +#include "test/jemalloc_test.h" + +#define TEST_BUF_SIZE 16 +#define UNIT_MAX (TEST_BUF_SIZE * 3) + +static size_t test_write_len; +static char test_buf[TEST_BUF_SIZE]; +static uint64_t arg_store; + +static void test_write_cb(void *cbopaque, const char *s) { + size_t prev_test_write_len = test_write_len; + test_write_len += strlen(s); /* only increase the length */ + arg_store = *(uint64_t *)cbopaque; /* only pass along the argument */ + assert_zu_le(prev_test_write_len, test_write_len, + "Test write overflowed"); +} + +TEST_BEGIN(test_buf_write) { + 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_arg_t test_buf_arg = + {test_write_cb, &arg, test_buf, TEST_BUF_SIZE - 1, 0}; + + memset(s, 'a', UNIT_MAX); + arg_store = arg; + for (unit = UNIT_MAX; unit >= 0; --unit) { + /* 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; + remain = 0; + for (i = 1; i <= n_unit; ++i) { + arg = prng_lg_range_u64(&arg, 64); + buffered_write_cb(&test_buf_arg, s); + remain += unit; + if (remain > test_buf_arg.buf_size) { + /* Flushes should have happened. */ + assert_u64_eq(arg_store, arg, "Call " + "back argument didn't get through"); + remain %= test_buf_arg.buf_size; + if (remain == 0) { + /* Last flush should be lazy. */ + remain += test_buf_arg.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(&test_buf_arg); + 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); + } + } +} +TEST_END + +int +main(void) { + return test(test_buf_write); +}