Detect if Jemalloc is linked with the binary (#4844)

Summary:
Declare Jemalloc non-standard APIs as weak symbols, so that if Jemalloc is linked with the binary, these symbols will be replaced by Jemalloc's, otherwise they will be nullptr. This is similar to how folly detect jemalloc, but we assume the main program use jemalloc as long as jemalloc is linked: https://github.com/facebook/folly/blob/master/folly/memory/Malloc.h#L147
Pull Request resolved: https://github.com/facebook/rocksdb/pull/4844

Differential Revision: D13574934

Pulled By: yiwu-arbug

fbshipit-source-id: 7ea871beb1be7d5a1259cc38f9b78078793db2db

TARGETS

@@ -67,13 +67,11 @@ is_opt_mode = build_mode.startswith("opt")
 if is_opt_mode:
     rocksdb_compiler_flags.append("-DNDEBUG")
 
-default_allocator = read_config("fbcode", "default_allocator")
-
 sanitizer = read_config("fbcode", "sanitizer")
 
-# Let RocksDB aware of jemalloc existence.
-# Do not enable it if sanitizer presents.
-if is_opt_mode and default_allocator.startswith("jemalloc") and sanitizer == "":
+# Do not enable jemalloc if sanitizer presents. RocksDB will further detect
+# whether the binary is linked with jemalloc at runtime.
+if sanitizer == "":
     rocksdb_compiler_flags.append("-DROCKSDB_JEMALLOC")
     rocksdb_external_deps.append(("jemalloc", None, "headers"))
 

buckifier/targets_cfg.py

@@ -71,13 +71,11 @@ is_opt_mode = build_mode.startswith("opt")
 if is_opt_mode:
     rocksdb_compiler_flags.append("-DNDEBUG")
 
-default_allocator = read_config("fbcode", "default_allocator")
-
 sanitizer = read_config("fbcode", "sanitizer")
 
-# Let RocksDB aware of jemalloc existence.
-# Do not enable it if sanitizer presents.
-if is_opt_mode and default_allocator.startswith("jemalloc") and sanitizer == "":
+# Do not enable jemalloc if sanitizer presents. RocksDB will further detect
+# whether the binary is linked with jemalloc at runtime.
+if sanitizer == "":
     rocksdb_compiler_flags.append("-DROCKSDB_JEMALLOC")
     rocksdb_external_deps.append(("jemalloc", None, "headers"))
 """

db/malloc_stats.cc

@@ -13,17 +13,16 @@
 #include <memory>
 #include <string.h>
 
+#include "port/jemalloc_helper.h"
+
+
 namespace rocksdb {
 
 #ifdef ROCKSDB_JEMALLOC
-#ifdef __FreeBSD__
-#include <malloc_np.h>
-#else
-#include "jemalloc/jemalloc.h"
+
 #ifdef JEMALLOC_NO_RENAME
 #define malloc_stats_print je_malloc_stats_print
 #endif
-#endif
 
 typedef struct {
   char* cur;
@@ -41,10 +40,10 @@ static void GetJemallocStatus(void* mstat_arg, const char* status) {
   snprintf(mstat->cur, buf_size, "%s", status);
   mstat->cur += status_len;
 }
-#endif  // ROCKSDB_JEMALLOC
-
-#ifdef ROCKSDB_JEMALLOC
 void DumpMallocStats(std::string* stats) {
+  if (!HasJemalloc()) {
+    return;
+  }
   MallocStatus mstat;
   const unsigned int kMallocStatusLen = 1000000;
   std::unique_ptr<char[]> buf{new char[kMallocStatusLen + 1]};
@@ -56,5 +55,5 @@ void DumpMallocStats(std::string* stats) {
 #else
 void DumpMallocStats(std::string*) {}
 #endif  // ROCKSDB_JEMALLOC
-}
+}  // namespace rocksdb
 #endif  // !ROCKSDB_LITE

port/jemalloc_helper.h

@@ -0,0 +1,49 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#pragma once
+
+#ifdef ROCKSDB_JEMALLOC
+#ifdef __FreeBSD__
+#include <malloc_np.h>
+#else
+#include <jemalloc/jemalloc.h>
+#endif
+
+// Declare non-standard jemalloc APIs as weak symbols. We can null-check these
+// symbols to detect whether jemalloc is linked with the binary.
+extern "C" void* mallocx(size_t, int) __attribute__((__weak__));
+extern "C" void* rallocx(void*, size_t, int) __attribute__((__weak__));
+extern "C" size_t xallocx(void*, size_t, size_t, int) __attribute__((__weak__));
+extern "C" size_t sallocx(const void*, int) __attribute__((__weak__));
+extern "C" void dallocx(void*, int) __attribute__((__weak__));
+extern "C" void sdallocx(void*, size_t, int) __attribute__((__weak__));
+extern "C" size_t nallocx(size_t, int) __attribute__((__weak__));
+extern "C" int mallctl(const char*, void*, size_t*, void*, size_t)
+    __attribute__((__weak__));
+extern "C" int mallctlnametomib(const char*, size_t*, size_t*)
+    __attribute__((__weak__));
+extern "C" int mallctlbymib(const size_t*, size_t, void*, size_t*, void*,
+                            size_t) __attribute__((__weak__));
+extern "C" void malloc_stats_print(void (*)(void*, const char*), void*,
+                                   const char*) __attribute__((__weak__));
+extern "C" size_t malloc_usable_size(JEMALLOC_USABLE_SIZE_CONST void*)
+    JEMALLOC_CXX_THROW __attribute__((__weak__));
+
+// Check if Jemalloc is linked with the binary. Note the main program might be
+// using a different memory allocator even this method return true.
+// It is loosely based on folly::usingJEMalloc(), minus the check that actually
+// allocate memory and see if it is through jemalloc, to handle the dlopen()
+// case:
+// https://github.com/facebook/folly/blob/76cf8b5841fb33137cfbf8b224f0226437c855bc/folly/memory/Malloc.h#L147
+static inline bool HasJemalloc() {
+  return mallocx != nullptr && rallocx != nullptr && xallocx != nullptr &&
+         sallocx != nullptr && dallocx != nullptr && sdallocx != nullptr &&
+         nallocx != nullptr && mallctl != nullptr &&
+         mallctlnametomib != nullptr && mallctlbymib != nullptr &&
+         malloc_stats_print != nullptr && malloc_usable_size != nullptr;
+}
+
+#endif  // ROCKSDB_JEMALLOC

util/jemalloc_nodump_allocator.cc

@@ -133,12 +133,16 @@ Status NewJemallocNodumpAllocator(
     JemallocAllocatorOptions& options,
     std::shared_ptr<MemoryAllocator>* memory_allocator) {
   *memory_allocator = nullptr;
-#ifndef ROCKSDB_JEMALLOC_NODUMP_ALLOCATOR
-  (void)options;
-  return Status::NotSupported(
+  Status unsupported = Status::NotSupported(
       "JemallocNodumpAllocator only available with jemalloc version >= 5 "
       "and MADV_DONTDUMP is available.");
+#ifndef ROCKSDB_JEMALLOC_NODUMP_ALLOCATOR
+  (void)options;
+  return unsupported;
 #else
+  if (!HasJemalloc()) {
+    return unsupported;
+  }
   if (memory_allocator == nullptr) {
     return Status::InvalidArgument("memory_allocator must be non-null.");
   }

util/jemalloc_nodump_allocator.h

@@ -8,6 +8,7 @@
 #include <atomic>
 #include <vector>
 
+#include "port/jemalloc_helper.h"
 #include "port/port.h"
 #include "rocksdb/memory_allocator.h"
 #include "util/core_local.h"
@@ -15,7 +16,6 @@
 
 #if defined(ROCKSDB_JEMALLOC) && defined(ROCKSDB_PLATFORM_POSIX)
 
-#include <jemalloc/jemalloc.h>
 #include <sys/mman.h>
 
 #if (JEMALLOC_VERSION_MAJOR >= 5) && defined(MADV_DONTDUMP)