From e125455bc33e2985a91ce231620c40c65658a71a Mon Sep 17 00:00:00 2001
From: Fei Deng <duke8253@gmail.com>
Date: Thu, 1 Jul 2021 13:34:27 -0500
Subject: [PATCH] use thread_local to speed things up

---
 include/tscore/JeAllocator.h | 10 +++---
 src/tscore/JeAllocator.cc    | 61 +++++++++++-------------------------
 2 files changed, 23 insertions(+), 48 deletions(-)

diff --git a/include/tscore/JeAllocator.h b/include/tscore/JeAllocator.h
index a1dc942b4c0..c83ba0ba029 100644
--- a/include/tscore/JeAllocator.h
+++ b/include/tscore/JeAllocator.h
@@ -70,15 +70,13 @@ class JemallocNodumpAllocator
 
 private:
 #if JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
-  static extent_alloc_t *original_alloc;
-  static extent_hooks_t extent_hooks;
-  static std::shared_mutex je_mutex;
-  static std::unordered_map<pthread_t, unsigned int> arenas;
-  static std::unordered_map<unsigned int, int> arena_flags;
+  thread_local static extent_alloc_t *original_alloc;
+  thread_local static extent_hooks_t extent_hooks;
+  thread_local static int arena_flags;
 
   static void *alloc(extent_hooks_t *extent, void *new_addr, size_t size, size_t alignment, bool *zero, bool *commit,
                      unsigned int arena_id);
-  unsigned int extend_and_setup_arena(pthread_t thread_id);
+  int extend_and_setup_arena();
 #endif /* JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED */
 };
 
diff --git a/src/tscore/JeAllocator.cc b/src/tscore/JeAllocator.cc
index f518f106c76..07ee5cfe1b0 100644
--- a/src/tscore/JeAllocator.cc
+++ b/src/tscore/JeAllocator.cc
@@ -29,11 +29,9 @@
 namespace jearena
 {
 #if JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
-extent_alloc_t *JemallocNodumpAllocator::original_alloc = nullptr;
-extent_hooks_t JemallocNodumpAllocator::extent_hooks;
-std::shared_mutex JemallocNodumpAllocator::je_mutex;
-std::unordered_map<pthread_t, unsigned int> JemallocNodumpAllocator::arenas;
-std::unordered_map<unsigned int, int> JemallocNodumpAllocator::arena_flags;
+thread_local extent_alloc_t *JemallocNodumpAllocator::original_alloc = nullptr;
+thread_local extent_hooks_t JemallocNodumpAllocator::extent_hooks;
+thread_local int JemallocNodumpAllocator::arena_flags = 0;
 
 void *
 JemallocNodumpAllocator::alloc(extent_hooks_t *extent, void *new_addr, size_t size, size_t alignment, bool *zero, bool *commit,
@@ -50,8 +48,8 @@ JemallocNodumpAllocator::alloc(extent_hooks_t *extent, void *new_addr, size_t si
   return result;
 }
 
-unsigned int
-JemallocNodumpAllocator::extend_and_setup_arena(pthread_t thread_id)
+int
+JemallocNodumpAllocator::extend_and_setup_arena()
 {
   unsigned int arena_id;
   size_t arena_id_len = sizeof(arena_id);
@@ -69,31 +67,21 @@ JemallocNodumpAllocator::extend_and_setup_arena(pthread_t thread_id)
     ink_abort("Unable to get the hooks: %s", std::strerror(ret));
   }
 
-  {
-    std::unique_lock lock(je_mutex);
-    if (original_alloc == nullptr) {
-      original_alloc = hooks->alloc;
-    } else {
-      ink_release_assert(original_alloc == hooks->alloc);
-    }
+  // Set the custom hook
+  original_alloc     = hooks->alloc;
+  extent_hooks       = *hooks;
+  extent_hooks.alloc = &JemallocNodumpAllocator::alloc;
 
-    // Set the custom hook
-    if (extent_hooks.alloc != &JemallocNodumpAllocator::alloc) {
-      extent_hooks       = *hooks;
-      extent_hooks.alloc = &JemallocNodumpAllocator::alloc;
-    }
-    extent_hooks_t *new_hooks = &extent_hooks;
-    if (auto ret = mallctl(key.c_str(), nullptr, nullptr, &new_hooks, sizeof(new_hooks))) {
-      ink_abort("Unable to set the hooks: %s", std::strerror(ret));
-    }
+  extent_hooks_t *new_hooks = &extent_hooks;
+  if (auto ret = mallctl(key.c_str(), nullptr, nullptr, &new_hooks, sizeof(new_hooks))) {
+    ink_abort("Unable to set the hooks: %s", std::strerror(ret));
+  }
 
-    Debug("JeAllocator", "arena \"%ud\" created", arena_id);
+  Debug("JeAllocator", "arena \"%ud\" created with flags \"%d\"", arena_id, flags);
 
-    arenas[thread_id]     = arena_id;
-    arena_flags[arena_id] = flags;
-  }
+  arena_flags = flags;
 
-  return arena_id;
+  return flags;
 }
 #endif /* JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED */
 
@@ -105,20 +93,9 @@ void *
 JemallocNodumpAllocator::allocate(InkFreeList *f)
 {
 #if JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
-  int flags           = 0;
-  pthread_t thread_id = pthread_self();
-  {
-    std::shared_lock lock(je_mutex);
-    unsigned int arena_id = 0;
-    auto arena            = arenas.find(thread_id);
-    if (unlikely(arena == arenas.end())) {
-      lock.unlock();
-      arena_id = extend_and_setup_arena(thread_id);
-      lock.lock();
-    } else {
-      arena_id = arena->second;
-    }
-    flags = arena_flags.find(arena_id)->second;
+  int flags = arena_flags;
+  if (unlikely(flags == 0)) {
+    flags = extend_and_setup_arena();
   }
 #endif