From e032a1a1de75cf7faf087406a21789ced2b2f650 Mon Sep 17 00:00:00 2001
From: Yinan Zhang <zyn8950@gmail.com>
Date: Fri, 31 Jul 2020 15:56:38 -0700
Subject: [PATCH] Add a stress test for batch allocation

---
 Makefile.in               |  3 +-
 test/stress/batch_alloc.c | 88 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 1 deletion(-)
 create mode 100644 test/stress/batch_alloc.c

diff --git a/Makefile.in b/Makefile.in
index da094f08..506d9da3 100644
--- a/Makefile.in
+++ b/Makefile.in
@@ -298,7 +298,8 @@ TESTS_ANALYZE := $(srcroot)test/analyze/rand.c \
 TESTS_STRESS := $(srcroot)test/stress/microbench.c \
 	$(srcroot)test/stress/fill_flush.c \
 	$(srcroot)test/stress/large_microbench.c \
-	$(srcroot)test/stress/hookbench.c
+	$(srcroot)test/stress/hookbench.c \
+	$(srcroot)test/stress/batch_alloc.c
 
 
 TESTS := $(TESTS_UNIT) $(TESTS_INTEGRATION) $(TESTS_INTEGRATION_CPP) \
diff --git a/test/stress/batch_alloc.c b/test/stress/batch_alloc.c
new file mode 100644
index 00000000..b203e05e
--- /dev/null
+++ b/test/stress/batch_alloc.c
@@ -0,0 +1,88 @@
+#include "test/jemalloc_test.h"
+#include "test/bench.h"
+
+#define BATCH (1000 * 1000)
+#define HUGE_BATCH (100 * BATCH)
+static void *batch_ptrs[HUGE_BATCH];
+static void *item_ptrs[HUGE_BATCH];
+
+#define SIZE 7
+
+typedef struct batch_alloc_packet_s batch_alloc_packet_t;
+struct batch_alloc_packet_s {
+	void **ptrs;
+	size_t num;
+	size_t size;
+	int flags;
+};
+
+static void
+batch_alloc_wrapper(size_t batch) {
+	batch_alloc_packet_t batch_alloc_packet = {batch_ptrs, batch, SIZE, 0};
+	size_t filled;
+	size_t len = sizeof(size_t);
+	assert_d_eq(mallctl("experimental.batch_alloc", &filled, &len,
+	    &batch_alloc_packet, sizeof(batch_alloc_packet)), 0, "");
+	assert_zu_eq(filled, batch, "");
+}
+
+static void
+item_alloc_wrapper(size_t batch) {
+	for (size_t i = 0; i < batch; ++i) {
+		item_ptrs[i] = malloc(SIZE);
+	}
+}
+
+static void
+release_and_clear(void **ptrs, size_t batch) {
+	for (size_t i = 0; i < batch; ++i) {
+		void *p = ptrs[i];
+		assert_ptr_not_null(p, "allocation failed");
+		sdallocx(p, SIZE, 0);
+		ptrs[i] = NULL;
+	}
+}
+
+static void
+batch_alloc_small_can_repeat() {
+	batch_alloc_wrapper(BATCH);
+	release_and_clear(batch_ptrs, BATCH);
+}
+
+static void
+item_alloc_small_can_repeat() {
+	item_alloc_wrapper(BATCH);
+	release_and_clear(item_ptrs, BATCH);
+}
+
+TEST_BEGIN(test_small_batch_with_free) {
+	compare_funcs(10, 100,
+	    "batch allocation", batch_alloc_small_can_repeat,
+	    "item allocation", item_alloc_small_can_repeat);
+}
+TEST_END
+
+static void
+batch_alloc_huge_cannot_repeat() {
+	batch_alloc_wrapper(HUGE_BATCH);
+}
+
+static void
+item_alloc_huge_cannot_repeat() {
+	item_alloc_wrapper(HUGE_BATCH);
+}
+
+TEST_BEGIN(test_huge_batch_without_free) {
+	compare_funcs(0, 1,
+	    "batch allocation", batch_alloc_huge_cannot_repeat,
+	    "item allocation", item_alloc_huge_cannot_repeat);
+	release_and_clear(batch_ptrs, HUGE_BATCH);
+	release_and_clear(item_ptrs, HUGE_BATCH);
+}
+TEST_END
+
+int main(void) {
+	return test_no_reentrancy(
+	    test_small_batch_with_free,
+	    test_huge_batch_without_free);
+}