From 3f0dc64c6b8c1fd77c819028013dacbc6d2ad6b6 Mon Sep 17 00:00:00 2001
From: Qi Wang <interwq@gwu.edu>
Date: Wed, 11 Apr 2018 15:15:26 -0700
Subject: [PATCH] Allow setting extent hooks on uninitialized auto arenas.

Setting extent hooks can result in initializing an unused auto arena.  This is
useful to install extent hooks on auto arenas from the beginning.
---
 src/ctl.c                 | 45 ++++++++++++++++++++++++++++-----------
 test/integration/extent.c | 38 +++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 12 deletions(-)

diff --git a/src/ctl.c b/src/ctl.c
index 02610cf0..86c2837a 100644
--- a/src/ctl.c
+++ b/src/ctl.c
@@ -2251,20 +2251,41 @@ arena_i_extent_hooks_ctl(tsd_t *tsd, const size_t *mib, size_t miblen,
 
 	malloc_mutex_lock(tsd_tsdn(tsd), &ctl_mtx);
 	MIB_UNSIGNED(arena_ind, 1);
-	if (arena_ind < narenas_total_get() && (arena =
-	    arena_get(tsd_tsdn(tsd), arena_ind, false)) != NULL) {
-		if (newp != NULL) {
-			extent_hooks_t *old_extent_hooks;
-			extent_hooks_t *new_extent_hooks
-			    JEMALLOC_CC_SILENCE_INIT(NULL);
-			WRITE(new_extent_hooks, extent_hooks_t *);
-			old_extent_hooks = extent_hooks_set(tsd, arena,
-			    new_extent_hooks);
+	if (arena_ind < narenas_total_get()) {
+		extent_hooks_t *old_extent_hooks;
+		arena = arena_get(tsd_tsdn(tsd), arena_ind, false);
+		if (arena == NULL) {
+			if (arena_ind >= narenas_auto) {
+				ret = EFAULT;
+				goto label_return;
+			}
+			old_extent_hooks =
+			    (extent_hooks_t *)&extent_hooks_default;
 			READ(old_extent_hooks, extent_hooks_t *);
+			if (newp != NULL) {
+				/* Initialize a new arena as a side effect. */
+				extent_hooks_t *new_extent_hooks
+				    JEMALLOC_CC_SILENCE_INIT(NULL);
+				WRITE(new_extent_hooks, extent_hooks_t *);
+				arena = arena_init(tsd_tsdn(tsd), arena_ind,
+				    new_extent_hooks);
+				if (arena == NULL) {
+					ret = EFAULT;
+					goto label_return;
+				}
+			}
 		} else {
-			extent_hooks_t *old_extent_hooks =
-			    extent_hooks_get(arena);
-			READ(old_extent_hooks, extent_hooks_t *);
+			if (newp != NULL) {
+				extent_hooks_t *new_extent_hooks
+				    JEMALLOC_CC_SILENCE_INIT(NULL);
+				WRITE(new_extent_hooks, extent_hooks_t *);
+				old_extent_hooks = extent_hooks_set(tsd, arena,
+				    new_extent_hooks);
+				READ(old_extent_hooks, extent_hooks_t *);
+			} else {
+				old_extent_hooks = extent_hooks_get(arena);
+				READ(old_extent_hooks, extent_hooks_t *);
+			}
 		}
 	} else {
 		ret = EFAULT;
diff --git a/test/integration/extent.c b/test/integration/extent.c
index c2dc1cb8..b5db0876 100644
--- a/test/integration/extent.c
+++ b/test/integration/extent.c
@@ -98,6 +98,43 @@ test_extent_body(unsigned arena_ind) {
 	dallocx(p, flags);
 }
 
+static void
+test_manual_hook_auto_arena(void) {
+	unsigned narenas;
+	size_t old_size, new_size, sz;
+	size_t hooks_mib[3];
+	size_t hooks_miblen;
+	extent_hooks_t *new_hooks, *old_hooks;
+
+	extent_hooks_prep();
+
+	sz = sizeof(unsigned);
+	/* Get number of auto arenas. */
+	assert_d_eq(mallctl("opt.narenas", (void *)&narenas, &sz, NULL, 0),
+	    0, "Unexpected mallctl() failure");
+	if (narenas == 1) {
+		return;
+	}
+
+	/* Install custom extent hooks on arena 1 (might not be initialized). */
+	hooks_miblen = sizeof(hooks_mib)/sizeof(size_t);
+	assert_d_eq(mallctlnametomib("arena.0.extent_hooks", hooks_mib,
+	    &hooks_miblen), 0, "Unexpected mallctlnametomib() failure");
+	hooks_mib[1] = 1;
+	old_size = sizeof(extent_hooks_t *);
+	new_hooks = &hooks;
+	new_size = sizeof(extent_hooks_t *);
+	assert_d_eq(mallctlbymib(hooks_mib, hooks_miblen, (void *)&old_hooks,
+	    &old_size, (void *)&new_hooks, new_size), 0,
+	    "Unexpected extent_hooks error");
+	static bool auto_arena_created = false;
+	if (old_hooks != &hooks) {
+		assert_b_eq(auto_arena_created, false,
+		    "Expected auto arena 1 created only once.");
+		auto_arena_created = true;
+	}
+}
+
 static void
 test_manual_hook_body(void) {
 	unsigned arena_ind;
@@ -169,6 +206,7 @@ test_manual_hook_body(void) {
 }
 
 TEST_BEGIN(test_extent_manual_hook) {
+	test_manual_hook_auto_arena();
 	test_manual_hook_body();
 
 	/* Test failure paths. */