From afdf38e0c81bb2b9e8e463ed093a1df09838241e Mon Sep 17 00:00:00 2001 From: antoniofilipovic Date: Mon, 9 Oct 2023 12:13:25 +0200 Subject: [PATCH] add per query memory limit --- src/memory/memory_control.cpp | 23 +++++++++++++++++++++++ src/memory/memory_control.hpp | 4 ++++ src/query/interpreter.cpp | 24 +++++++++++++++--------- src/utils/memory_tracker.cpp | 7 +++++++ src/utils/memory_tracker.hpp | 2 ++ 5 files changed, 51 insertions(+), 9 deletions(-) diff --git a/src/memory/memory_control.cpp b/src/memory/memory_control.cpp index 01e9a88b7..4ce4819f1 100644 --- a/src/memory/memory_control.cpp +++ b/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(size)); + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Alloc(static_cast(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(size)); + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Free(static_cast(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(size)); + + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Free(static_cast(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(size)); + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Free(static_cast(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(length)); + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Alloc(static_cast(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(length)); + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Free(static_cast(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(length)); + if (query_limit) [[unlikely]] { + memory_tracker_per_thread.Free(static_cast(length)); + } + return false; } diff --git a/src/memory/memory_control.hpp b/src/memory/memory_control.hpp index 471acf774..f34242d0b 100644 --- a/src/memory/memory_control.hpp +++ b/src/memory/memory_control.hpp @@ -13,9 +13,13 @@ #include #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 diff --git a/src/query/interpreter.cpp b/src/query/interpreter.cpp index 56c7b36d6..84b7a85ef 100644 --- a/src/query/interpreter.cpp +++ b/src/query/interpreter.cpp @@ -1265,6 +1265,17 @@ PullPlan::PullPlan(const std::shared_ptr plan, const Parameters &par std::optional PullPlan::Pull(AnyStream *stream, std::optional n, const std::vector &output_symbols, std::map *summary) { + if (memory_limit_) { + memgraph::memory::query_limit = true; + memgraph::memory::memory_tracker_per_thread.SetMaximumHardLimit(static_cast(*memory_limit_)); + memgraph::memory::memory_tracker_per_thread.SetHardLimit(static_cast(*memory_limit_)); + } + utils::OnScopeExit> 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 PullPlan::Pull(AnyStream *strea pool_memory.emplace(kMaxBlockPerChunks, 1024, &monotonic_memory, &resource_with_exception); } - std::optional 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 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(info.memory_usage))}, {TypedValue("disk_usage"), TypedValue(static_cast(info.disk_usage))}, - // {TypedValue("jemalloc_memory_allocated"), - // TypedValue(utils::GetReadableSize(static_cast(utils::total_memory_tracker.Amount())))}, + // {TypedValue("readable_memory_allocated"), + // TypedValue(utils::GetReadableSize(static_cast(utils::total_memory_tracker.Amount())))}, {TypedValue("memory_allocated"), TypedValue(static_cast(utils::total_memory_tracker.Amount()))}, {TypedValue("allocation_limit"), TypedValue(static_cast(utils::total_memory_tracker.HardLimit()))}, {TypedValue("global_isolation_level"), TypedValue(IsolationLevelToString(storage->GetIsolationLevel()))}, diff --git a/src/utils/memory_tracker.cpp b/src/utils/memory_tracker.cpp index fa9910e5e..774faf72c 100644 --- a/src/utils/memory_tracker.cpp +++ b/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."); diff --git a/src/utils/memory_tracker.hpp b/src/utils/memory_tracker.hpp index ce45435d5..4db3182b9 100644 --- a/src/utils/memory_tracker.hpp +++ b/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 {