From 8b24cb8fdf2bf210e243c1d676484a4ffa5c3f6c Mon Sep 17 00:00:00 2001
From: Qi Wang <interwq@gwu.edu>
Date: Mon, 23 Aug 2021 15:55:54 -0700
Subject: [PATCH] Don't assume initialized arena in the default alloc hook.

Specifically, this change allows the default alloc hook to used during
arenas.create.  One use case is to invoke the default alloc hook in a customized
hook arena, i.e. the default hooks can be read out of a default arena, then
create customized ones based on these hooks.  Note that mixing the default with
customized hooks is not recommended, and should only be considered when the
customization is simple and straightforward.
---
 src/ehooks.c            | 25 ++++++++-----------------
 test/unit/arena_reset.c | 15 +++++++++++++++
 2 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/src/ehooks.c b/src/ehooks.c
index 535066e7..5d12d003 100644
--- a/src/ehooks.c
+++ b/src/ehooks.c
@@ -52,9 +52,12 @@ void *
 ehooks_default_alloc_impl(tsdn_t *tsdn, void *new_addr, size_t size,
     size_t alignment, bool *zero, bool *commit, unsigned arena_ind) {
 	arena_t *arena = arena_get(tsdn, arena_ind, false);
-	void *ret = extent_alloc_core(tsdn, arena, new_addr, size, alignment, zero,
-	    commit, (dss_prec_t)atomic_load_u(&arena->dss_prec,
-	    ATOMIC_RELAXED));
+	/* NULL arena indicates arena_create. */
+	assert(arena != NULL || alignment == HUGEPAGE);
+	dss_prec_t dss = (arena == NULL) ? dss_prec_disabled :
+	    (dss_prec_t)atomic_load_u(&arena->dss_prec, ATOMIC_RELAXED);
+	void *ret = extent_alloc_core(tsdn, arena, new_addr, size, alignment,
+	    zero, commit, dss);
 	if (have_madvise_huge && ret) {
 		pages_set_thp_state(ret, size);
 	}
@@ -64,20 +67,8 @@ ehooks_default_alloc_impl(tsdn_t *tsdn, void *new_addr, size_t size,
 static void *
 ehooks_default_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size,
     size_t alignment, bool *zero, bool *commit, unsigned arena_ind) {
-	tsdn_t *tsdn;
-	arena_t *arena;
-
-	tsdn = tsdn_fetch();
-	arena = arena_get(tsdn, arena_ind, false);
-	/*
-	 * The arena we're allocating on behalf of must have been initialized
-	 * already.
-	 */
-	assert(arena != NULL);
-
-	return ehooks_default_alloc_impl(tsdn, new_addr, size,
-	    ALIGNMENT_CEILING(alignment, PAGE), zero, commit,
-	    arena_ind_get(arena));
+	return ehooks_default_alloc_impl(tsdn_fetch(), new_addr, size,
+	    ALIGNMENT_CEILING(alignment, PAGE), zero, commit, arena_ind);
 }
 
 bool
diff --git a/test/unit/arena_reset.c b/test/unit/arena_reset.c
index a2cf3e54..589689c0 100644
--- a/test/unit/arena_reset.c
+++ b/test/unit/arena_reset.c
@@ -255,6 +255,21 @@ TEST_BEGIN(test_arena_destroy_hooks_default) {
 	do_arena_reset_post(ptrs, nptrs, arena_ind);
 
 	do_arena_destroy(arena_ind_another);
+
+	/* Try arena.create with custom hooks. */
+	size_t sz = sizeof(extent_hooks_t *);
+	extent_hooks_t *default_hooks;
+	expect_d_eq(mallctl("arena.0.extent_hooks", (void *)&default_hooks,
+	    &sz, NULL, 0), 0, "Unexpected mallctlnametomib() failure");
+
+	/* Default impl; but wrapped as "customized". */
+	extent_hooks_t new_hooks = *default_hooks;
+	extent_hooks_t *hook = &new_hooks;
+	sz = sizeof(unsigned);
+	expect_d_eq(mallctl("arenas.create", (void *)&arena_ind, &sz,
+	    (void *)&hook, sizeof(void *)), 0,
+	    "Unexpected mallctl() failure");
+	do_arena_destroy(arena_ind);
 }
 TEST_END