Add jemalloc purge (#239)

CMakeLists.txt

@@ -265,7 +265,11 @@ if (MG_ENTERPRISE)
   add_definitions(-DMG_ENTERPRISE)
 endif()
 
+set(ENABLE_JEMALLOC ON)
+
 if (ASAN)
+  message(WARNING "Disabling jemalloc as it doesn't work well with ASAN")
+  set(ENABLE_JEMALLOC OFF)
   # Enable Addres sanitizer and get nicer stack traces in error messages.
   # NOTE: AddressSanitizer uses llvm-symbolizer binary from the Clang
   # distribution to symbolize the stack traces (note that ideally the

src/CMakeLists.txt

@@ -8,6 +8,7 @@ add_subdirectory(io)
 add_subdirectory(kvstore)
 add_subdirectory(telemetry)
 add_subdirectory(communication)
+add_subdirectory(memory)
 add_subdirectory(storage/v2)
 add_subdirectory(integrations)
 add_subdirectory(query)
@@ -39,7 +40,7 @@ if (MG_ENTERPRISE)
 endif()
 
 set(mg_single_node_v2_libs stdc++fs Threads::Threads
-    telemetry_lib mg-query mg-communication mg-new-delete mg-utils)
+    telemetry_lib mg-query mg-communication mg-memory mg-utils)
 if (MG_ENTERPRISE)
   # These are enterprise subsystems
   set(mg_single_node_v2_libs ${mg_single_node_v2_libs} mg-auth mg-audit)

src/memory/CMakeLists.txt

@@ -0,0 +1,10 @@
+set(memory_src_files
+    new_delete.cpp
+    memory_control.cpp)
+
+add_library(mg-memory STATIC ${memory_src_files})
+target_link_libraries(mg-memory mg-utils fmt)
+
+if (ENABLE_JEMALLOC)
+  target_link_libraries(mg-memory jemalloc)
+endif()

src/memory/memory_control.cpp

@@ -0,0 +1,22 @@
+#include "memory_control.hpp"
+
+#if USE_JEMALLOC
+#include <jemalloc/jemalloc.h>
+#endif
+
+namespace memory {
+
+// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
+#define STRINGIFY_HELPER(x) #x
+// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
+#define STRINGIFY(x) STRINGIFY_HELPER(x)
+
+void PurgeUnusedMemory() {
+#if USE_JEMALLOC
+  mallctl("arena." STRINGIFY(MALLCTL_ARENAS_ALL) ".purge", nullptr, nullptr, nullptr, 0);
+#endif
+}
+
+#undef STRINGIFY
+#undef STRINGIFY_HELPER
+}  // namespace memory

src/memory/memory_control.hpp

@@ -0,0 +1,5 @@
+#pragma once
+
+namespace memory {
+void PurgeUnusedMemory();
+}  // namespace memory

src/memory/new_delete.cpp

@@ -4,6 +4,7 @@
 #if USE_JEMALLOC
 #include <jemalloc/jemalloc.h>
 #else
+#include <malloc.h>
 #include <cstdlib>
 #endif
 

src/query/CMakeLists.txt

@@ -41,7 +41,7 @@ add_library(mg-query STATIC ${mg_query_sources})
 add_dependencies(mg-query generate_lcp_query)
 target_include_directories(mg-query PUBLIC ${CMAKE_SOURCE_DIR}/include)
 target_link_libraries(mg-query dl cppitertools)
-target_link_libraries(mg-query mg-integrations-kafka mg-storage-v2 mg-utils mg-kvstore)
+target_link_libraries(mg-query mg-integrations-kafka mg-storage-v2 mg-utils mg-kvstore mg-memory)
 if("${MG_PYTHON_VERSION}" STREQUAL "")
     find_package(Python3 3.5 REQUIRED COMPONENTS Development)
 else()

src/query/interpreter.cpp

@@ -6,6 +6,7 @@
 #include <optional>
 
 #include "glue/communication.hpp"
+#include "memory/memory_control.hpp"
 #include "query/constants.hpp"
 #include "query/context.hpp"
 #include "query/cypher_query_interpreter.hpp"
@@ -1203,6 +1204,7 @@ PreparedQuery PrepareFreeMemoryQuery(ParsedQuery parsed_query, const bool in_exp
       std::move(parsed_query.required_privileges),
       [interpreter_context](AnyStream *stream, std::optional<int> n) -> std::optional<QueryHandlerResult> {
         interpreter_context->db->FreeMemory();
+        memory::PurgeUnusedMemory();
         return QueryHandlerResult::COMMIT;
       },
       RWType::NONE};

src/utils/CMakeLists.txt

@@ -15,6 +15,3 @@ set(utils_src_files
 
 add_library(mg-utils STATIC ${utils_src_files})
 target_link_libraries(mg-utils stdc++fs Threads::Threads spdlog fmt gflags uuid rt)
-
-add_library(mg-new-delete STATIC new_delete.cpp)
-target_link_libraries(mg-new-delete jemalloc fmt)

tests/unit/CMakeLists.txt

@@ -32,7 +32,7 @@ function(_add_unit_test test_cpp custom_main)
   # used to help create two targets of the same name even though CMake
   # requires unique logical target names
   set_target_properties(${target_name} PROPERTIES OUTPUT_NAME ${exec_name})
-  target_link_libraries(${target_name} mg-utils mg-new-delete gtest gmock Threads::Threads dl)
+  target_link_libraries(${target_name} mg-memory mg-utils gtest gmock Threads::Threads dl)
   # register test
   if(TEST_COVERAGE)
     add_test(${target_name} env LLVM_PROFILE_FILE=${exec_name}.profraw ./${exec_name})