Enable cacheline_aligned_alloc() to allocate from jemalloc if enabled.

Summary:
Reuse WITH_JEMALLOC option in preparation for module search unification.
  Move jemalloc overrides into a separate .cc
  Remote obsolete JEMALLOC_NOINIT option.
Closes https://github.com/facebook/rocksdb/pull/3078

Differential Revision: D6174826

Pulled By: yiwu-arbug

fbshipit-source-id: 9970a0289b4490272d15853920d9d7531af91140

CMakeLists.txt

@@ -41,10 +41,10 @@ endif()
 
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_SOURCE_DIR}/cmake/modules/")
 
+option(WITH_JEMALLOC "build with JeMalloc" OFF)
 if(MSVC)
   include(${CMAKE_CURRENT_SOURCE_DIR}/thirdparty.inc)
 else()
-  option(WITH_JEMALLOC "build with JeMalloc" OFF)
   if(CMAKE_SYSTEM_NAME MATCHES "FreeBSD")
     # FreeBSD has jemaloc as default malloc
     # but it does not have all the jemalloc files in include/...
@@ -576,6 +576,12 @@ if(WIN32)
     port/win/win_logger.cc
     port/win/win_thread.cc
     port/win/xpress_win.cc)
+	
+if(WITH_JEMALLOC)
+  list(APPEND SOURCES
+    port/win/win_jemalloc.cc)
+endif()
+	
 else()
   list(APPEND SOURCES
     port/port_posix.cc

port/win/port_win.cc

@@ -228,80 +228,3 @@ int GetMaxOpenFiles() { return -1; }
 
 }  // namespace port
 }  // namespace rocksdb
-
-#ifdef JEMALLOC
-
-#include "jemalloc/jemalloc.h"
-
-#ifndef JEMALLOC_NON_INIT
-
-namespace rocksdb {
-
-namespace port {
-
-__declspec(noinline) void WINAPI InitializeJemalloc() {
-  je_init();
-  atexit(je_uninit);
-}
-
-}  // port
-}  // rocksdb
-
-extern "C" {
-
-#ifdef _WIN64
-
-#pragma comment(linker, "/INCLUDE:p_rocksdb_init_jemalloc")
-
-typedef void(WINAPI* CRT_Startup_Routine)(void);
-
-// .CRT section is merged with .rdata on x64 so it must be constant data.
-// must be of external linkage
-// We put this into XCT since we want to run this earlier than C++ static
-// constructors
-// which are placed into XCU
-#pragma const_seg(".CRT$XCT")
-extern const CRT_Startup_Routine p_rocksdb_init_jemalloc;
-const CRT_Startup_Routine p_rocksdb_init_jemalloc =
-    rocksdb::port::InitializeJemalloc;
-#pragma const_seg()
-
-#else  // _WIN64
-
-// x86 untested
-
-#pragma comment(linker, "/INCLUDE:_p_rocksdb_init_jemalloc")
-
-#pragma section(".CRT$XCT", read)
-JEMALLOC_SECTION(".CRT$XCT") JEMALLOC_ATTR(used) static const void(
-    WINAPI* p_rocksdb_init_jemalloc)(void) = rocksdb::port::InitializeJemalloc;
-
-#endif  // _WIN64
-
-}  // extern "C"
-
-#endif // JEMALLOC_NON_INIT
-
-// Global operators to be replaced by a linker
-
-void* operator new(size_t size) {
-  void* p = je_malloc(size);
-  if (!p) {
-    throw std::bad_alloc();
-  }
-  return p;
-}
-
-void* operator new[](size_t size) {
-  void* p = je_malloc(size);
-  if (!p) {
-    throw std::bad_alloc();
-  }
-  return p;
-}
-
-void operator delete(void* p) { je_free(p); }
-
-void operator delete[](void* p) { je_free(p); }
-
-#endif  // JEMALLOC

port/win/port_win.h

@@ -241,13 +241,31 @@ extern void InitOnce(OnceType* once, void (*initializer)());
 #define CACHE_LINE_SIZE 64U
 #endif
 
+#ifdef ROCKSDB_JEMALLOC
+#include "jemalloc/jemalloc.h"
+// Separate inlines so they can be replaced if needed
+inline void* jemalloc_aligned_alloc( size_t size, size_t alignment) {
+  return je_aligned_alloc(alignment, size);
+}
+inline void jemalloc_aligned_free(void* p) {
+  je_free(p);
+}
+#endif
 
 inline void *cacheline_aligned_alloc(size_t size) {
+#ifdef ROCKSDB_JEMALLOC
+  return jemalloc_aligned_alloc(size, CACHE_LINE_SIZE);
+#else
   return _aligned_malloc(size, CACHE_LINE_SIZE);
+#endif
 }
 
 inline void cacheline_aligned_free(void *memblock) {
+#ifdef ROCKSDB_JEMALLOC
+  jemalloc_aligned_free(memblock);
+#else
   _aligned_free(memblock);
+#endif
 }
 
 // https://gcc.gnu.org/bugzilla/show_bug.cgi?id=52991 for MINGW32

port/win/win_jemalloc.cc

@@ -0,0 +1,47 @@
+//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
+//  This source code is licensed under the BSD-style license found in the
+//  LICENSE file in the root directory of this source tree. An additional grant
+//  of patent rights can be found in the PATENTS file in the same directory.
+//
+// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file. See the AUTHORS file for names of contributors.
+
+#ifndef ROCKSDB_JEMALLOC
+# error This file can only be part of jemalloc aware build
+#endif
+
+#include <stdexcept>
+#include "jemalloc/jemalloc.h"
+
+// Global operators to be replaced by a linker when this file is
+// a part of the build
+
+void* operator new(size_t size) {
+  void* p = je_malloc(size);
+  if (!p) {
+    throw std::bad_alloc();
+  }
+  return p;
+}
+
+void* operator new[](size_t size) {
+  void* p = je_malloc(size);
+  if (!p) {
+    throw std::bad_alloc();
+  }
+  return p;
+}
+
+void operator delete(void* p) {
+  if (p) {
+    je_free(p);
+  }
+}
+
+void operator delete[](void* p) {
+  if (p) {
+    je_free(p);
+  }
+}
+

port/win/xpress_win.cc

@@ -17,10 +17,6 @@
 
 #ifdef XPRESS
 
-#ifdef JEMALLOC
-#include <jemalloc/jemalloc.h>
-#endif
-
 // Put this under ifdef so windows systems w/o this
 // can still build
 #include <compressapi.h>
@@ -43,22 +39,6 @@ auto CloseDecompressorFun = [](void* h) {
     ::CloseDecompressor(reinterpret_cast<DECOMPRESSOR_HANDLE>(h));
   }
 };
-
-
-#ifdef JEMALLOC
-// Make sure compressors use our jemalloc if redirected
-PVOID CompressorAlloc(PVOID, SIZE_T size) {
-  return je_malloc(size);
-}
-
-VOID CompressorFree(PVOID, PVOID p) {
-  if (p != NULL) {
-    je_free(p);
-  }
-}
-
-#endif
-
 }
 
 bool Compress(const char* input, size_t length, std::string* output) {
@@ -73,17 +53,6 @@ bool Compress(const char* input, size_t length, std::string* output) {
 
   COMPRESS_ALLOCATION_ROUTINES* allocRoutinesPtr = nullptr;
 
-#ifdef JEMALLOC
-  COMPRESS_ALLOCATION_ROUTINES allocationRoutines;
-
-  //  Init. allocation routines
-  allocationRoutines.Allocate = CompressorAlloc;
-  allocationRoutines.Free = CompressorFree;
-  allocationRoutines.UserContext = NULL;
-
-  allocRoutinesPtr = &allocationRoutines;
-#endif
-
   COMPRESSOR_HANDLE compressor = NULL;
 
   BOOL success = CreateCompressor(
@@ -94,17 +63,17 @@ bool Compress(const char* input, size_t length, std::string* output) {
   if (!success) {
 #ifdef _DEBUG
     std::cerr << "XPRESS: Failed to create Compressor LastError: " <<
-       GetLastError() << std::endl;
+      GetLastError() << std::endl;
 #endif
     return false;
   }
 
   std::unique_ptr<void, decltype(CloseCompressorFun)>
-     compressorGuard(compressor, CloseCompressorFun);
+    compressorGuard(compressor, CloseCompressorFun);
 
   SIZE_T compressedBufferSize = 0;
 
- //  Query compressed buffer size.
+  //  Query compressed buffer size.
   success = ::Compress(
     compressor,                 //  Compressor Handle
     const_cast<char*>(input),   //  Input buffer
@@ -123,8 +92,8 @@ bool Compress(const char* input, size_t length, std::string* output) {
         "XPRESS: Failed to estimate compressed buffer size LastError " <<
         lastError << std::endl;
 #endif
-       return false;
+      return false;
-     }
+    }
   }
 
   assert(compressedBufferSize > 0);
@@ -146,7 +115,7 @@ bool Compress(const char* input, size_t length, std::string* output) {
   if (!success) {
 #ifdef _DEBUG
     std::cerr << "XPRESS: Failed to compress LastError " <<
-       GetLastError() << std::endl;
+      GetLastError() << std::endl;
 #endif
     return false;
   }
@@ -169,16 +138,6 @@ char* Decompress(const char* input_data, size_t input_length,
 
   COMPRESS_ALLOCATION_ROUTINES* allocRoutinesPtr = nullptr;
 
-#ifdef JEMALLOC
-  COMPRESS_ALLOCATION_ROUTINES allocationRoutines;
-
-  //  Init. allocation routines
-  allocationRoutines.Allocate = CompressorAlloc;
-  allocationRoutines.Free = CompressorFree;
-  allocationRoutines.UserContext = NULL;
-  allocRoutinesPtr = &allocationRoutines;
-#endif
-
   DECOMPRESSOR_HANDLE decompressor = NULL;
 
   BOOL success = CreateDecompressor(
@@ -190,7 +149,7 @@ char* Decompress(const char* input_data, size_t input_length,
   if (!success) {
 #ifdef _DEBUG
     std::cerr << "XPRESS: Failed to create Decompressor LastError "
-              << GetLastError() << std::endl;
+      << GetLastError() << std::endl;
 #endif
     return nullptr;
   }
@@ -215,8 +174,8 @@ char* Decompress(const char* input_data, size_t input_length,
     if (lastError != ERROR_INSUFFICIENT_BUFFER) {
 #ifdef _DEBUG
       std::cerr
-          << "XPRESS: Failed to estimate decompressed buffer size LastError "
+        << "XPRESS: Failed to estimate decompressed buffer size LastError "
-          << lastError << std::endl;
+        << lastError << std::endl;
 #endif
       return nullptr;
     }

thirdparty.inc

@@ -8,8 +8,6 @@ set(USE_SNAPPY_DEFAULT 0)        # SNAPPY is disabled by default, enable with -D
 set(USE_LZ4_DEFAULT 0)           # LZ4 is disabled by default, enable with -DLZ4=1 cmake command line agrument
 set(USE_ZLIB_DEFAULT 0)          # ZLIB is disabled by default, enable with -DZLIB=1 cmake command line agrument
 set(USE_XPRESS_DEFAULT 0)        # XPRESS is disabled by default, enable with -DXPRESS=1 cmake command line agrument
-set(USE_JEMALLOC_DEFAULT 0)      # JEMALLOC is disabled by default, enable with -DJEMALLOC=1 cmake command line agrument
-set(USE_JENONINIT_DEFAULT 1)     # Default is enabled do not call je_init/je_uninit as the newer versions do not have it disable with -DJENONINIT=0
 
 #
 # This example assumes all the libraries locate in directories under THIRDPARTY_HOME environment variable
@@ -219,15 +217,15 @@ set(JEMALLOC_LIB_RELEASE ${JEMALLOC_HOME}/bin/retail/amd64/jemalloc.lib)
 #
 # Don't touch these lines
 #
-if (DEFINED JEMALLOC)
-  set(USE_JEMALLOC ${JEMALLOC})
-else ()
-  set(USE_JEMALLOC ${USE_JEMALLOC_DEFAULT})
-endif ()
 
-if (${USE_JEMALLOC} EQUAL 1)
+# For compatibilty with previous
+if(JEMALLOC)
+  set(WITH_JEMALLOC ON)
+endif()
+
+if (WITH_JEMALLOC)
   message(STATUS "JEMALLOC library is enabled")
-  set(JEMALLOC_CXX_FLAGS "-DJEMALLOC -DJEMALLOC_EXPORT= ")
+  set(JEMALLOC_CXX_FLAGS "-DROCKSDB_JEMALLOC -DJEMALLOC_EXPORT= ")
   
   if(DEFINED ENV{JEMALLOC_INCLUDE})
     set(JEMALLOC_INCLUDE $ENV{JEMALLOC_INCLUDE})
@@ -248,16 +246,7 @@ if (${USE_JEMALLOC} EQUAL 1)
   set (THIRDPARTY_LIBS ${THIRDPARTY_LIBS} ${JEMALLOC_LIBS})
   set (ARTIFACT_SUFFIX "_je")
   
-  set(USE_JENONINIT USE_JENONINIT_DEFAULT)
+  set(WITH_JEMALLOC ON)
- 
-  if(JENONINIT)
-    set(USE_JENONINIT ${JENONINIT})
-  endif()
-  
-  if(${USE_JENONINIT} EQUAL 1)
-    add_definitions(-DJEMALLOC_NON_INIT)
-    message(STATUS "JEMALLOC NONINIT version")
-  endif()
   
 else ()
   set (ARTIFACT_SUFFIX "")