From cf9724531af2864b243668d82aa63114e9737bfd Mon Sep 17 00:00:00 2001
From: David CARLIER <devnexen@gmail.com>
Date: Wed, 17 Feb 2021 20:40:11 +0000
Subject: [PATCH] Darwin malloc_size override support proposal.

Darwin has similar api than Linux/FreeBSD's malloc_usable_size.
---
 configure.ac                                  |  3 ++
 .../internal/jemalloc_internal_defs.h.in      |  5 +++
 include/jemalloc/jemalloc_protos.h.in         |  4 +++
 src/jemalloc.c                                | 35 ++++++++++++++-----
 test/include/test/jemalloc_test.h.in          |  5 +++
 test/integration/aligned_alloc.c              |  4 +--
 test/integration/allocated.c                  |  2 +-
 test/integration/malloc.c                     |  2 +-
 test/integration/posix_memalign.c             |  2 +-
 test/integration/rallocx.c                    |  4 +--
 test/stress/microbench.c                      |  2 +-
 test/unit/junk.c                              |  2 +-
 test/unit/prof_stats.c                        |  2 +-
 13 files changed, 53 insertions(+), 19 deletions(-)

diff --git a/configure.ac b/configure.ac
index 7e2b44c5..7a49e84f 100644
--- a/configure.ac
+++ b/configure.ac
@@ -1056,6 +1056,9 @@ AC_CHECK_FUNC([memalign],
 AC_CHECK_FUNC([valloc],
 	      [AC_DEFINE([JEMALLOC_OVERRIDE_VALLOC], [ ])
 	       public_syms="${public_syms} valloc"])
+AC_CHECK_FUNC([malloc_size],
+	      [AC_DEFINE([JEMALLOC_HAVE_MALLOC_SIZE], [ ])
+	       public_syms="${public_syms} malloc_size"])
 
 dnl Check for allocator-related functions that should be wrapped.
 wrap_syms=
diff --git a/include/jemalloc/internal/jemalloc_internal_defs.h.in b/include/jemalloc/internal/jemalloc_internal_defs.h.in
index 418b0cb2..a4be549b 100644
--- a/include/jemalloc/internal/jemalloc_internal_defs.h.in
+++ b/include/jemalloc/internal/jemalloc_internal_defs.h.in
@@ -337,6 +337,11 @@
  */
 #undef JEMALLOC_HAVE_MEMCNTL
 
+/*
+ * Defined if malloc_size is supported
+ */
+#undef JEMALLOC_HAVE_MALLOC_SIZE
+
 /* Define if operating system has alloca.h header. */
 #undef JEMALLOC_HAS_ALLOCA_H
 
diff --git a/include/jemalloc/jemalloc_protos.h.in b/include/jemalloc/jemalloc_protos.h.in
index d75b2224..356221cc 100644
--- a/include/jemalloc/jemalloc_protos.h.in
+++ b/include/jemalloc/jemalloc_protos.h.in
@@ -53,6 +53,10 @@ JEMALLOC_EXPORT void JEMALLOC_NOTHROW	@je_@malloc_stats_print(
     const char *opts);
 JEMALLOC_EXPORT size_t JEMALLOC_NOTHROW	@je_@malloc_usable_size(
     JEMALLOC_USABLE_SIZE_CONST void *ptr) JEMALLOC_CXX_THROW;
+#ifdef JEMALLOC_HAVE_MALLOC_SIZE
+JEMALLOC_EXPORT size_t JEMALLOC_NOTHROW	@je_@malloc_size(
+    const void *ptr);
+#endif
 
 #ifdef JEMALLOC_OVERRIDE_MEMALIGN
 JEMALLOC_EXPORT JEMALLOC_ALLOCATOR JEMALLOC_RESTRICT_RETURN
diff --git a/src/jemalloc.c b/src/jemalloc.c
index 8e04fa6f..469a4910 100644
--- a/src/jemalloc.c
+++ b/src/jemalloc.c
@@ -3904,18 +3904,14 @@ je_malloc_stats_print(void (*write_cb)(void *, const char *), void *cbopaque,
 }
 #undef STATS_PRINT_BUFSIZE
 
-JEMALLOC_EXPORT size_t JEMALLOC_NOTHROW
-je_malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void *ptr) {
-	size_t ret;
-	tsdn_t *tsdn;
-
-	LOG("core.malloc_usable_size.entry", "ptr: %p", ptr);
-
+JEMALLOC_ALWAYS_INLINE size_t
+je_malloc_usable_size_impl(JEMALLOC_USABLE_SIZE_CONST void *ptr) {
 	assert(malloc_initialized() || IS_INITIALIZER);
 
-	tsdn = tsdn_fetch();
+	tsdn_t *tsdn = tsdn_fetch();
 	check_entry_exit_locking(tsdn);
 
+	size_t ret;
 	if (unlikely(ptr == NULL)) {
 		ret = 0;
 	} else {
@@ -3926,12 +3922,33 @@ je_malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void *ptr) {
 			ret = isalloc(tsdn, ptr);
 		}
 	}
-
 	check_entry_exit_locking(tsdn);
+
+	return ret;
+}
+
+JEMALLOC_EXPORT size_t JEMALLOC_NOTHROW
+je_malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void *ptr) {
+	LOG("core.malloc_usable_size.entry", "ptr: %p", ptr);
+
+	size_t ret = je_malloc_usable_size_impl(ptr);
+
 	LOG("core.malloc_usable_size.exit", "result: %zu", ret);
 	return ret;
 }
 
+#ifdef JEMALLOC_HAVE_MALLOC_SIZE
+JEMALLOC_EXPORT size_t JEMALLOC_NOTHROW
+je_malloc_size(const void *ptr) {
+	LOG("core.malloc_size.entry", "ptr: %p", ptr);
+
+	size_t ret = je_malloc_usable_size_impl(ptr);
+
+	LOG("core.malloc_size.exit", "result: %zu", ret);
+	return ret;
+}
+#endif
+
 static void
 batch_alloc_prof_sample_assert(tsd_t *tsd, size_t batch, size_t usize) {
 	assert(config_prof && opt_prof);
diff --git a/test/include/test/jemalloc_test.h.in b/test/include/test/jemalloc_test.h.in
index 0e332165..3f8c0da7 100644
--- a/test/include/test/jemalloc_test.h.in
+++ b/test/include/test/jemalloc_test.h.in
@@ -132,6 +132,11 @@ static const bool config_debug =
 #define MEXP 19937
 #include "test/SFMT.h"
 
+#ifndef JEMALLOC_HAVE_MALLOC_SIZE
+#define TEST_MALLOC_SIZE malloc_usable_size
+#else
+#define TEST_MALLOC_SIZE malloc_size
+#endif
 /******************************************************************************/
 /*
  * Define always-enabled assertion macros, so that test assertions execute even
diff --git a/test/integration/aligned_alloc.c b/test/integration/aligned_alloc.c
index 3f619e7e..b37d5ba0 100644
--- a/test/integration/aligned_alloc.c
+++ b/test/integration/aligned_alloc.c
@@ -120,7 +120,7 @@ TEST_BEGIN(test_alignment_and_size) {
 					    "size=%zu (%#zx): %s",
 					    alignment, size, size, buf);
 				}
-				total += malloc_usable_size(ps[i]);
+				total += TEST_MALLOC_SIZE(ps[i]);
 				if (total >= (MAXALIGN << 1)) {
 					break;
 				}
@@ -141,7 +141,7 @@ TEST_END
 TEST_BEGIN(test_zero_alloc) {
 	void *res = aligned_alloc(8, 0);
 	assert(res);
-	size_t usable = malloc_usable_size(res);
+	size_t usable = TEST_MALLOC_SIZE(res);
 	assert(usable > 0);
 	free(res);
 }
diff --git a/test/integration/allocated.c b/test/integration/allocated.c
index 8f2f21d5..0c64272c 100644
--- a/test/integration/allocated.c
+++ b/test/integration/allocated.c
@@ -70,7 +70,7 @@ thd_start(void *arg) {
 	expect_ptr_eq(ap0, ap1,
 	    "Pointer returned by \"thread.allocatedp\" should not change");
 
-	usize = malloc_usable_size(p);
+	usize = TEST_MALLOC_SIZE(p);
 	expect_u64_le(a0 + usize, a1,
 	    "Allocated memory counter should increase by at least the amount "
 	    "explicitly allocated");
diff --git a/test/integration/malloc.c b/test/integration/malloc.c
index 8b33bc8f..ef449163 100644
--- a/test/integration/malloc.c
+++ b/test/integration/malloc.c
@@ -3,7 +3,7 @@
 TEST_BEGIN(test_zero_alloc) {
 	void *res = malloc(0);
 	assert(res);
-	size_t usable = malloc_usable_size(res);
+	size_t usable = TEST_MALLOC_SIZE(res);
 	assert(usable > 0);
 	free(res);
 }
diff --git a/test/integration/posix_memalign.c b/test/integration/posix_memalign.c
index 6f8a1b03..2da0549b 100644
--- a/test/integration/posix_memalign.c
+++ b/test/integration/posix_memalign.c
@@ -101,7 +101,7 @@ TEST_BEGIN(test_alignment_and_size) {
 					    "size=%zu (%#zx): %s",
 					    alignment, size, size, buf);
 				}
-				total += malloc_usable_size(ps[i]);
+				total += TEST_MALLOC_SIZE(ps[i]);
 				if (total >= (MAXALIGN << 1)) {
 					break;
 				}
diff --git a/test/integration/rallocx.c b/test/integration/rallocx.c
index 57c7967f..d4a48fce 100644
--- a/test/integration/rallocx.c
+++ b/test/integration/rallocx.c
@@ -185,7 +185,7 @@ TEST_BEGIN(test_align_enum) {
 				assert_ptr_not_null(p,
 				    "Unexpected mallocx() error");
 				assert_zu_eq(nallocx(1, flags),
-				    malloc_usable_size(p),
+				    TEST_MALLOC_SIZE(p),
 				    "Wrong mallocx() usable size");
 				int flags_next =
 				    MALLOCX_LG_ALIGN(lg_align_next);
@@ -193,7 +193,7 @@ TEST_BEGIN(test_align_enum) {
 				assert_ptr_not_null(p,
 				    "Unexpected rallocx() error");
 				expect_zu_eq(nallocx(size, flags_next),
-				    malloc_usable_size(p),
+				    TEST_MALLOC_SIZE(p),
 				    "Wrong rallocx() usable size");
 				free(p);
 			}
diff --git a/test/stress/microbench.c b/test/stress/microbench.c
index 226677f7..062e32fd 100644
--- a/test/stress/microbench.c
+++ b/test/stress/microbench.c
@@ -69,7 +69,7 @@ malloc_mus_free(void) {
 		test_fail("Unexpected malloc() failure");
 		return;
 	}
-	malloc_usable_size(p);
+	TEST_MALLOC_SIZE(p);
 	free(p);
 }
 
diff --git a/test/unit/junk.c b/test/unit/junk.c
index 314da3ce..543092f1 100644
--- a/test/unit/junk.c
+++ b/test/unit/junk.c
@@ -30,7 +30,7 @@ do_allocs(size_t size, bool zero, size_t lg_align) {
 		if (opt_junk_alloc && !zero) {				\
 			expect_ptr_eq(ptr, last_junked_ptr, "");	\
 			expect_zu_eq(last_junked_usize,			\
-			    malloc_usable_size(ptr), "");		\
+			    TEST_MALLOC_SIZE(ptr), "");			\
 		}							\
 	} while (0)
 	if (!zero && lg_align == 0) {
diff --git a/test/unit/prof_stats.c b/test/unit/prof_stats.c
index a9145871..c88c4ae0 100644
--- a/test/unit/prof_stats.c
+++ b/test/unit/prof_stats.c
@@ -43,7 +43,7 @@ test_combinations(szind_t ind, size_t sizes_array[N_PTRS],
 		int flags = flags_array[i];
 		void *p = mallocx(sz, flags);
 		assert_ptr_not_null(p, "malloc() failed");
-		assert(malloc_usable_size(p) == sz_index2size(ind));
+		assert(TEST_MALLOC_SIZE(p) == sz_index2size(ind));
 		ptrs[i] = p;
 		live_req_sum += sz;
 		live_count++;