Harden usingJEMalloc()

Summary:
In a shared library, we're using whichever memory allocator the main program
uses, but we may still believe we're using jemalloc if we load libjemalloc.so
(because rallocm is defined). This means we'll pass to rallocm pointers that
were not allocated by jemalloc, which is bad.

Test Plan: tests added

Reviewed By: ngbronson@fb.com

FB internal diff: D1228483

diff --git a/folly/Makefile.am b/folly/Makefile.am
index b420da275fdf919e8b33b6ebe2ef62bb98760834..cb6158b84ba4f08f55b45aa198c22927e3ef430a 100644 (file)
--- a/folly/Makefile.am
+++ b/folly/Makefile.am
@@ -189,6 +189,7 @@ libfolly_la_SOURCES = \
        io/async/Request.cpp \
        json.cpp \
        detail/MemoryIdler.cpp \
+       Malloc.cpp \
        MemoryMapping.cpp \
        Random.cpp \
        Range.cpp \

diff --git a/folly/Malloc.cpp b/folly/Malloc.cpp
new file mode 100644 (file)
index 0000000..5f7144d
--- /dev/null
+++ b/folly/Malloc.cpp
@@ -0,0 +1,61 @@
+/*
+ * Copyright 2014 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "folly/Malloc.h"
+
+namespace folly {
+
+// How do we determine that we're using jemalloc?
+// In the hackiest way possible. We allocate memory using malloc() and see if
+// the per-thread counter of allocated memory increases. This makes me feel
+// dirty inside. Also note that this requires jemalloc to have been compiled
+// with --enable-stats.
+bool usingJEMallocSlow() {
+  // Some platforms (*cough* OSX *cough*) require weak symbol checks to be
+  // in the form if (mallctl != NULL). Not if (mallctl) or if (!mallctl) (!!).
+  // http://goo.gl/xpmctm
+  if (allocm == NULL || rallocm == NULL || mallctl == NULL) {
+    return false;
+  }
+
+  // "volatile" because gcc optimizes out the reads from *counter, because
+  // it "knows" malloc doesn't modify global state...
+  volatile uint64_t* counter;
+  size_t counterLen = sizeof(uint64_t*);
+
+  if (mallctl("thread.allocatedp", static_cast<void*>(&counter), &counterLen,
+              nullptr, 0) != 0) {
+    return false;
+  }
+
+  if (counterLen != sizeof(uint64_t*)) {
+    return false;
+  }
+
+  uint64_t origAllocated = *counter;
+
+  void* ptr = malloc(1);
+  if (!ptr) {
+    // wtf, failing to allocate 1 byte
+    return false;
+  }
+  free(ptr);
+
+  return (origAllocated != *counter);
+}
+
+}  // namespaces
+

diff --git a/folly/Malloc.h b/folly/Malloc.h
index 1e02b764ba3835e80effbc01da8781dd2db35e2e..046d82c77b01616dc483be77d6ce7c544a39f038 100644 (file)
--- a/folly/Malloc.h
+++ b/folly/Malloc.h
@@ -45,14 +45,16 @@ namespace folly {
 #pragma GCC system_header
 
 /**
- * Declare rallocm() and allocm() as weak symbols. These will be provided by
- * jemalloc if we are using jemalloc, or will be NULL if we are using another
- * malloc implementation.
+ * Declare rallocm(), allocm(), and mallctl() as weak symbols. These will be
+ * provided by jemalloc if we are using jemalloc, or will be NULL if we are
+ * using another malloc implementation.
  */
 extern "C" int rallocm(void**, size_t*, size_t, size_t, int)
 __attribute__((weak));
 extern "C" int allocm(void**, size_t*, size_t, int)
 __attribute__((weak));
+extern "C" int mallctl(const char*, void*, size_t*, void*, size_t)
+__attribute__((weak));
 
 #include <bits/functexcept.h>
 #define FOLLY_HAVE_MALLOC_H 1
@@ -107,12 +109,18 @@ _GLIBCXX_BEGIN_NAMESPACE_VERSION
 namespace folly {
 #endif
 
+bool usingJEMallocSlow();
 
 /**
  * Determine if we are using jemalloc or not.
  */
 inline bool usingJEMalloc() {
-  return rallocm != NULL;
+  // Checking for rallocm != NULL is not sufficient; we may be in a dlopen()ed
+  // module that depends on libjemalloc, so rallocm is resolved, but the main
+  // program might be using a different memory allocator. Look at the
+  // implementation of usingJEMallocSlow() for the (hacky) details.
+  static bool result = usingJEMallocSlow();
+  return result;
 }
 
 /**

diff --git a/folly/detail/Malloc.cpp b/folly/detail/Malloc.cpp
index 9882dd84c0bce0d06f1ee73d99a7f03a47905557..6a322314fae0fdaf255419c42e745ebfafb993ff 100644 (file)
--- a/folly/detail/Malloc.cpp
+++ b/folly/detail/Malloc.cpp
@@ -21,6 +21,7 @@ extern "C" {
 #if !FOLLY_HAVE_WEAK_SYMBOLS
 int (*rallocm)(void**, size_t*, size_t, size_t, int) = nullptr;
 int (*allocm)(void**, size_t*, size_t, int) = nullptr;
+int (*mallctl)(const char*, void*, size_t*, void*, size_t) = nullptr;
 #endif
 
 }

diff --git a/folly/detail/Malloc.h b/folly/detail/Malloc.h
index fc61b9cf70a16556641d26e1a3d3321698b48ebb..c4ab1b2d11f1ad9a95ce7e844714ba5a4be2d7b3 100644 (file)
--- a/folly/detail/Malloc.h
+++ b/folly/detail/Malloc.h
@@ -28,9 +28,11 @@ extern "C" {
 #if FOLLY_HAVE_WEAK_SYMBOLS
 int rallocm(void**, size_t*, size_t, size_t, int) __attribute__((weak));
 int allocm(void**, size_t*, size_t, int) __attribute__((weak));
+int mallctl(const char*, void*, size_t*, void*, size_t) __attribute__((weak));
 #else
 extern int (*rallocm)(void**, size_t*, size_t, size_t, int);
 extern int (*allocm)(void**, size_t*, size_t, int);
+extern int (*mallctl)(const char*, void*, size_t*, void*, size_t);
 #endif
 
 }