add poc memory tracker per procedure

This commit introduces memory tracker with jemalloc extent hooks which fully works in case when jemalloc config is following:
MALLOC_CONF="retain:false,percpu_arena:percpu,oversize_threshold:1000000000000,muzzy_decay_ms:0,dirty_decay_ms:0" \
./configure \
    --disable-cxx \
    $COMMON_CONFIGURE_FLAGS \
    --with-malloc-conf="retain:false,percpu_arena:percpu,oversize_threshold:1000000000000,muzzy_decay_ms:0,dirty_decay_ms:0"

This config will for jemalloc not to use lazy purge or MADV_FREE(muzzy_decay_ms=0 and dirty_decay_ms=0), it will force jemalloc not to use
custom arena for huge allocations (oversize_threshold) and it will force jemalloc not extend virtual memory indefinitely (retain=false)
and therefore call alloc hook when allocation actually takes place.

Only problem is if we do huge allocations which are not mapped on alloc directly, in that case jemalloc uses cache and we can overcounter
allocation size.

cmake/FindJemalloc.cmake

@@ -1,55 +0,0 @@
-# Try to find jemalloc library
-#
-# Use this module as:
-#   find_package(Jemalloc)
-#
-# or:
-#   find_package(Jemalloc REQUIRED)
-#
-# This will define the following variables:
-#
-#   Jemalloc_FOUND           True if the system has the jemalloc library.
-#   Jemalloc_INCLUDE_DIRS    Include directories needed to use jemalloc.
-#   Jemalloc_LIBRARIES       Libraries needed to link to jemalloc.
-#
-# The following cache variables may also be set:
-#
-# Jemalloc_INCLUDE_DIR     The directory containing jemalloc/jemalloc.h.
-# Jemalloc_LIBRARY         The path to the jemalloc static library.
-
-find_path(Jemalloc_INCLUDE_DIR NAMES jemalloc/jemalloc.h PATH_SUFFIXES include)
-
-find_library(Jemalloc_LIBRARY NAMES libjemalloc.a PATH_SUFFIXES lib)
-
-include(FindPackageHandleStandardArgs)
-find_package_handle_standard_args(Jemalloc
-  FOUND_VAR Jemalloc_FOUND
-  REQUIRED_VARS
-    Jemalloc_LIBRARY
-    Jemalloc_INCLUDE_DIR
-)
-
-if(Jemalloc_FOUND)
-  set(Jemalloc_LIBRARIES ${Jemalloc_LIBRARY})
-  set(Jemalloc_INCLUDE_DIRS ${Jemalloc_INCLUDE_DIR})
-else()
-  if(Jemalloc_FIND_REQUIRED)
-    message(FATAL_ERROR "Cannot find jemalloc!")
-  else()
-    message(WARNING "jemalloc is not found!")
-  endif()
-endif()
-
-if(Jemalloc_FOUND AND NOT TARGET Jemalloc::Jemalloc)
-  add_library(Jemalloc::Jemalloc UNKNOWN IMPORTED)
-  set_target_properties(Jemalloc::Jemalloc
-    PROPERTIES
-      IMPORTED_LOCATION "${Jemalloc_LIBRARY}"
-      INTERFACE_INCLUDE_DIRECTORIES "${Jemalloc_INCLUDE_DIR}"
-  )
-endif()
-
-mark_as_advanced(
-  Jemalloc_INCLUDE_DIR
-  Jemalloc_LIBRARY
-)

cmake/Findjemalloc.cmake

@@ -0,0 +1,67 @@
+# Try to find jemalloc library
+#
+# Use this module as:
+#   find_package(jemalloc)
+#
+# or:
+#   find_package(jemalloc REQUIRED)
+#
+# This will define the following variables:
+#
+#   JEMALLOC_FOUND           True if the system has the jemalloc library.
+#   Jemalloc_INCLUDE_DIRS    Include directories needed to use jemalloc.
+#   Jemalloc_LIBRARIES       Libraries needed to link to jemalloc.
+#
+# The following cache variables may also be set:
+#
+# Jemalloc_INCLUDE_DIR     The directory containing jemalloc/jemalloc.h.
+# Jemalloc_LIBRARY         The path to the jemalloc static library.
+
+
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(jemalloc
+  FOUND_VAR JEMALLOC_FOUND
+  REQUIRED_VARS
+    JEMALLOC_LIBRARY
+    JEMALLOC_INCLUDE_DIR
+)
+
+if(JEMALLOC_INCLUDE_DIR)
+  message(STATUS "Found jemalloc include dir: ${JEMALLOC_INCLUDE_DIR}")
+else()
+  message(WARNING "jemalloc not found!")
+endif()
+
+if(JEMALLOC_LIBRARY)
+  message(STATUS "Found jemalloc library: ${JEMALLOC_LIBRARY}")
+else()
+  message(WARNING "jemalloc library not found!")
+endif()
+
+if(JEMALLOC_FOUND)
+  set(Jemalloc_LIBRARIES ${JEMALLOC_LIBRARY})
+  set(Jemalloc_INCLUDE_DIRS ${JEMALLOC_INCLUDE_DIR})
+else()
+  if(Jemalloc_FIND_REQUIRED)
+    message(FATAL_ERROR "Cannot find jemalloc!")
+  else()
+    message(WARNING "jemalloc is not found!")
+  endif()
+endif()
+
+if(JEMALLOC_FOUND AND NOT TARGET Jemalloc::Jemalloc)
+  message(STATUS "JEMALLOC NOT TARGET")
+
+  add_library(Jemalloc::Jemalloc UNKNOWN IMPORTED)
+  set_target_properties(Jemalloc::Jemalloc
+    PROPERTIES
+      IMPORTED_LOCATION "${JEMALLOC_LIBRARY}"
+      INTERFACE_INCLUDE_DIRECTORIES "${JEMALLOC_INCLUDE_DIR}"
+  )
+endif()
+
+mark_as_advanced(
+  JEMALLOC_INCLUDE_DIR
+  JEMALLOC_LIBRARY
+)

include/mg_procedure.h

@@ -111,6 +111,35 @@ enum mgp_error mgp_global_aligned_alloc(size_t size_in_bytes, size_t alignment,
 /// The behavior is undefined if `ptr` is not a value returned from a prior
 /// mgp_global_alloc() or mgp_global_aligned_alloc().
 void mgp_global_free(void *p);
+
+/// State of the graph database.
+struct mgp_graph;
+
+/// Allocations are tracked only for master thread. If new threads are spawned
+/// inside procedure, by calling following function with thread id
+/// you can start tracking allocations for that thread too. This
+/// is important if you need query memory limit to work
+/// for given procedure or per procedure memory limit.
+enum mgp_error mgp_track_thread_allocations(struct mgp_graph *graph, const char *thread_id);
+
+/// Once allocations are tracked for custom thread, you need to stop tracking allocations
+/// for given thread, before thread finishes with execution, or is detached.
+/// Otherwise it might result in slowdown of system due to unnecessary tracking of
+/// allocations.
+enum mgp_error mgp_untrack_thread_allocations(struct mgp_graph *graph, const char *thread_id);
+
+/// Allocations are tracked only for master thread. If new threads are spawned
+/// inside procedure, by calling following function
+/// you can start tracking allocations for current thread too. This
+/// is important if you need query memory limit to work
+/// for given procedure or per procedure memory limit.
+enum mgp_error mgp_track_current_thread_allocations(struct mgp_graph *graph);
+
+/// Once allocations are tracked for current thread, you need to stop tracking allocations
+/// for given thread, before thread finishes with execution, or is detached.
+/// Otherwise it might result in slowdown of system due to unnecessary tracking of
+/// allocations.
+enum mgp_error mgp_untrack_current_thread_allocations(struct mgp_graph *graph);
 ///@}
 
 /// @name Operations on mgp_value
@@ -851,9 +880,6 @@ enum mgp_error mgp_edge_set_properties(struct mgp_edge *e, struct mgp_map *prope
 enum mgp_error mgp_edge_iter_properties(struct mgp_edge *e, struct mgp_memory *memory,
                                         struct mgp_properties_iterator **result);
 
-/// State of the graph database.
-struct mgp_graph;
-
 /// Get the vertex corresponding to given ID, or NULL if no such vertex exists.
 /// Resulting vertex must be freed using mgp_vertex_destroy.
 /// Return mgp_error::MGP_ERROR_UNABLE_TO_ALLOCATE if unable to allocate the vertex.

libs/CMakeLists.txt

@@ -15,7 +15,6 @@ set(GFLAGS_NOTHREADS OFF)
 # NOTE: config/generate.py depends on the gflags help XML format.
 find_package(gflags REQUIRED)
 find_package(fmt 8.0.1)
-find_package(Jemalloc REQUIRED)
 find_package(ZLIB 1.2.11 REQUIRED)
 
 set(LIB_DIR ${CMAKE_CURRENT_SOURCE_DIR})
@@ -99,6 +98,17 @@ macro(import_external_library name type library_location include_dir)
   import_library(${name} ${type} ${${_upper_name}_LIBRARY} ${${_upper_name}_INCLUDE_DIR})
 endmacro(import_external_library)
 
+
+macro(set_path_external_library name type library_location include_dir)
+  string(TOUPPER ${name} _upper_name)
+  set(${_upper_name}_LIBRARY ${library_location} CACHE FILEPATH
+      "Path to ${name} library" FORCE)
+  set(${_upper_name}_INCLUDE_DIR ${include_dir} CACHE FILEPATH
+      "Path to ${name} include directory" FORCE)
+  mark_as_advanced(${name}_LIBRARY ${name}_INCLUDE_DIR)
+endmacro(set_path_external_library)
+
+
 # setup antlr
 import_external_library(antlr4 STATIC
   ${CMAKE_CURRENT_SOURCE_DIR}/antlr4/runtime/Cpp/lib/libantlr4-runtime.a
@@ -265,3 +275,8 @@ import_header_library(ctre ${CMAKE_CURRENT_SOURCE_DIR})
 # setup absl (cmake sub_directory tolerant)
 set(ABSL_PROPAGATE_CXX_STD ON)
 add_subdirectory(absl EXCLUDE_FROM_ALL)
+
+# set Jemalloc
+set_path_external_library(jemalloc STATIC
+  ${CMAKE_CURRENT_SOURCE_DIR}/jemalloc/lib/libjemalloc.a
+  ${CMAKE_CURRENT_SOURCE_DIR}/jemalloc/include/)

libs/setup.sh

@@ -124,6 +124,7 @@ declare -A primary_urls=(
   ["librdtsc"]="http://$local_cache_host/git/librdtsc.git"
   ["ctre"]="http://$local_cache_host/file/hanickadot/compile-time-regular-expressions/v3.7.2/single-header/ctre.hpp"
   ["absl"]="https://$local_cache_host/git/abseil-cpp.git"
+  ["jemalloc"]="https://$local_cache_host/git/jemalloc.git"
 )
 
 # The goal of secondary urls is to have links to the "source of truth" of
@@ -151,6 +152,7 @@ declare -A secondary_urls=(
   ["librdtsc"]="https://github.com/gabrieleara/librdtsc.git"
   ["ctre"]="https://raw.githubusercontent.com/hanickadot/compile-time-regular-expressions/v3.7.2/single-header/ctre.hpp"
   ["absl"]="https://github.com/abseil/abseil-cpp.git"
+  ["jemalloc"]="https://github.com/jemalloc/jemalloc.git"
 )
 
 # antlr
@@ -252,3 +254,21 @@ cd ..
 # abseil 20230125.3
 absl_ref="20230125.3"
 repo_clone_try_double "${primary_urls[absl]}" "${secondary_urls[absl]}" "absl" "$absl_ref"
+
+
+# jemalloc ea6b3e973b477b8061e0076bb257dbd7f3faa756
+JEMALLOC_COMMIT_VERSION="5.2.1"
+repo_clone_try_double "${secondary_urls[jemalloc]}" "${secondary_urls[jemalloc]}" "jemalloc" "$JEMALLOC_COMMIT_VERSION"
+
+# this is hack for cmake in libs to set path, and for FindJemalloc to use Jemalloc_INCLUDE_DIR
+pushd jemalloc
+
+./autogen.sh
+MALLOC_CONF="retain:false,percpu_arena:percpu,oversize_threshold:0,muzzy_decay_ms:5000,dirty_decay_ms:5000" \
+./configure \
+  --disable-cxx \
+  --enable-shared=no --prefix=$working_dir \
+  --with-malloc-conf="retain:false,percpu_arena:percpu,oversize_threshold:0,muzzy_decay_ms:5000,dirty_decay_ms:5000"
+
+make -j$CPUS install
+popd

src/memgraph.cpp

@@ -22,6 +22,7 @@
 #include "glue/run_id.hpp"
 #include "helpers.hpp"
 #include "license/license_sender.hpp"
+#include "memory/memory_control.hpp"
 #include "query/config.hpp"
 #include "query/discard_value_stream.hpp"
 #include "query/interpreter.hpp"
@@ -106,6 +107,7 @@ void InitSignalHandlers(const std::function<void()> &shutdown_fun) {
 }
 
 int main(int argc, char **argv) {
+  memgraph::memory::SetHooks();
   google::SetUsageMessage("Memgraph database server");
   gflags::SetVersionString(version_string);
 
@@ -197,7 +199,6 @@ int main(int argc, char **argv) {
           "won't be available.");
     }
   }
-
   std::cout << "You are running Memgraph v" << gflags::VersionString() << std::endl;
   std::cout << "To get started with Memgraph, visit https://memgr.ph/start" << std::endl;
 

src/memory/CMakeLists.txt

@@ -2,7 +2,9 @@ set(memory_src_files
     new_delete.cpp
     memory_control.cpp)
 
-find_package(Jemalloc REQUIRED)
+
+
+find_package(jemalloc REQUIRED)
 
 add_library(mg-memory STATIC ${memory_src_files})
 target_link_libraries(mg-memory mg-utils fmt)

src/memory/memory_control.cpp

@@ -1,4 +1,4 @@
-// Copyright 2022 Memgraph Ltd.
+// Copyright 2023 Memgraph Ltd.
 //
 // Use of this software is governed by the Business Source License
 // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -10,6 +10,11 @@
 // licenses/APL.txt.
 
 #include "memory_control.hpp"
+#include <cstdint>
+#include <sstream>
+#include <thread>
+#include "utils/logging.hpp"
+#include "utils/memory_tracker.hpp"
 
 #if USE_JEMALLOC
 #include <jemalloc/jemalloc.h>
@@ -22,6 +27,274 @@ namespace memgraph::memory {
 // NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
 #define STRINGIFY(x) STRINGIFY_HELPER(x)
 
+std::string get_string_thread_id(const std::thread::id &thread_id) {
+  std::ostringstream oss;
+  oss << thread_id;
+  return oss.str();
+}
+
+// TODO (af) think if following implementation would make sense
+/*
+std::string get_thread_id() {
+  static thread_local std::thread::id current_thread_id = std::this_thread::get_id();
+  // figure out how to cache this part
+  std::ostringstream oss;
+  oss << current_thread_id;
+  return oss.str();
+}
+*/
+std::string get_thread_id() { return get_string_thread_id(std::this_thread::get_id()); }
+
+#if USE_JEMALLOC
+
+static void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t alignment, bool *zero,
+                      bool *commit, unsigned arena_ind);
+static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind);
+static void my_destroy(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind);
+static bool my_commit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
+                      unsigned arena_ind);
+static bool my_decommit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
+                        unsigned arena_ind);
+static bool my_purge_forced(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
+                            unsigned arena_ind);
+extent_hooks_t *old_hooks = nullptr;
+
+static extent_hooks_t custom_hooks = {
+    .alloc = &my_alloc,
+    .dalloc = &my_dalloc,
+    .destroy = &my_destroy,
+    .commit = &my_commit,
+    .decommit = &my_decommit,
+    .purge_lazy = nullptr,
+    .purge_forced = &my_purge_forced,
+    .split = nullptr,
+    .merge = nullptr,
+};
+
+static const extent_hooks_t *new_hooks = &custom_hooks;
+
+void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t alignment, bool *zero, bool *commit,
+               unsigned arena_ind) {
+  // This needs to be before, to throw exception in case of too big alloc
+  if (*commit) [[likely]] {
+    memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
+    if (arena_tracking[arena_ind]) [[unlikely]] {
+      transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Alloc(static_cast<int64_t>(size));
+    }
+  }
+
+  auto *ptr = old_hooks->alloc(extent_hooks, new_addr, size, alignment, zero, commit, arena_ind);
+  if (ptr == nullptr) [[unlikely]] {
+    if (*commit) {
+      memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
+      if (arena_tracking[arena_ind]) [[unlikely]] {
+        transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Free(static_cast<int64_t>(size));
+      }
+    }
+    return ptr;
+  }
+
+  return ptr;
+}
+
+static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind) {
+  auto err = old_hooks->dalloc(extent_hooks, addr, size, committed, arena_ind);
+
+  if (err) [[unlikely]] {
+    return err;
+  }
+
+  if (committed) [[likely]] {
+    memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
+
+    if (arena_tracking[arena_ind]) [[unlikely]] {
+      transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Free(static_cast<int64_t>(size));
+    }
+  }
+
+  return false;
+}
+
+static void my_destroy(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind) {
+  if (committed) [[likely]] {
+    memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
+    if (arena_tracking[arena_ind]) [[unlikely]] {
+      transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Free(static_cast<int64_t>(size));
+    }
+  }
+
+  old_hooks->destroy(extent_hooks, addr, size, committed, arena_ind);
+}
+
+static bool my_commit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
+                      unsigned arena_ind) {
+  auto err = old_hooks->commit(extent_hooks, addr, size, offset, length, arena_ind);
+
+  if (err) {
+    return err;
+  }
+
+  memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(length));
+  if (arena_tracking[arena_ind]) [[unlikely]] {
+    transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Alloc(static_cast<int64_t>(size));
+  }
+
+  return false;
+}
+
+static bool my_decommit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
+                        unsigned arena_ind) {
+  MG_ASSERT(old_hooks && old_hooks->decommit);
+  auto err = old_hooks->decommit(extent_hooks, addr, size, offset, length, arena_ind);
+
+  if (err) {
+    return err;
+  }
+
+  memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
+  if (arena_tracking[arena_ind]) [[unlikely]] {
+    transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Free(static_cast<int64_t>(size));
+  }
+
+  return false;
+}
+
+static bool my_purge_forced(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
+                            unsigned arena_ind) {
+  MG_ASSERT(old_hooks && old_hooks->purge_forced);
+  auto err = old_hooks->purge_forced(extent_hooks, addr, size, offset, length, arena_ind);
+
+  if (err) [[unlikely]] {
+    return err;
+  }
+  memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
+
+  if (arena_tracking[arena_ind]) [[unlikely]] {
+    transaction_id_tracker[thread_id_to_transaction_id[get_thread_id()]].Free(static_cast<int64_t>(size));
+  }
+
+  return false;
+}
+
+#endif
+
+void SetHooks() {
+#if USE_JEMALLOC
+
+  uint64_t allocated{0};
+  uint64_t sz{sizeof(allocated)};
+
+  sz = sizeof(unsigned);
+  unsigned n_arenas{0};
+  int err = mallctl("opt.narenas", (void *)&n_arenas, &sz, nullptr, 0);
+
+  if (err) {
+    return;
+  }
+
+  spdlog::trace("n areanas {}", n_arenas);
+
+  if (nullptr != old_hooks) {
+    return;
+  }
+
+  for (int i = 0; i < n_arenas; i++) {
+    arena_tracking[i] = 0;
+    std::string func_name = "arena." + std::to_string(i) + ".extent_hooks";
+
+    size_t hooks_len = sizeof(old_hooks);
+
+    int err = mallctl(func_name.c_str(), &old_hooks, &hooks_len, nullptr, 0);
+
+    if (err) {
+      LOG_FATAL("Error getting hooks for jemalloc arena {}", i);
+    }
+
+    // Due to the way jemalloc works, we need first to set their hooks
+    // which will trigger creating arena, then we can set our custom hook wrappers
+
+    err = mallctl(func_name.c_str(), nullptr, nullptr, &old_hooks, sizeof(old_hooks));
+
+    MG_ASSERT(old_hooks);
+    MG_ASSERT(old_hooks->alloc);
+    MG_ASSERT(old_hooks->dalloc);
+    MG_ASSERT(old_hooks->destroy);
+    MG_ASSERT(old_hooks->commit);
+    MG_ASSERT(old_hooks->decommit);
+    MG_ASSERT(old_hooks->purge_forced);
+    MG_ASSERT(old_hooks->purge_lazy);
+    MG_ASSERT(old_hooks->split);
+    MG_ASSERT(old_hooks->merge);
+
+    custom_hooks.purge_lazy = old_hooks->purge_lazy;
+    custom_hooks.split = old_hooks->split;
+    custom_hooks.merge = old_hooks->merge;
+
+    if (err) {
+      LOG_FATAL("Error setting jemalloc hooks for jemalloc arena {}", i);
+    }
+
+    err = mallctl(func_name.c_str(), nullptr, nullptr, &new_hooks, sizeof(new_hooks));
+
+    if (err) {
+      LOG_FATAL("Error setting custom hooks for jemalloc arena {}", i);
+    }
+  }
+
+#endif
+}
+
+unsigned GetArenaForThread() {
+#if USE_JEMALLOC
+  unsigned thread_arena{0};
+  size_t size_thread_arena = sizeof(thread_arena);
+  int err = mallctl("thread.arena", &thread_arena, &size_thread_arena, nullptr, 0);
+  if (err) {
+    return -1;
+  }
+  return thread_arena;
+#endif
+  return -1;
+}
+
+bool AddTrackingOnArena(unsigned arena_id) {
+#if USE_JEMALLOC
+  arena_tracking[arena_id].fetch_add(1);
+#endif
+  return false;
+}
+
+bool RemoveTrackingOnArena(unsigned arena_id) {
+#if USE_JEMALLOC
+  arena_tracking[arena_id].fetch_sub(1);
+#endif
+  return true;
+}
+
+void UpdateThreadToTransactionId(const std::thread::id &thread_id, uint64_t transaction_id) {
+  thread_id_to_transaction_id[get_string_thread_id(thread_id)] = transaction_id;
+}
+
+void UpdateThreadToTransactionId(const char *thread_id, uint64_t transaction_id) {
+  thread_id_to_transaction_id[std::string(thread_id)] = transaction_id;
+}
+
+void ResetThreadToTransactionId(const std::thread::id &thread_id) {
+  thread_id_to_transaction_id.erase(get_string_thread_id(thread_id));
+}
+
+void ResetThreadToTransactionId(const char *thread_id) { thread_id_to_transaction_id.erase(std::string(thread_id)); }
+
+void AddTrackingsOnCurrentThread(uint64_t transaction_id) {
+  UpdateThreadToTransactionId(std::this_thread::get_id(), transaction_id);
+  AddTrackingOnArena(memgraph::memory::GetArenaForThread());
+}
+
+void RemoveTrackingsOnCurrentThread() {
+  ResetThreadToTransactionId(std::this_thread::get_id());
+  RemoveTrackingOnArena(GetArenaForThread());
+}
+
 void PurgeUnusedMemory() {
 #if USE_JEMALLOC
   mallctl("arena." STRINGIFY(MALLCTL_ARENAS_ALL) ".purge", nullptr, nullptr, nullptr, 0);
@@ -30,4 +303,5 @@ void PurgeUnusedMemory() {
 
 #undef STRINGIFY
 #undef STRINGIFY_HELPER
+
 }  // namespace memgraph::memory

src/memory/memory_control.hpp

@@ -1,4 +1,4 @@
-// Copyright 2022 Memgraph Ltd.
+// Copyright 2023 Memgraph Ltd.
 //
 // Use of this software is governed by the Business Source License
 // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -11,6 +11,34 @@
 
 #pragma once
 
+#include <cstddef>
+#include <cstdint>
+#include <unordered_map>
+#include "utils/logging.hpp"
+#include "utils/memory_tracker.hpp"
+
+#include <thread>
 namespace memgraph::memory {
+
 void PurgeUnusedMemory();
+void SetHooks();
+
+// TODO(af) This should all be part of memgraph::memory::thread namespace, and moved to different file
+// This should be part of class
+unsigned GetArenaForThread();
+bool AddTrackingOnArena(unsigned);
+bool RemoveTrackingOnArena(unsigned);
+void UpdateThreadToTransactionId(const std::thread::id &, uint64_t);
+void ResetThreadToTransactionId(const std::thread::id &);
+void UpdateThreadToTransactionId(const char *, uint64_t);
+void ResetThreadToTransactionId(const char *);
+void AddTrackingsOnCurrentThread(uint64_t);
+void RemoveTrackingsOnCurrentThread();
+
+inline std::unordered_map<std::string, uint64_t> thread_id_to_transaction_id;
+// TODO(af): think if we need to solve issue of tracking allocations for arena
+// if user forgets to unregister tracking for that thread before it dies.
+inline std::unordered_map<unsigned, std::atomic<int>> arena_tracking;
+inline std::unordered_map<uint64_t, utils::MemoryTracker> transaction_id_tracker;
+
 }  // namespace memgraph::memory

src/memory/new_delete.cpp

@@ -1,4 +1,4 @@
-// Copyright 2022 Memgraph Ltd.
+// Copyright 2023 Memgraph Ltd.
 //
 // Use of this software is governed by the Business Source License
 // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -87,21 +87,15 @@ void deleteSized(void *ptr, const std::size_t /*unused*/, const std::align_val_t
 #endif
 
 void TrackMemory(std::size_t size) {
-#if USE_JEMALLOC
-  if (size != 0) [[likely]] {
-    size = nallocx(size, 0);
-  }
-#endif
+#if !USE_JEMALLOC
   memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
+#endif
 }
 
 void TrackMemory(std::size_t size, const std::align_val_t align) {
-#if USE_JEMALLOC
-  if (size != 0) [[likely]] {
-    size = nallocx(size, MALLOCX_ALIGN(align));  // NOLINT(hicpp-signed-bitwise)
-  }
-#endif
+#if !USE_JEMALLOC
   memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
+#endif
 }
 
 bool TrackMemoryNoExcept(const std::size_t size) {
@@ -126,11 +120,7 @@ bool TrackMemoryNoExcept(const std::size_t size, const std::align_val_t align) {
 
 void UntrackMemory([[maybe_unused]] void *ptr, [[maybe_unused]] std::size_t size = 0) noexcept {
   try {
-#if USE_JEMALLOC
-    if (ptr != nullptr) [[likely]] {
-      memgraph::utils::total_memory_tracker.Free(sallocx(ptr, 0));
-    }
-#else
+#if !USE_JEMALLOC
     if (size) {
       memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
     } else {
@@ -144,11 +134,7 @@ void UntrackMemory([[maybe_unused]] void *ptr, [[maybe_unused]] std::size_t size
 
 void UntrackMemory(void *ptr, const std::align_val_t align, [[maybe_unused]] std::size_t size = 0) noexcept {
   try {
-#if USE_JEMALLOC
-    if (ptr != nullptr) [[likely]] {
-      memgraph::utils::total_memory_tracker.Free(sallocx(ptr, MALLOCX_ALIGN(align)));  // NOLINT(hicpp-signed-bitwise)
-    }
-#else
+#if !USE_JEMALLOC
     if (size) {
       memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
     } else {

src/query/db_accessor.cpp

@@ -21,6 +21,18 @@ namespace memgraph::query {
 SubgraphDbAccessor::SubgraphDbAccessor(query::DbAccessor db_accessor, Graph *graph)
     : db_accessor_(db_accessor), graph_(graph) {}
 
+void SubgraphDbAccessor::TrackThreadAllocations(const char *thread_id) {
+  return db_accessor_.TrackThreadAllocations(thread_id);
+}
+
+void SubgraphDbAccessor::TrackCurrentThreadAllocations() { return db_accessor_.TrackCurrentThreadAllocations(); }
+
+void SubgraphDbAccessor::UntrackThreadAllocations(const char *thread_id) {
+  return db_accessor_.UntrackThreadAllocations(thread_id);
+}
+
+void SubgraphDbAccessor::UntrackCurrentThreadAllocations() { return db_accessor_.TrackCurrentThreadAllocations(); }
+
 storage::PropertyId SubgraphDbAccessor::NameToProperty(const std::string_view name) {
   return db_accessor_.NameToProperty(name);
 }

src/query/db_accessor.hpp

@@ -17,6 +17,7 @@
 #include <cppitertools/filter.hpp>
 #include <cppitertools/imap.hpp>
 
+#include "memory/memory_control.hpp"
 #include "query/exceptions.hpp"
 #include "storage/v2/edge_accessor.hpp"
 #include "storage/v2/id_types.hpp"
@@ -372,6 +373,26 @@ class DbAccessor final {
 
   void FinalizeTransaction() { accessor_->FinalizeTransaction(); }
 
+  void TrackThreadAllocations(const char *thread_id) {
+    memgraph::memory::UpdateThreadToTransactionId(thread_id, *accessor_->GetTransactionId());
+    auto arena = memgraph::memory::GetArenaForThread();
+    memgraph::memory::AddTrackingOnArena(arena);
+  }
+
+  void TrackCurrentThreadAllocations() {
+    memgraph::memory::AddTrackingsOnCurrentThread(*accessor_->GetTransactionId());
+  }
+
+  void UntrackThreadAllocations(const char *thread_id) {
+    memgraph::memory::ResetThreadToTransactionId(thread_id);
+    auto arena = memgraph::memory::GetArenaForThread();
+    memgraph::memory::RemoveTrackingOnArena(arena);
+  }
+
+  void UntrackCurrentThreadAllocations() { memgraph::memory::RemoveTrackingsOnCurrentThread(); }
+
+  std::optional<uint64_t> GetTransactionId() { return accessor_->GetTransactionId(); }
+
   VerticesIterable Vertices(storage::View view) { return VerticesIterable(accessor_->Vertices(view)); }
 
   VerticesIterable Vertices(storage::View view, storage::LabelId label) {
@@ -643,6 +664,14 @@ class SubgraphDbAccessor final {
 
   static SubgraphDbAccessor *MakeSubgraphDbAccessor(DbAccessor *db_accessor, Graph *graph);
 
+  void TrackThreadAllocations(const char *thread_id);
+
+  void TrackCurrentThreadAllocations();
+
+  void UntrackThreadAllocations(const char *thread_id);
+
+  void UntrackCurrentThreadAllocations();
+
   storage::PropertyId NameToProperty(std::string_view name);
 
   storage::LabelId NameToLabel(std::string_view name);

src/query/interpreter.cpp

@@ -26,6 +26,7 @@
 #include <optional>
 #include <stdexcept>
 #include <thread>
+#include <tuple>
 #include <unordered_map>
 #include <utility>
 #include <variant>
@@ -1265,6 +1266,29 @@ PullPlan::PullPlan(const std::shared_ptr<CachedPlan> plan, const Parameters &par
 std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::Pull(AnyStream *stream, std::optional<int> n,
                                                                 const std::vector<Symbol> &output_symbols,
                                                                 std::map<std::string, TypedValue> *summary) {
+  std::optional<uint64_t> transaction_id = ctx_.db_accessor->GetTransactionId();
+  MG_ASSERT(transaction_id.has_value());
+  unsigned arena_ind{0};
+  if (memory_limit_) {
+    // TODO (AF) think to isolate this in namespace or make a class
+    memgraph::memory::transaction_id_tracker.emplace(std::piecewise_construct, std::forward_as_tuple(*transaction_id),
+                                                     std::forward_as_tuple());
+    auto &memory_tracker = memgraph::memory::transaction_id_tracker[*transaction_id];
+    memory_tracker.SetMaximumHardLimit(static_cast<int64_t>(*memory_limit_));
+    memory_tracker.SetHardLimit(static_cast<int64_t>(*memory_limit_));
+    arena_ind = memgraph::memory::GetArenaForThread();
+    memgraph::memory::AddTrackingOnArena(arena_ind);
+    memgraph::memory::UpdateThreadToTransactionId(std::this_thread::get_id(), *transaction_id);
+  }
+  utils::OnScopeExit<std::function<void()>> reset_query_limit{
+      [memory_limit = memory_limit_, transaction_id = *transaction_id, arena_ind]() {
+        if (memory_limit) {
+          // TODO (AF) think to isolate this in namespace or make a class
+          memgraph::memory::transaction_id_tracker.erase(transaction_id);
+          memgraph::memory::RemoveTrackingOnArena(arena_ind);
+          memgraph::memory::ResetThreadToTransactionId(std::this_thread::get_id());
+        }
+      }};
   // Set up temporary memory for a single Pull. Initial memory comes from the
   // stack. 256 KiB should fit on the stack and should be more than enough for a
   // single `Pull`.
@@ -1288,13 +1312,7 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::Pull(AnyStream *strea
     pool_memory.emplace(kMaxBlockPerChunks, 1024, &monotonic_memory, &resource_with_exception);
   }
 
-  std::optional<utils::LimitedMemoryResource> maybe_limited_resource;
-  if (memory_limit_) {
-    maybe_limited_resource.emplace(&*pool_memory, *memory_limit_);
-    ctx_.evaluation_context.memory = &*maybe_limited_resource;
-  } else {
-    ctx_.evaluation_context.memory = &*pool_memory;
-  }
+  ctx_.evaluation_context.memory = &*pool_memory;
 
   // Returns true if a result was pulled.
   const auto pull_result = [&]() -> bool { return cursor_->Pull(frame_, ctx_); };
@@ -1361,6 +1379,7 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::Pull(AnyStream *strea
   }
   cursor_->Shutdown();
   ctx_.profile_execution_time = execution_time_;
+
   return GetStatsWithTotalTime(ctx_);
 }
 
@@ -2978,7 +2997,6 @@ PreparedQuery PrepareDatabaseInfoQuery(ParsedQuery parsed_query, bool in_explici
           results.push_back({TypedValue(label_property_index_mark), TypedValue(storage->LabelToName(item.first)),
                              TypedValue(storage->PropertyToName(item.second))});
         }
-
         std::sort(results.begin(), results.end(), [&label_index_mark](const auto &record_1, const auto &record_2) {
           const auto type_1 = record_1[0].ValueString();
           const auto type_2 = record_2[0].ValueString();
@@ -3068,6 +3086,8 @@ PreparedQuery PrepareSystemInfoQuery(ParsedQuery parsed_query, bool in_explicit_
             {TypedValue("average_degree"), TypedValue(info.average_degree)},
             {TypedValue("memory_usage"), TypedValue(static_cast<int64_t>(info.memory_usage))},
             {TypedValue("disk_usage"), TypedValue(static_cast<int64_t>(info.disk_usage))},
+            {TypedValue("readable_memory_allocated"),
+             TypedValue(utils::GetReadableSize(static_cast<double>(utils::total_memory_tracker.Amount())))},
             {TypedValue("memory_allocated"), TypedValue(static_cast<int64_t>(utils::total_memory_tracker.Amount()))},
             {TypedValue("allocation_limit"), TypedValue(static_cast<int64_t>(utils::total_memory_tracker.HardLimit()))},
             {TypedValue("global_isolation_level"), TypedValue(IsolationLevelToString(storage->GetIsolationLevel()))},
@@ -3097,7 +3117,6 @@ PreparedQuery PrepareSystemInfoQuery(ParsedQuery parsed_query, bool in_explicit_
                            action = action_on_complete;
                            pull_plan = std::make_shared<PullPlanVector>(std::move(results));
                          }
-
                          if (pull_plan->Pull(stream, n)) {
                            return action;
                          }

src/query/procedure/mg_procedure_impl.cpp

@@ -3540,3 +3540,28 @@ mgp_error mgp_log(const mgp_log_level log_level, const char *output) {
     throw std::invalid_argument{fmt::format("Invalid log level: {}", log_level)};
   });
 }
+
+mgp_error mgp_track_thread_allocations(mgp_graph *graph, const char *thread_id) {
+  return WrapExceptions([&]() {
+    std::visit([thread_id](auto *db_accessor) -> void { db_accessor->TrackThreadAllocations(thread_id); }, graph->impl);
+  });
+}
+
+mgp_error mgp_track_current_thread_allocations(mgp_graph *graph) {
+  return WrapExceptions([&]() {
+    std::visit([](auto *db_accessor) -> void { db_accessor->TrackCurrentThreadAllocations(); }, graph->impl);
+  });
+}
+
+mgp_error mgp_untrack_thread_allocations(mgp_graph *graph, const char *thread_id) {
+  return WrapExceptions([&]() {
+    std::visit([thread_id](auto *db_accessor) -> void { db_accessor->UntrackThreadAllocations(thread_id); },
+               graph->impl);
+  });
+}
+
+mgp_error mgp_untrack_current_thread_allocations(mgp_graph *graph) {
+  return WrapExceptions([&]() {
+    std::visit([](auto *db_accessor) -> void { db_accessor->UntrackCurrentThreadAllocations(); }, graph->impl);
+  });
+}

src/storage/v2/inmemory/storage.cpp

@@ -272,6 +272,7 @@ InMemoryStorage::InMemoryAccessor::DetachDelete(std::vector<VertexAccessor *> no
 
   // Need to inform the next CollectGarbage call that there are some
   // non-transactional deletions that need to be collected
+
   auto const inform_gc_vertex_deletion = utils::OnScopeExit{[this, &deleted_vertices = deleted_vertices]() {
     if (!deleted_vertices.empty() && transaction_.storage_mode == StorageMode::IN_MEMORY_ANALYTICAL) {
       auto *mem_storage = static_cast<InMemoryStorage *>(storage_);

src/utils/memory_tracker.cpp

@@ -89,6 +89,13 @@ void MemoryTracker::TryRaiseHardLimit(const int64_t limit) {
     ;
 }
 
+void MemoryTracker::ResetTrackings() {
+  hard_limit_.store(0, std::memory_order_relaxed);
+  peak_.store(0, std::memory_order_relaxed);
+  amount_.store(0, std::memory_order_relaxed);
+  maximum_hard_limit_ = 0;
+}
+
 void MemoryTracker::SetMaximumHardLimit(const int64_t limit) {
   if (maximum_hard_limit_ < 0) {
     spdlog::warn("Invalid maximum hard limit.");

src/utils/memory_tracker.hpp

@@ -1,4 +1,4 @@
-// Copyright 2022 Memgraph Ltd.
+// Copyright 2023 Memgraph Ltd.
 //
 // Use of this software is governed by the Business Source License
 // included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -22,8 +22,15 @@ class OutOfMemoryException : public utils::BasicException {
   explicit OutOfMemoryException(const std::string &msg) : utils::BasicException(msg) {}
 };
 
+// exapand on tracking procedure allocations by switching mode of allocations for memory tracker
+enum class TrackingMode { Procedure, Default };
+
 class MemoryTracker final {
  private:
+  std::atomic<int64_t> proc_amount_{0};
+  std::atomic<int64_t> proc_peak_{0};
+  std::atomic<int64_t> proc_hard_limit_{0};
+
   std::atomic<int64_t> amount_{0};
   std::atomic<int64_t> peak_{0};
   std::atomic<int64_t> hard_limit_{0};
@@ -58,6 +65,8 @@ class MemoryTracker final {
   void TryRaiseHardLimit(int64_t limit);
   void SetMaximumHardLimit(int64_t limit);
 
+  void ResetTrackings();
+
   // By creating an object of this class, every allocation in its scope that goes over
   // the set hard limit produces an OutOfMemoryException.
   class OutOfMemoryExceptionEnabler final {

tests/e2e/disk_storage/workloads.yaml

@@ -1,7 +1,7 @@
 disk_storage: &disk_storage
   cluster:
     main:
-      args: ["--bolt-port", "7687", "--log-level", "TRACE", "--memory-limit", "50"]
+      args: ["--bolt-port", "7687", "--log-level", "TRACE", "--memory-limit", "125"]
       log_file: "disk_storage.log"
       setup_queries: []
       validation_queries: []

tests/e2e/memgraph.py

@@ -130,9 +130,14 @@ class MemgraphInstanceRunner:
     def stop(self):
         if not self.is_running():
             return
-        self.proc_mg.terminate()
-        code = self.proc_mg.wait()
-        assert code == 0, "The Memgraph process exited with non-zero!"
+
+        pid = self.proc_mg.pid
+        try:
+            os.kill(pid, 15)  # 15 is the signal number for SIGTERM
+        except os.OSError:
+            assert False
+
+        time.sleep(1)
 
     def kill(self):
         if not self.is_running():

tests/e2e/memory/CMakeLists.txt

@@ -11,3 +11,5 @@ target_link_libraries(memgraph__e2e__memory__limit_global_alloc gflags mgclient
 add_executable(memgraph__e2e__memory__limit_global_alloc_proc memory_limit_global_alloc_proc.cpp)
 target_link_libraries(memgraph__e2e__memory__limit_global_alloc_proc gflags mgclient mg-utils mg-io Threads::Threads)
 
+add_executable(memgraph__e2e__memory__limit_query_alloc_proc_multi_thread query_memory_limit.cpp)
+target_link_libraries(memgraph__e2e__memory__limit_query_alloc_proc_multi_thread gflags mgclient mg-utils mg-io Threads::Threads)

tests/e2e/memory/procedures/CMakeLists.txt

@@ -3,3 +3,8 @@ target_include_directories(global_memory_limit PRIVATE ${CMAKE_SOURCE_DIR}/inclu
 
 add_library(global_memory_limit_proc SHARED global_memory_limit_proc.c)
 target_include_directories(global_memory_limit_proc PRIVATE ${CMAKE_SOURCE_DIR}/include)
+
+
+add_library(query_memory_limit_proc_multi_thread SHARED query_memory_limit_proc_multi_thread.cpp)
+target_include_directories(query_memory_limit_proc_multi_thread PRIVATE ${CMAKE_SOURCE_DIR}/include)
+target_link_libraries(query_memory_limit_proc_multi_thread mg-utils)

tests/e2e/memory/procedures/query_memory_limit_proc_multi_thread.cpp

@@ -0,0 +1,125 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#include <atomic>
+#include <cassert>
+#include <exception>
+#include <functional>
+#include <mgp.hpp>
+#include <mutex>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#include "mg_procedure.h"
+#include "utils/on_scope_exit.hpp"
+
+enum mgp_error Alloc(void *ptr) {
+  const size_t two_sixty_eight_mb = 1 << 28;
+
+  return mgp_global_alloc(two_sixty_eight_mb, (void **)(&ptr));
+}
+
+// change communication between threads with feature and promise
+std::atomic<int> num_allocations{0};
+std::vector<void *> ptrs_;
+
+void AllocFunc(mgp_graph *graph) {
+  [[maybe_unused]] const enum mgp_error tracking_error = mgp_track_current_thread_allocations(graph);
+
+  void *ptr = nullptr;
+
+  ptrs_.emplace_back(ptr);
+  try {
+    const enum mgp_error alloc_err = Alloc(ptr);
+    if (alloc_err != mgp_error::MGP_ERROR_UNABLE_TO_ALLOCATE) {
+      num_allocations.fetch_add(1, std::memory_order_relaxed);
+    }
+  } catch (const std::exception &e) {
+    // not to terminate std::thread here
+    assert(false);
+  }
+
+  [[maybe_unused]] const enum mgp_error untracking_error = mgp_untrack_current_thread_allocations(graph);
+}
+
+void DualThread(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory) {
+  mgp::MemoryDispatcherGuard guard{memory};
+  const auto arguments = mgp::List(args);
+  const auto record_factory = mgp::RecordFactory(result);
+  num_allocations.store(0, std::memory_order_relaxed);
+  try {
+    std::vector<std::thread> threads;
+
+    for (int i = 0; i < 2; i++) {
+      threads.emplace_back(AllocFunc, memgraph_graph);
+    }
+
+    for (int i = 0; i < 2; i++) {
+      threads[i].join();
+    }
+    for (void *ptr : ptrs_) {
+      if (ptr != nullptr) {
+        mgp_global_free(ptr);
+      }
+    }
+
+    auto new_record = record_factory.NewRecord();
+
+    new_record.Insert("allocated_all", num_allocations.load(std::memory_order_relaxed) == 2);
+  } catch (std::exception &e) {
+    record_factory.SetErrorMessage(e.what());
+  }
+}
+
+void Regular(mgp_list *args, mgp_graph *memgraph_graph, mgp_result *result, mgp_memory *memory) {
+  mgp::MemoryDispatcherGuard guard{memory};
+  const auto arguments = mgp::List(args);
+  const auto record_factory = mgp::RecordFactory(result);
+
+  try {
+    void *ptr{nullptr};
+
+    memgraph::utils::OnScopeExit<std::function<void(void)>> cleanup{[&ptr]() {
+      if (nullptr == ptr) {
+        return;
+      }
+      mgp_global_free(ptr);
+    }};
+
+    const enum mgp_error alloc_err = Alloc(ptr);
+    auto new_record = record_factory.NewRecord();
+    new_record.Insert("allocated", alloc_err != mgp_error::MGP_ERROR_UNABLE_TO_ALLOCATE);
+  } catch (std::exception &e) {
+    record_factory.SetErrorMessage(e.what());
+  }
+}
+
+extern "C" int mgp_init_module(struct mgp_module *module, struct mgp_memory *memory) {
+  try {
+    mgp::memory = memory;
+
+    AddProcedure(DualThread, std::string("dual_thread").c_str(), mgp::ProcedureType::Read, {},
+                 {mgp::Return(std::string("allocated_all").c_str(), mgp::Type::Bool)}, module, memory);
+
+    AddProcedure(Regular, std::string("regular").c_str(), mgp::ProcedureType::Read, {},
+                 {mgp::Return(std::string("allocated").c_str(), mgp::Type::Bool)}, module, memory);
+
+  } catch (const std::exception &e) {
+    return 1;
+  }
+
+  return 0;
+}
+
+extern "C" int mgp_shutdown_module() { return 0; }

tests/e2e/memory/query_memory_limit.cpp

@@ -0,0 +1,68 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#include <gflags/gflags.h>
+#include <algorithm>
+#include <exception>
+#include <ios>
+#include <iostream>
+#include <mgclient.hpp>
+
+#include "utils/logging.hpp"
+#include "utils/timer.hpp"
+
+DEFINE_uint64(bolt_port, 7687, "Bolt port");
+DEFINE_uint64(timeout, 120, "Timeout seconds");
+DEFINE_bool(multi_db, false, "Run test in multi db environment");
+
+int main(int argc, char **argv) {
+  google::SetUsageMessage("Memgraph E2E Query Memory Limit In Multi-Thread For Global Allocators");
+  gflags::ParseCommandLineFlags(&argc, &argv, true);
+  memgraph::logging::RedirectToStderr();
+
+  mg::Client::Init();
+
+  auto client =
+      mg::Client::Connect({.host = "127.0.0.1", .port = static_cast<uint16_t>(FLAGS_bolt_port), .use_ssl = false});
+  if (!client) {
+    LOG_FATAL("Failed to connect!");
+  }
+
+  if (FLAGS_multi_db) {
+    client->Execute("CREATE DATABASE clean;");
+    client->DiscardAll();
+    client->Execute("USE DATABASE clean;");
+    client->DiscardAll();
+    client->Execute("MATCH (n) DETACH DELETE n;");
+    client->DiscardAll();
+  }
+
+  MG_ASSERT(
+      client->Execute("CALL libquery_memory_limit_proc_multi_thread.dual_thread() YIELD allocated_all RETURN "
+                      "allocated_all QUERY MEMORY LIMIT 500MB"));
+  bool error{false};
+  try {
+    auto result_rows = client->FetchAll();
+    if (result_rows) {
+      auto row = *result_rows->begin();
+      error = row[0].ValueBool() == false;
+      std::cout << std::boolalpha << error << ", error status" << std::endl;
+    }
+
+  } catch (const std::exception &e) {
+    std::cout << e.what() << std::endl;
+    error = true;
+  }
+
+  MG_ASSERT(error, "Error should have happend");
+
+  return 0;
+}

tests/e2e/memory/workloads.yaml

@@ -2,7 +2,7 @@ bolt_port: &bolt_port "7687"
 args: &args
   - "--bolt-port"
   - *bolt_port
-  - "--memory-limit=1000"
+  - "--memory-limit=1024"
   - "--storage-gc-cycle-sec=180"
   - "--log-level=TRACE"
 
@@ -23,6 +23,19 @@ disk_cluster: &disk_cluster
         - "STORAGE MODE ON_DISK_TRANSACTIONAL"
       validation_queries: []
 
+args_query_limit: &args_query_limit
+  - "--bolt-port"
+  - *bolt_port
+  - "--storage-gc-cycle-sec=180"
+  - "--log-level=TRACE"
+
+in_memory_query_limit_cluster: &in_memory_query_limit_cluster
+  cluster:
+    main:
+      args: *args_query_limit
+      log_file: "memory-e2e.log"
+      setup_queries: []
+      validation_queries: []
 
 workloads:
   - name: "Memory control"
@@ -70,3 +83,9 @@ workloads:
     args: ["--bolt-port", *bolt_port, "--timeout", "180"]
     proc: "tests/e2e/memory/procedures/"
     <<: *disk_cluster
+
+  - name: "Memory control query limit"
+    binary: "tests/e2e/memory/memgraph__e2e__memory__limit_query_alloc_proc_multi_thread"
+    args: ["--bolt-port", *bolt_port, "--timeout", "180"]
+    proc: "tests/e2e/memory/procedures/"
+    <<: *in_memory_query_limit_cluster