Encapsulate buffer allocation failure in buffered writer

include/jemalloc/internal/buf_writer.h

@@ -13,10 +13,8 @@
 typedef void (write_cb_t)(void *, const char *);
 
 typedef struct {
-	write_cb_t *public_write_cb;
+	write_cb_t *write_cb;
-	void *public_cbopaque;
+	void *cbopaque;
-	write_cb_t *private_write_cb;
-	void *private_cbopaque;
 	char *buf;
 	size_t buf_size;
 	size_t buf_end;
@@ -25,9 +23,8 @@ typedef struct {
 
 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);
+write_cb_t buf_writer_cb;
 void buf_writer_terminate(tsdn_t *tsdn, buf_writer_t *buf_writer);
 
 #endif /* JEMALLOC_INTERNAL_BUF_WRITER_H */

src/buf_writer.c

@@ -25,28 +25,29 @@ buf_writer_free_internal_buf(tsdn_t *tsdn, void *buf) {
 	}
 }
 
-static write_cb_t buf_writer_cb;
-
 static void
 buf_writer_assert(buf_writer_t *buf_writer) {
+	assert(buf_writer != NULL);
+	assert(buf_writer->write_cb != NULL);
 	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);
+		assert(buf_writer->internal_buf);
 	}
+	assert(buf_writer->buf_end <= buf_writer->buf_size);
 }
 
 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) {
+	if (write_cb != NULL) {
+		buf_writer->write_cb = write_cb;
+	} else {
+		buf_writer->write_cb = je_malloc_message != NULL ?
+		    je_malloc_message : wrtmessage;
+	}
+	buf_writer->cbopaque = cbopaque;
 	assert(buf_len >= 2);
 	if (buf != NULL) {
 		buf_writer->buf = buf;
@@ -56,36 +57,14 @@ buf_writer_init(tsdn_t *tsdn, buf_writer_t *buf_writer, write_cb_t *write_cb,
 		    buf_len);
 		buf_writer->internal_buf = true;
 	}
-	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;
 	}
-}
+	buf_writer->buf_end = 0;
-
-write_cb_t *
-buf_writer_get_write_cb(buf_writer_t *buf_writer) {
 	buf_writer_assert(buf_writer);
-	return buf_writer->public_write_cb;
+	return buf_writer->buf == NULL;
-}
-
-void *
-buf_writer_get_cbopaque(buf_writer_t *buf_writer) {
-	buf_writer_assert(buf_writer);
-	return buf_writer->public_cbopaque;
 }
 
 void
@@ -94,34 +73,31 @@ buf_writer_flush(buf_writer_t *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->write_cb(buf_writer->cbopaque, buf_writer->buf);
-		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;
+	buf_writer_assert(buf_writer);
 }
 
-static void
+void
 buf_writer_cb(void *buf_writer_arg, const char *s) {
 	buf_writer_t *buf_writer = (buf_writer_t *)buf_writer_arg;
 	buf_writer_assert(buf_writer);
-	assert(buf_writer->buf != NULL);
+	if (buf_writer->buf == NULL) {
-	assert(buf_writer->buf_end <= buf_writer->buf_size);
+		buf_writer->write_cb(buf_writer->cbopaque, s);
-	size_t i, slen, n, s_remain, buf_remain;
+		return;
+	}
+	size_t i, slen, n;
 	for (i = 0, slen = strlen(s); i < slen; i += n) {
 		if (buf_writer->buf_end == buf_writer->buf_size) {
 			buf_writer_flush(buf_writer);
 		}
-		s_remain = slen - i;
+		size_t s_remain = slen - i;
-		buf_remain = buf_writer->buf_size - buf_writer->buf_end;
+		size_t buf_remain = buf_writer->buf_size - buf_writer->buf_end;
 		n = s_remain < buf_remain ? s_remain : buf_remain;
 		memcpy(buf_writer->buf + buf_writer->buf_end, s + i, n);
 		buf_writer->buf_end += n;
+		buf_writer_assert(buf_writer);
 	}
 	assert(i == slen);
 }

src/jemalloc.c

@@ -3847,8 +3847,7 @@ je_malloc_stats_print(void (*write_cb)(void *, const char *), void *cbopaque,
 		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),
+		stats_print(buf_writer_cb, &buf_writer, opts);
-		    buf_writer_get_cbopaque(&buf_writer), opts);
 		buf_writer_terminate(tsdn, &buf_writer);
 	}
 

src/prof_log.c

@@ -632,9 +632,8 @@ prof_log_stop(tsdn_t *tsdn) {
 	buf_writer_t 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,
+	emitter_init(&emitter, emitter_output_json_compact, buf_writer_cb,
-	    buf_writer_get_write_cb(&buf_writer),
+	    &buf_writer);
-	    buf_writer_get_cbopaque(&buf_writer));
 
 	emitter_begin(&emitter);
 	prof_log_emit_metadata(&emitter);

src/prof_recent.c

@@ -466,9 +466,8 @@ prof_recent_alloc_dump(tsd_t *tsd, void (*write_cb)(void *, const char *),
 	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,
+	emitter_init(&emitter, emitter_output_json_compact, buf_writer_cb,
-	    buf_writer_get_write_cb(&buf_writer),
+	    &buf_writer);
-	    buf_writer_get_cbopaque(&buf_writer));
 	emitter_begin(&emitter);
 
 	malloc_mutex_lock(tsd_tsdn(tsd), &prof_recent_alloc_mtx);

test/unit/buf_writer.c

@@ -24,10 +24,8 @@ 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;
-	assert(buf_writer->buf != NULL);
-	write_cb_t *write_cb = buf_writer_get_write_cb(buf_writer);
-	void *cbopaque = buf_writer_get_cbopaque(buf_writer);
 
+	assert(buf_writer->buf != NULL);
 	memset(s, 'a', UNIT_MAX);
 	arg = 4; /* Starting value of random argument. */
 	arg_store = arg;
@@ -39,7 +37,7 @@ test_buf_writer_body(tsdn_t *tsdn, buf_writer_t *buf_writer) {
 			remain = 0;
 			for (i = 1; i <= n_unit; ++i) {
 				arg = prng_lg_range_u64(&arg, 64);
-				write_cb(cbopaque, s);
+				buf_writer_cb(buf_writer, s);
 				remain += unit;
 				if (remain > buf_writer->buf_size) {
 					/* Flushes should have happened. */
@@ -89,10 +87,6 @@ TEST_BEGIN(test_buf_write_oom) {
 	assert_true(buf_writer_init(tsdn, &buf_writer, test_write_cb, &arg,
 	    NULL, SC_LARGE_MAXCLASS + 1), "buf_writer_init() should OOM");
 	assert(buf_writer.buf == 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;
@@ -108,7 +102,7 @@ TEST_BEGIN(test_buf_write_oom) {
 			test_write_len = 0;
 			for (i = 1; i <= n_unit; ++i) {
 				arg = prng_lg_range_u64(&arg, 64);
-				write_cb(cbopaque, s);
+				buf_writer_cb(&buf_writer, s);
 				assert_u64_eq(arg_store, arg,
 				    "Call back argument didn't get through");
 				assert_zu_eq(test_write_len, i * unit,