Encapsulate buffer allocation in buffered writer

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 *);
+	write_cb_t *public_write_cb;
-	void *cbopaque;
+	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
+bool buf_writer_init(tsdn_t *tsdn, buf_writer_t *buf_writer,
-buf_writer_init(buf_writer_t *buf_writer, void (*write_cb)(void *,
+    write_cb_t *write_cb, void *cbopaque, char *buf, size_t buf_len);
-    const char *), void *cbopaque, char *buf, size_t buf_len) {
+write_cb_t *buf_writer_get_write_cb(buf_writer_t *buf_writer);
-	buf_writer->write_cb = write_cb;
+void *buf_writer_get_cbopaque(buf_writer_t *buf_writer);
-	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;
-}
-
 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 */

src/buf_writer.c

@@ -5,23 +5,114 @@
 #include "jemalloc/internal/buf_writer.h"
 #include "jemalloc/internal/malloc_io.h"
 
-void
+static void *
-buf_writer_flush(buf_writer_t *buf_writer) {
+buf_writer_allocate_internal_buf(tsdn_t *tsdn, size_t buf_len) {
-	assert(buf_writer->buf_end <= buf_writer->buf_size);
+#ifdef JEMALLOC_JET
-	buf_writer->buf[buf_writer->buf_end] = '\0';
+	if (buf_len > SC_LARGE_MAXCLASS) {
-	if (buf_writer->write_cb == NULL) {
+		return NULL;
-		buf_writer->write_cb = je_malloc_message != NULL ?
+	}
-		    je_malloc_message : wrtmessage;
+#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);
+	}
+}

src/jemalloc.c

@@ -3739,20 +3739,13 @@ 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,
+		buf_writer_init(tsdn, &buf_writer, write_cb, cbopaque, NULL,
 		    STATS_PRINT_BUFSIZE);
-			stats_print(buf_writer_cb, &buf_writer, opts);
+		stats_print(buf_writer_get_write_cb(&buf_writer),
-			buf_writer_flush(&buf_writer);
+		    buf_writer_get_cbopaque(&buf_writer), opts);
-			idalloctm(tsdn, buf, NULL, NULL, true, true);
+		buf_writer_terminate(tsdn, &buf_writer);
-		}
 	}
 
 	check_entry_exit_locking(tsdn);

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) {
+	buf_writer_init(tsdn, &buf_writer, prof_emitter_write_cb, &arg, 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_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_terminate(tsdn, &buf_writer);
-		buf_writer_flush(&buf_writer);
-		idalloctm(tsdn, buf, NULL, NULL, true, true);
-	}
 
 	/* Reset global state. */
 	if (log_tables_initialized) {

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) {
+	buf_writer_init(tsd_tsdn(tsd), &buf_writer, write_cb, cbopaque, 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_t emitter;
 	emitter_init(&emitter, emitter_output_json_compact,
-		    buf_writer_cb, &buf_writer);
+	    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_terminate(tsd_tsdn(tsd), &buf_writer);
-		buf_writer_flush(&buf_writer);
-		idalloctm(tsd_tsdn(tsd), buf, NULL, NULL, true, true);
-	}
 }
 #undef PROF_RECENT_PRINT_BUFSIZE
 

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. */
+	assert_ptr_not_null(buf_writer->buf, "Buffer is null");
-	buf_writer_t buf_writer;
+	write_cb_t *write_cb = buf_writer_get_write_cb(buf_writer);
-	buf_writer_init(&buf_writer, test_write_cb, &arg, test_buf,
+	void *cbopaque = buf_writer_get_cbopaque(buf_writer);
-	    TEST_BUF_SIZE);
 
 	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);
 }