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

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 \

folly/Malloc.cpp

+/*
+ * 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
+

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;
 }
 
 /**

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
 
 }

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
 
 }