JemallocAllocator: thread-local tcache (#4603)

Summary:
Add option to support  thread-local tcache to reduce mutex contention inside Jemalloc arena.
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4603

Differential Revision: D12830738

Pulled By: yiwu-arbug

fbshipit-source-id: 59bd25b165b903f23a6a8531b18d72e140d69f65

util/jemalloc_nodump_allocator.cc

@@ -6,7 +6,10 @@
 #include "util/jemalloc_nodump_allocator.h"
 
 #include <string>
+#include <thread>
 
+#include "port/likely.h"
+#include "port/port.h"
 #include "util/string_util.h"
 
 namespace rocksdb {
@@ -16,16 +19,44 @@ namespace rocksdb {
 std::atomic<extent_alloc_t*> JemallocNodumpAllocator::original_alloc_{nullptr};
 
 JemallocNodumpAllocator::JemallocNodumpAllocator(
-    unsigned arena_index, int flags, std::unique_ptr<extent_hooks_t>&& hooks)
-    : arena_index_(arena_index), flags_(flags), hooks_(std::move(hooks)) {
-  assert(arena_index != 0);
+    std::unique_ptr<extent_hooks_t>&& arena_hooks, unsigned arena_index)
+    : arena_hooks_(std::move(arena_hooks)),
+      arena_index_(arena_index),
+      tcache_(&JemallocNodumpAllocator::DestroyThreadSpecificCache) {}
+
+int JemallocNodumpAllocator::GetThreadSpecificCache() {
+  // We always enable tcache. The only corner case is when there are a ton of
+  // threads accessing with low frequency, then it could consume a lot of
+  // memory (may reach # threads * ~1MB) without bringing too much benefit.
+  unsigned* tcache_index = reinterpret_cast<unsigned*>(tcache_.Get());
+  if (UNLIKELY(tcache_index == nullptr)) {
+    // Instantiate tcache.
+    tcache_index = new unsigned(0);
+    size_t tcache_index_size = sizeof(unsigned);
+    int ret =
+        mallctl("tcache.create", tcache_index, &tcache_index_size, nullptr, 0);
+    if (ret != 0) {
+      // No good way to expose the error. Silently disable tcache.
+      delete tcache_index;
+      return MALLOCX_TCACHE_NONE;
+    }
+    tcache_.Reset(static_cast<void*>(tcache_index));
+  }
+  return MALLOCX_TCACHE(*tcache_index);
 }
 
 void* JemallocNodumpAllocator::Allocate(size_t size) {
-  return mallocx(size, flags_);
+  int tcache_flag = GetThreadSpecificCache();
+  return mallocx(size, MALLOCX_ARENA(arena_index_) | tcache_flag);
 }
 
-void JemallocNodumpAllocator::Deallocate(void* p) { dallocx(p, flags_); }
+void JemallocNodumpAllocator::Deallocate(void* p) {
+  // Obtain tcache.
+  int tcache_flag = GetThreadSpecificCache();
+  // No need to pass arena index to dallocx(). Jemalloc will find arena index
+  // from its own metadata.
+  dallocx(p, tcache_flag);
+}
 
 void* JemallocNodumpAllocator::Alloc(extent_hooks_t* extent, void* new_addr,
                                      size_t size, size_t alignment, bool* zero,
@@ -48,13 +79,38 @@ void* JemallocNodumpAllocator::Alloc(extent_hooks_t* extent, void* new_addr,
   return result;
 }
 
-JemallocNodumpAllocator::~JemallocNodumpAllocator() {
-  assert(arena_index_ != 0);
-  std::string key = "arena." + ToString(arena_index_) + ".destroy";
+Status JemallocNodumpAllocator::DestroyArena(unsigned arena_index) {
+  assert(arena_index != 0);
+  std::string key = "arena." + ToString(arena_index) + ".destroy";
   int ret = mallctl(key.c_str(), nullptr, 0, nullptr, 0);
   if (ret != 0) {
-    fprintf(stderr, "Failed to destroy jemalloc arena, error code: %d\n", ret);
+    return Status::Incomplete("Failed to destroy jemalloc arena, error code: " +
+                              ToString(ret));
   }
+  return Status::OK();
+}
+
+void JemallocNodumpAllocator::DestroyThreadSpecificCache(void* ptr) {
+  assert(ptr != nullptr);
+  unsigned* tcache_index = static_cast<unsigned*>(ptr);
+  size_t tcache_index_size = sizeof(unsigned);
+  int ret __attribute__((__unused__)) =
+      mallctl("tcache.destroy", nullptr, 0, tcache_index, tcache_index_size);
+  // Silently ignore error.
+  assert(ret == 0);
+  delete tcache_index;
+}
+
+JemallocNodumpAllocator::~JemallocNodumpAllocator() {
+  // Destroy tcache before destroying arena.
+  autovector<void*> tcache_list;
+  tcache_.Scrape(&tcache_list, nullptr);
+  for (void* tcache_index : tcache_list) {
+    DestroyThreadSpecificCache(tcache_index);
+  }
+  // Destroy arena. Silently ignore error.
+  Status s __attribute__((__unused__)) = DestroyArena(arena_index_);
+  assert(s.ok());
 }
 
 size_t JemallocNodumpAllocator::UsableSize(void* p,
@@ -65,8 +121,8 @@ size_t JemallocNodumpAllocator::UsableSize(void* p,
 
 Status NewJemallocNodumpAllocator(
     std::shared_ptr<MemoryAllocator>* memory_allocator) {
-#ifndef ROCKSDB_JEMALLOC_NODUMP_ALLOCATOR
   *memory_allocator = nullptr;
+#ifndef ROCKSDB_JEMALLOC_NODUMP_ALLOCATOR
   return Status::NotSupported(
       "JemallocNodumpAllocator only available with jemalloc version >= 5 "
       "and MADV_DONTDUMP is available.");
@@ -74,6 +130,7 @@ Status NewJemallocNodumpAllocator(
   if (memory_allocator == nullptr) {
     return Status::InvalidArgument("memory_allocator must be non-null.");
   }
+
   // Create arena.
   unsigned arena_index = 0;
   size_t arena_index_size = sizeof(arena_index);
@@ -84,16 +141,14 @@ Status NewJemallocNodumpAllocator(
                               ToString(ret));
   }
   assert(arena_index != 0);
-  int flags = MALLOCX_ARENA(arena_index) | MALLOCX_TCACHE_NONE;
-  std::string key = "arena." + ToString(arena_index) + ".extent_hooks";
 
   // Read existing hooks.
+  std::string key = "arena." + ToString(arena_index) + ".extent_hooks";
   extent_hooks_t* hooks;
   size_t hooks_size = sizeof(hooks);
   ret = mallctl(key.c_str(), &hooks, &hooks_size, nullptr, 0);
   if (ret != 0) {
-    std::string msg =
-        "Failed to read existing hooks, error code: " + ToString(ret);
+    JemallocNodumpAllocator::DestroyArena(arena_index);
     return Status::Incomplete("Failed to read existing hooks, error code: " +
                               ToString(ret));
   }
@@ -101,10 +156,13 @@ Status NewJemallocNodumpAllocator(
   // Store existing alloc.
   extent_alloc_t* original_alloc = hooks->alloc;
   extent_alloc_t* expected = nullptr;
-  bool success __attribute__((__unused__)) =
+  bool success =
       JemallocNodumpAllocator::original_alloc_.compare_exchange_strong(
           expected, original_alloc);
-  assert(success || original_alloc == expected);
+  if (!success && original_alloc != expected) {
+    JemallocNodumpAllocator::DestroyArena(arena_index);
+    return Status::Incomplete("Original alloc conflict.");
+  }
 
   // Set the custom hook.
   std::unique_ptr<extent_hooks_t> new_hooks(new extent_hooks_t(*hooks));
@@ -112,13 +170,14 @@ Status NewJemallocNodumpAllocator(
   extent_hooks_t* hooks_ptr = new_hooks.get();
   ret = mallctl(key.c_str(), nullptr, nullptr, &hooks_ptr, sizeof(hooks_ptr));
   if (ret != 0) {
+    JemallocNodumpAllocator::DestroyArena(arena_index);
     return Status::Incomplete("Failed to set custom hook, error code: " +
                               ToString(ret));
   }
 
   // Create cache allocator.
   memory_allocator->reset(
-      new JemallocNodumpAllocator(arena_index, flags, std::move(new_hooks)));
+      new JemallocNodumpAllocator(std::move(new_hooks), arena_index));
   return Status::OK();
 #endif  // ROCKSDB_JEMALLOC_NODUMP_ALLOCATOR
 }

util/jemalloc_nodump_allocator.h

@@ -6,8 +6,12 @@
 #pragma once
 
 #include <atomic>
+#include <vector>
 
+#include "port/port.h"
 #include "rocksdb/memory_allocator.h"
+#include "util/core_local.h"
+#include "util/thread_local.h"
 
 #if defined(ROCKSDB_JEMALLOC) && defined(ROCKSDB_PLATFORM_POSIX)
 
@@ -21,8 +25,8 @@ namespace rocksdb {
 
 class JemallocNodumpAllocator : public MemoryAllocator {
  public:
-  JemallocNodumpAllocator(unsigned arena_index, int flags,
-                          std::unique_ptr<extent_hooks_t>&& hooks);
+  JemallocNodumpAllocator(std::unique_ptr<extent_hooks_t>&& arena_hooks,
+                          unsigned arena_index);
   ~JemallocNodumpAllocator();
 
   const char* Name() const override { return "JemallocNodumpAllocator"; }
@@ -39,6 +43,16 @@ class JemallocNodumpAllocator : public MemoryAllocator {
                      size_t alignment, bool* zero, bool* commit,
                      unsigned arena_ind);
 
+  // Destroy arena on destruction of the allocator, or on failure.
+  static Status DestroyArena(unsigned arena_index);
+
+  // Destroy tcache on destruction of the allocator, or thread exit.
+  static void DestroyThreadSpecificCache(void* ptr);
+
+  // Get or create tcache. Return flag suitable to use with `mallocx`:
+  // either MALLOCX_TCACHE_NONE or MALLOCX_TCACHE(tc).
+  int GetThreadSpecificCache();
+
   // A function pointer to jemalloc default alloc. Use atomic to make sure
   // NewJemallocNodumpAllocator is thread-safe.
   //
@@ -46,9 +60,14 @@ class JemallocNodumpAllocator : public MemoryAllocator {
   // alloc needs to be static to pass to jemalloc as function pointer.
   static std::atomic<extent_alloc_t*> original_alloc_;
 
-  unsigned arena_index_;
-  int flags_;
-  const std::unique_ptr<extent_hooks_t> hooks_;
+  // Custom hooks has to outlive corresponding arena.
+  const std::unique_ptr<extent_hooks_t> arena_hooks_;
+
+  // Arena index.
+  const unsigned arena_index_;
+
+  // Hold thread-local tcache index.
+  ThreadLocalPtr tcache_;
 };
 
 }  // namespace rocksdb