From 59e371f46331a3f4b688d6622a0af7ccc4f96be6 Mon Sep 17 00:00:00 2001
From: David Goldblatt <davidgoldblatt@fb.com>
Date: Mon, 30 Apr 2018 16:24:36 -0700
Subject: [PATCH] Hooks: Add a hook exhaustion test.

When we run out of space in which to store hooks, we should return EAGAIN from
the mallctl, but not otherwise misbehave.
---
 include/jemalloc/internal/hook.h |  2 ++
 src/hook.c                       |  9 ++++---
 test/unit/mallctl.c              | 40 +++++++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 6 deletions(-)

diff --git a/include/jemalloc/internal/hook.h b/include/jemalloc/internal/hook.h
index 678c6018..9ea9c6f0 100644
--- a/include/jemalloc/internal/hook.h
+++ b/include/jemalloc/internal/hook.h
@@ -46,6 +46,8 @@
  *   order.
  */
 
+#define HOOK_MAX 4
+
 enum hook_alloc_e {
 	hook_alloc_malloc,
 	hook_alloc_posix_memalign,
diff --git a/src/hook.c b/src/hook.c
index 564c2a0c..24afe999 100644
--- a/src/hook.c
+++ b/src/hook.c
@@ -14,9 +14,8 @@ struct hooks_internal_s {
 
 seq_define(hooks_internal_t, hooks)
 
-#define HOOKS_MAX 4
 static atomic_u_t nhooks = ATOMIC_INIT(0);
-static seq_hooks_t hooks[HOOKS_MAX];
+static seq_hooks_t hooks[HOOK_MAX];
 static malloc_mutex_t hooks_mu;
 
 bool
@@ -28,7 +27,7 @@ hook_boot() {
 static void *
 hook_install_locked(hooks_t *to_install) {
 	hooks_internal_t hooks_internal;
-	for (int i = 0; i < HOOKS_MAX; i++) {
+	for (int i = 0; i < HOOK_MAX; i++) {
 		bool success = seq_try_load_hooks(&hooks_internal, &hooks[i]);
 		/* We hold mu; no concurrent access. */
 		assert(success);
@@ -74,7 +73,7 @@ void
 hook_remove(tsdn_t *tsdn, void *opaque) {
 	if (config_debug) {
 		char *hooks_begin = (char *)&hooks[0];
-		char *hooks_end = (char *)&hooks[HOOKS_MAX];
+		char *hooks_end = (char *)&hooks[HOOK_MAX];
 		char *hook = (char *)opaque;
 		assert(hooks_begin <= hook && hook < hooks_end
 		    && (hook - hooks_begin) % sizeof(seq_hooks_t) == 0);
@@ -87,7 +86,7 @@ hook_remove(tsdn_t *tsdn, void *opaque) {
 
 #define FOR_EACH_HOOK_BEGIN(hooks_internal_ptr)				\
 for (int for_each_hook_counter = 0;					\
-    for_each_hook_counter < HOOKS_MAX;					\
+    for_each_hook_counter < HOOK_MAX;					\
     for_each_hook_counter++) {						\
 	bool for_each_hook_success = seq_try_load_hooks(		\
 	    (hooks_internal_ptr), &hooks[for_each_hook_counter]);	\
diff --git a/test/unit/mallctl.c b/test/unit/mallctl.c
index 34a4d67c..8a36c0a4 100644
--- a/test/unit/mallctl.c
+++ b/test/unit/mallctl.c
@@ -1,5 +1,6 @@
 #include "test/jemalloc_test.h"
 
+#include "jemalloc/internal/hook.h"
 #include "jemalloc/internal/util.h"
 
 TEST_BEGIN(test_mallctl_errors) {
@@ -810,6 +811,42 @@ TEST_BEGIN(test_hooks) {
 }
 TEST_END
 
+TEST_BEGIN(test_hooks_exhaustion) {
+	bool hook_called = false;
+	hooks_t hooks = {&alloc_hook, &dalloc_hook, NULL, &hook_called};
+
+	void *handle;
+	void *handles[HOOK_MAX];
+	size_t sz = sizeof(handle);
+	int err;
+	for (int i = 0; i < HOOK_MAX; i++) {
+		handle = NULL;
+		err = mallctl("experimental.hooks.install", &handle, &sz,
+		    &hooks, sizeof(hooks));
+		assert_d_eq(err, 0, "Error installation hooks");
+		assert_ptr_ne(handle, NULL, "Got NULL handle");
+		handles[i] = handle;
+	}
+	err = mallctl("experimental.hooks.install", &handle, &sz, &hooks,
+	    sizeof(hooks));
+	assert_d_eq(err, EAGAIN, "Should have failed hook installation");
+	for (int i = 0; i < HOOK_MAX; i++) {
+		err = mallctl("experimental.hooks.remove", NULL, NULL,
+		    &handles[i], sizeof(handles[i]));
+		assert_d_eq(err, 0, "Hook removal failed");
+	}
+	/* Insertion failed, but then we removed some; it should work now. */
+	handle = NULL;
+	err = mallctl("experimental.hooks.install", &handle, &sz, &hooks,
+	    sizeof(hooks));
+	assert_d_eq(err, 0, "Hook insertion failed");
+	assert_ptr_ne(handle, NULL, "Got NULL handle");
+	err = mallctl("experimental.hooks.remove", NULL, NULL, &handle,
+	    sizeof(handle));
+	assert_d_eq(err, 0, "Hook removal failed");
+}
+TEST_END
+
 int
 main(void) {
 	return test(
@@ -839,5 +876,6 @@ main(void) {
 	    test_arenas_create,
 	    test_arenas_lookup,
 	    test_stats_arenas,
-	    test_hooks);
+	    test_hooks,
+	    test_hooks_exhaustion);
 }