Update folly/JemallocNodumpAllocator to jemalloc 4 & 5 compatible.

Summary:
jemalloc 5.0 comes replaced chunks with extents APIs. Make the API compatible
with both jemalloc 4 and 5.

Reviewed By: yfeldblum

Differential Revision: D5170925

fbshipit-source-id: ce25723975729b0b63371f89a96842a1b2e3b3cc

folly/experimental/JemallocNodumpAllocator.cpp

@@ -37,12 +37,22 @@ bool JemallocNodumpAllocator::extend_and_setup_arena() {
   }
 
   size_t len = sizeof(arena_index_);
-  if (auto ret = mallctl("arenas.extend", &arena_index_, &len, nullptr, 0)) {
+  if (auto ret = mallctl(
+#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
+          "arenas.extend"
+#else
+          "arenas.create"
+#endif
+          ,
+          &arena_index_,
+          &len,
+          nullptr,
+          0)) {
     LOG(FATAL) << "Unable to extend arena: " << errnoStr(ret);
   }
   flags_ = MALLOCX_ARENA(arena_index_) | MALLOCX_TCACHE_NONE;
 
-  // Set the custom alloc hook
+#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
   const auto key =
       folly::to<std::string>("arena.", arena_index_, ".chunk_hooks");
   chunk_hooks_t hooks;
@@ -51,18 +61,36 @@ bool JemallocNodumpAllocator::extend_and_setup_arena() {
   if (auto ret = mallctl(key.c_str(), &hooks, &len, nullptr, 0)) {
     LOG(FATAL) << "Unable to get the hooks: " << errnoStr(ret);
   }
-  if (original_chunk_alloc_ == nullptr) {
-    original_chunk_alloc_ = hooks.alloc;
+  if (original_alloc_ == nullptr) {
+    original_alloc_ = hooks.alloc;
   } else {
-    DCHECK_EQ(original_chunk_alloc_, hooks.alloc);
+    DCHECK_EQ(original_alloc_, hooks.alloc);
   }
 
   // Set the custom hook
-  hooks.alloc = &JemallocNodumpAllocator::chunk_alloc;
+  hooks.alloc = &JemallocNodumpAllocator::alloc;
   if (auto ret =
           mallctl(key.c_str(), nullptr, nullptr, &hooks, sizeof(hooks))) {
     LOG(FATAL) << "Unable to set the hooks: " << errnoStr(ret);
   }
+#else
+  const auto key =
+      folly::to<std::string>("arena.", arena_index_, ".extent_hooks");
+  extent_hooks_t* hooks;
+  len = sizeof(hooks);
+  // Read the existing hooks
+  if (auto ret = mallctl(key.c_str(), &hooks, &len, nullptr, 0)) {
+    LOG(FATAL) << "Unable to get the hooks: " << errnoStr(ret);
+  }
+  if (original_alloc_ == nullptr) {
+    original_alloc_ = hooks->alloc;
+  } else {
+    DCHECK_EQ(original_alloc_, hooks->alloc);
+  }
+
+  // Set the custom hook
+  hooks->alloc = &JemallocNodumpAllocator::alloc;
+#endif
 
   return true;
 #else // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
@@ -80,17 +108,31 @@ void* JemallocNodumpAllocator::reallocate(void* p, size_t size) {
 
 #ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
 
-chunk_alloc_t* JemallocNodumpAllocator::original_chunk_alloc_ = nullptr;
-
-void* JemallocNodumpAllocator::chunk_alloc(
+#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
+chunk_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr;
+void* JemallocNodumpAllocator::alloc(
     void* chunk,
+#else
+extent_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr;
+void* JemallocNodumpAllocator::alloc(
+    extent_hooks_t* extent,
+    void* new_addr,
+#endif
     size_t size,
     size_t alignment,
     bool* zero,
     bool* commit,
     unsigned arena_ind) {
-  void* result =
-      original_chunk_alloc_(chunk, size, alignment, zero, commit, arena_ind);
+  void* result = original_alloc_(
+      JEMALLOC_CHUNK_OR_EXTENT,
+#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_EXTENT
+      new_addr,
+#endif
+      size,
+      alignment,
+      zero,
+      commit,
+      arena_ind);
   if (result != nullptr) {
     if (auto ret = madvise(result, size, MADV_DONTDUMP)) {
       VLOG(1) << "Unable to madvise(MADV_DONTDUMP): " << errnoStr(ret);

folly/experimental/JemallocNodumpAllocator.h

@@ -27,6 +27,13 @@
 
 #if (JEMALLOC_VERSION_MAJOR > 3) && defined(MADV_DONTDUMP)
 #define FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED 1
+#if (JEMALLOC_VERSION_MAJOR == 4)
+#define FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
+#define JEMALLOC_CHUNK_OR_EXTENT chunk
+#else
+#define FOLLY_JEMALLOC_NODUMP_ALLOCATOR_EXTENT
+#define JEMALLOC_CHUNK_OR_EXTENT extent
+#endif
 #endif
 
 #endif // FOLLY_HAVE_LIBJEMALLOC
@@ -41,13 +48,13 @@ namespace folly {
  * the memory is not dump-able.
  *
  * This is done by setting MADV_DONTDUMP using the `madvise` system call. A
- * custom hook installed which is called when allocating a new chunk of memory.
- * All it does is call the original jemalloc hook to allocate the memory and
- * then set the advise on it before returning the pointer to the allocated
- * memory. Jemalloc does not use allocated chunks across different arenas,
- * without `munmap`-ing them first, and the advises are not sticky i.e. they are
- * unset if `munmap` is done. Also this arena can't be used by any other part of
- * the code by just calling `malloc`.
+ * custom hook installed which is called when allocating a new chunk / extent of
+ * memory.  All it does is call the original jemalloc hook to allocate the
+ * memory and then set the advise on it before returning the pointer to the
+ * allocated memory. Jemalloc does not use allocated chunks / extents across
+ * different arenas, without `munmap`-ing them first, and the advises are not
+ * sticky i.e. they are unset if `munmap` is done. Also this arena can't be used
+ * by any other part of the code by just calling `malloc`.
  *
  * If target system doesn't support MADV_DONTDUMP or jemalloc doesn't support
  * custom arena hook, JemallocNodumpAllocator would fall back to using malloc /
@@ -79,13 +86,21 @@ class JemallocNodumpAllocator {
 
  private:
 #ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
-  static chunk_alloc_t* original_chunk_alloc_;
-  static void* chunk_alloc(void* chunk,
-                           size_t size,
-                           size_t alignment,
-                           bool* zero,
-                           bool* commit,
-                           unsigned arena_ind);
+#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
+  static chunk_alloc_t* original_alloc_;
+  static void* alloc(
+      void* chunk,
+#else
+  static extent_alloc_t* original_alloc_;
+  static void* alloc(
+      extent_hooks_t* extent,
+      void* new_addr,
+#endif
+      size_t size,
+      size_t alignment,
+      bool* zero,
+      bool* commit,
+      unsigned arena_ind);
 #endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
 
   bool extend_and_setup_arena();