From be5e49f4fa09247a91557690cdaef42a82a83d6a Mon Sep 17 00:00:00 2001
From: Yinan Zhang <zyn8950@gmail.com>
Date: Thu, 22 Oct 2020 16:07:25 -0700
Subject: [PATCH] Add a batch mode for cache_bin_alloc()

---
 include/jemalloc/internal/cache_bin.h | 12 ++++++
 test/unit/cache_bin.c                 | 53 +++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/include/jemalloc/internal/cache_bin.h b/include/jemalloc/internal/cache_bin.h
index 551afc85..c1b8fc42 100644
--- a/include/jemalloc/internal/cache_bin.h
+++ b/include/jemalloc/internal/cache_bin.h
@@ -317,6 +317,18 @@ cache_bin_alloc(cache_bin_t *bin, bool *success) {
 	return cache_bin_alloc_impl(bin, success, true);
 }
 
+JEMALLOC_ALWAYS_INLINE cache_bin_sz_t
+cache_bin_alloc_batch(cache_bin_t *bin, size_t num, void **out) {
+	size_t n = cache_bin_ncached_get_internal(bin);
+	if (n > num) {
+		n = num;
+	}
+	memcpy(out, bin->stack_head, n * sizeof(void *));
+	bin->stack_head += n;
+	cache_bin_low_water_adjust(bin);
+	return n;
+}
+
 /*
  * Free an object into the given bin.  Fails only if the bin is full.
  */
diff --git a/test/unit/cache_bin.c b/test/unit/cache_bin.c
index 7798bfab..b31d07d2 100644
--- a/test/unit/cache_bin.c
+++ b/test/unit/cache_bin.c
@@ -52,6 +52,34 @@ do_flush_test(cache_bin_t *bin, cache_bin_info_t *info, void **ptrs,
 	}
 }
 
+static void
+do_batch_alloc_test(cache_bin_t *bin, cache_bin_info_t *info, void **ptrs,
+    cache_bin_sz_t nfill, size_t batch) {
+	assert_true(cache_bin_ncached_get(bin, info) == 0, "");
+	CACHE_BIN_PTR_ARRAY_DECLARE(arr, nfill);
+	cache_bin_init_ptr_array_for_fill(bin, info, &arr, nfill);
+	for (cache_bin_sz_t i = 0; i < nfill; i++) {
+		arr.ptr[i] = &ptrs[i];
+	}
+	cache_bin_finish_fill(bin, info, &arr, nfill);
+	assert_true(cache_bin_ncached_get(bin, info) == nfill, "");
+	cache_bin_low_water_set(bin);
+
+	void **out = malloc((batch + 1) * sizeof(void *));
+	size_t n = cache_bin_alloc_batch(bin, batch, out);
+	assert_true(n == ((size_t)nfill < batch ? (size_t)nfill : batch), "");
+	for (cache_bin_sz_t i = 0; i < (cache_bin_sz_t)n; i++) {
+		expect_ptr_eq(out[i], &ptrs[i], "");
+	}
+	expect_true(cache_bin_low_water_get(bin, info) == nfill -
+	    (cache_bin_sz_t)n, "");
+	while (cache_bin_ncached_get(bin, info) > 0) {
+		bool success;
+		cache_bin_alloc(bin, &success);
+	}
+	free(out);
+}
+
 TEST_BEGIN(test_cache_bin) {
 	const int ncached_max = 100;
 	bool success;
@@ -197,6 +225,31 @@ TEST_BEGIN(test_cache_bin) {
 	do_flush_test(&bin, &info, ptrs, ncached_max / 2, ncached_max / 4);
 	do_flush_test(&bin, &info, ptrs, ncached_max / 2, 0);
 
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max, ncached_max);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max, ncached_max * 2);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max, ncached_max / 2);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max, 2);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max, 1);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max, 0);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max / 2,
+	    ncached_max / 2);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max / 2, ncached_max);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max / 2,
+	    ncached_max / 4);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max / 2, 2);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max / 2, 1);
+	do_batch_alloc_test(&bin, &info, ptrs, ncached_max / 2, 0);
+	do_batch_alloc_test(&bin, &info, ptrs, 2, ncached_max);
+	do_batch_alloc_test(&bin, &info, ptrs, 2, 2);
+	do_batch_alloc_test(&bin, &info, ptrs, 2, 1);
+	do_batch_alloc_test(&bin, &info, ptrs, 2, 0);
+	do_batch_alloc_test(&bin, &info, ptrs, 1, 2);
+	do_batch_alloc_test(&bin, &info, ptrs, 1, 1);
+	do_batch_alloc_test(&bin, &info, ptrs, 1, 0);
+	do_batch_alloc_test(&bin, &info, ptrs, 0, 2);
+	do_batch_alloc_test(&bin, &info, ptrs, 0, 1);
+	do_batch_alloc_test(&bin, &info, ptrs, 0, 0);
+
 	free(ptrs);
 }
 TEST_END