add per query memory limit

src/memory/memory_control.cpp

@@ -57,12 +57,18 @@ void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t
   // 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 (query_limit) [[unlikely]] {
+      memory_tracker_per_thread.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 (query_limit) [[unlikely]] {
+        memory_tracker_per_thread.Free(static_cast<int64_t>(size));
+      }
     }
     return ptr;
   }
@@ -79,6 +85,10 @@ static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, boo
 
   if (committed) [[likely]] {
     memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
+
+    if (query_limit) [[unlikely]] {
+      memory_tracker_per_thread.Free(static_cast<int64_t>(size));
+    }
   }
 
   return false;
@@ -87,6 +97,9 @@ static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, boo
 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 (query_limit) [[unlikely]] {
+      memory_tracker_per_thread.Free(static_cast<int64_t>(size));
+    }
   }
 
   old_hooks->destroy(extent_hooks, addr, size, committed, arena_ind);
@@ -101,6 +114,9 @@ static bool my_commit(extent_hooks_t *extent_hooks, void *addr, size_t size, siz
   }
 
   memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(length));
+  if (query_limit) [[unlikely]] {
+    memory_tracker_per_thread.Alloc(static_cast<int64_t>(length));
+  }
 
   return false;
 }
@@ -115,6 +131,9 @@ static bool my_decommit(extent_hooks_t *extent_hooks, void *addr, size_t size, s
   }
 
   memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
+  if (query_limit) [[unlikely]] {
+    memory_tracker_per_thread.Free(static_cast<int64_t>(length));
+  }
 
   return false;
 }
@@ -129,6 +148,10 @@ static bool my_purge_forced(extent_hooks_t *extent_hooks, void *addr, size_t siz
   }
   memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
 
+  if (query_limit) [[unlikely]] {
+    memory_tracker_per_thread.Free(static_cast<int64_t>(length));
+  }
+
   return false;
 }
 

src/memory/memory_control.hpp

@@ -13,9 +13,13 @@
 
 #include <cstddef>
 #include "utils/logging.hpp"
+#include "utils/memory_tracker.hpp"
 namespace memgraph::memory {
 
 void PurgeUnusedMemory();
 void SetHooks();
 
+inline thread_local bool query_limit{false};
+inline thread_local utils::MemoryTracker memory_tracker_per_thread{};
+
 }  // namespace memgraph::memory

src/query/interpreter.cpp

@@ -1265,6 +1265,17 @@ 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) {
+  if (memory_limit_) {
+    memgraph::memory::query_limit = true;
+    memgraph::memory::memory_tracker_per_thread.SetMaximumHardLimit(static_cast<int64_t>(*memory_limit_));
+    memgraph::memory::memory_tracker_per_thread.SetHardLimit(static_cast<int64_t>(*memory_limit_));
+  }
+  utils::OnScopeExit<std::function<void()>> reset_query_limit{[memory_limit = memory_limit_]() {
+    if (memory_limit) {
+      memgraph::memory::query_limit = false;
+      memgraph::memory::memory_tracker_per_thread.ResetTrackings();
+    }
+  }};
   // 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 +1299,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 +1366,7 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::Pull(AnyStream *strea
   }
   cursor_->Shutdown();
   ctx_.profile_execution_time = execution_time_;
+
   return GetStatsWithTotalTime(ctx_);
 }
 
@@ -3067,8 +3073,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("jemalloc_memory_allocated"),
-            //  TypedValue(utils::GetReadableSize(static_cast<double>(utils::total_memory_tracker.Amount())))},
+            // {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()))},

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

@@ -58,6 +58,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 {