From a32b7bd5676e669821d15d319f686c3add451f4b Mon Sep 17 00:00:00 2001
From: Latchesar Ionkov <lionkov@lanl.gov>
Date: Tue, 1 May 2018 12:31:09 -0600
Subject: [PATCH] Mallctl: Add arenas.lookup

Implement a new mallctl operation that allows looking up the arena a
region of memory belongs to.
---
 doc/jemalloc.xml.in |  9 +++++++++
 src/ctl.c           | 34 +++++++++++++++++++++++++++++++++-
 test/unit/mallctl.c | 17 +++++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/doc/jemalloc.xml.in b/doc/jemalloc.xml.in
index 3ec43ce1..1e12fd3a 100644
--- a/doc/jemalloc.xml.in
+++ b/doc/jemalloc.xml.in
@@ -2130,6 +2130,15 @@ struct extent_hooks_s {
         and return the new arena index.</para></listitem>
       </varlistentry>
 
+      <varlistentry id="arenas.lookup">
+        <term>
+          <mallctl>arenas.lookup</mallctl>
+          (<type>unsigned</type>, <type>void*</type>)
+          <literal>rw</literal>
+        </term>
+        <listitem><para>Index of the arena to which an allocation belongs to.</para></listitem>
+      </varlistentry>
+
       <varlistentry id="prof.thread_active_init">
         <term>
           <mallctl>prof.thread_active_init</mallctl>
diff --git a/src/ctl.c b/src/ctl.c
index 86c2837a..1e713a3d 100644
--- a/src/ctl.c
+++ b/src/ctl.c
@@ -139,6 +139,7 @@ CTL_PROTO(arenas_nbins)
 CTL_PROTO(arenas_nhbins)
 CTL_PROTO(arenas_nlextents)
 CTL_PROTO(arenas_create)
+CTL_PROTO(arenas_lookup)
 CTL_PROTO(prof_thread_active_init)
 CTL_PROTO(prof_active)
 CTL_PROTO(prof_dump)
@@ -373,7 +374,8 @@ static const ctl_named_node_t arenas_node[] = {
 	{NAME("bin"),		CHILD(indexed, arenas_bin)},
 	{NAME("nlextents"),	CTL(arenas_nlextents)},
 	{NAME("lextent"),	CHILD(indexed, arenas_lextent)},
-	{NAME("create"),	CTL(arenas_create)}
+	{NAME("create"),	CTL(arenas_create)},
+	{NAME("lookup"),	CTL(arenas_lookup)}
 };
 
 static const ctl_named_node_t	prof_node[] = {
@@ -2471,6 +2473,36 @@ label_return:
 	return ret;
 }
 
+static int
+arenas_lookup_ctl(tsd_t *tsd, const size_t *mib, size_t miblen, void *oldp,
+    size_t *oldlenp, void *newp, size_t newlen) {
+	int ret;
+	unsigned arena_ind;
+	void *ptr;
+	extent_t *extent;
+	arena_t *arena;
+
+	ptr = NULL;
+	ret = EINVAL;
+	malloc_mutex_lock(tsd_tsdn(tsd), &ctl_mtx);
+	WRITE(ptr, void *);
+	extent = iealloc(tsd_tsdn(tsd), ptr);
+	if (extent == NULL)
+		goto label_return;
+
+	arena = extent_arena_get(extent);
+	if (arena == NULL)
+		goto label_return;
+
+	arena_ind = arena_ind_get(arena);
+	READ(arena_ind, unsigned);
+
+	ret = 0;
+label_return:
+	malloc_mutex_unlock(tsd_tsdn(tsd), &ctl_mtx);
+	return ret;
+}
+
 /******************************************************************************/
 
 static int
diff --git a/test/unit/mallctl.c b/test/unit/mallctl.c
index f8b180b1..1ecbab08 100644
--- a/test/unit/mallctl.c
+++ b/test/unit/mallctl.c
@@ -738,6 +738,22 @@ TEST_BEGIN(test_arenas_create) {
 }
 TEST_END
 
+TEST_BEGIN(test_arenas_lookup) {
+	unsigned arena, arena1;
+	void *ptr;
+	size_t sz = sizeof(unsigned);
+
+	assert_d_eq(mallctl("arenas.create", (void *)&arena, &sz, NULL, 0), 0,
+	    "Unexpected mallctl() failure");
+	ptr = mallocx(42, MALLOCX_ARENA(arena) | MALLOCX_TCACHE_NONE);
+	assert_ptr_not_null(ptr, "Unexpected mallocx() failure");
+	assert_d_eq(mallctl("arenas.lookup", &arena1, &sz, &ptr, sizeof(ptr)),
+	    0, "Unexpected mallctl() failure");
+	assert_u_eq(arena, arena1, "Unexpected arena index");
+	dallocx(ptr, 0);
+}
+TEST_END
+
 TEST_BEGIN(test_stats_arenas) {
 #define TEST_STATS_ARENAS(t, name) do {					\
 	t name;								\
@@ -784,5 +800,6 @@ main(void) {
 	    test_arenas_bin_constants,
 	    test_arenas_lextent_constants,
 	    test_arenas_create,
+	    test_arenas_lookup,
 	    test_stats_arenas);
 }