tursom/memgraph
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Use extent hooks for memory procedure limit (#1443)

Browse Source
This commit is contained in:
Antonio Filipovic
2023-11-07 16:04:29 +01:00
committed by GitHub
parent ece4b0dba8
commit 4d5ea03dfa
19 changed files with 656 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
src/memory/global_memory_control.cpp
View File

@@ -62,10 +62,7 @@ void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t
if (*commit) [[likely]] {
memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Alloc(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackAllocOnCurrentThread(size);
}
}
@@ -74,10 +71,7 @@ void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t
if (*commit) {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Free(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
}
}
return ptr;
@@ -97,10 +91,7 @@ static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, boo
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Free(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
}
}
@@ -111,10 +102,7 @@ static void my_destroy(extent_hooks_t *extent_hooks, void *addr, size_t size, bo
if (committed) [[likely]] {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
if (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Free(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
}
}
@@ -131,10 +119,7 @@ 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 (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Alloc(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
}
return false;
@@ -151,10 +136,7 @@ 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 (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Free(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
}
return false;
@@ -171,10 +153,7 @@ 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 (GetQueriesMemoryControl().IsArenaTracked(arena_ind)) [[unlikely]] {
auto *memory_tracker = GetQueriesMemoryControl().GetTrackerCurrentThread();
if (memory_tracker != nullptr) [[likely]] {
memory_tracker->Alloc(static_cast<int64_t>(size));
}
GetQueriesMemoryControl().TrackFreeOnCurrentThread(size);
}
return false;

111
src/memory/query_memory_control.cpp
View File

@@ -10,6 +10,7 @@
// licenses/APL.txt.
#include <atomic>
#include <cassert>
#include <cstdint>
#include <iostream>
#include <optional>
@@ -58,34 +59,58 @@ void QueriesMemoryControl::EraseThreadToTransactionId(const std::thread::id &thr
accessor.remove(thread_id);
}
utils::MemoryTracker *QueriesMemoryControl::GetTrackerCurrentThread() {
void QueriesMemoryControl::TrackAllocOnCurrentThread(size_t size) {
auto thread_id_to_transaction_id_accessor = thread_id_to_transaction_id.access();
// we might be just constructing mapping between thread id and transaction id
// so we miss this allocation
auto thread_id_to_transaction_id_elem = thread_id_to_transaction_id_accessor.find(std::this_thread::get_id());
if (thread_id_to_transaction_id_elem == thread_id_to_transaction_id_accessor.end()) {
return nullptr;
return;
}
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
auto transaction_id_to_tracker =
transaction_id_to_tracker_accessor.find(thread_id_to_transaction_id_elem->transaction_id);
return &transaction_id_to_tracker->tracker;
// It can happen that some allocation happens between mapping thread to
// transaction id, so we miss this allocation
if (transaction_id_to_tracker == transaction_id_to_tracker_accessor.end()) [[unlikely]] {
return;
}
auto &query_tracker = transaction_id_to_tracker->tracker;
query_tracker.TrackAlloc(size);
}
void QueriesMemoryControl::TrackFreeOnCurrentThread(size_t size) {
auto thread_id_to_transaction_id_accessor = thread_id_to_transaction_id.access();
// we might be just constructing mapping between thread id and transaction id
// so we miss this allocation
auto thread_id_to_transaction_id_elem = thread_id_to_transaction_id_accessor.find(std::this_thread::get_id());
if (thread_id_to_transaction_id_elem == thread_id_to_transaction_id_accessor.end()) {
return;
}
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
auto transaction_id_to_tracker =
transaction_id_to_tracker_accessor.find(thread_id_to_transaction_id_elem->transaction_id);
// It can happen that some allocation happens between mapping thread to
// transaction id, so we miss this allocation
if (transaction_id_to_tracker == transaction_id_to_tracker_accessor.end()) [[unlikely]] {
return;
}
auto &query_tracker = transaction_id_to_tracker->tracker;
query_tracker.TrackFree(size);
}
void QueriesMemoryControl::CreateTransactionIdTracker(uint64_t transaction_id, size_t inital_limit) {
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
auto [elem, result] = transaction_id_to_tracker_accessor.insert({transaction_id, utils::MemoryTracker{}});
auto [elem, result] = transaction_id_to_tracker_accessor.insert({transaction_id, utils::QueryMemoryTracker{}});
elem->tracker.SetMaximumHardLimit(inital_limit);
elem->tracker.SetHardLimit(inital_limit);
}
bool QueriesMemoryControl::CheckTransactionIdTrackerExists(uint64_t transaction_id) {
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
return transaction_id_to_tracker_accessor.contains(transaction_id);
elem->tracker.SetQueryLimit(inital_limit);
}
bool QueriesMemoryControl::EraseTransactionIdTracker(uint64_t transaction_id) {
@@ -102,6 +127,45 @@ void QueriesMemoryControl::InitializeArenaCounter(unsigned arena_ind) {
arena_tracking[arena_ind].store(0, std::memory_order_relaxed);
}
bool QueriesMemoryControl::CheckTransactionIdTrackerExists(uint64_t transaction_id) {
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
return transaction_id_to_tracker_accessor.contains(transaction_id);
}
void QueriesMemoryControl::TryCreateTransactionProcTracker(uint64_t transaction_id, int64_t procedure_id,
size_t limit) {
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
auto query_tracker = transaction_id_to_tracker_accessor.find(transaction_id);
if (query_tracker == transaction_id_to_tracker_accessor.end()) {
return;
}
query_tracker->tracker.TryCreateProcTracker(procedure_id, limit);
}
void QueriesMemoryControl::SetActiveProcIdTracker(uint64_t transaction_id, int64_t procedure_id) {
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
auto query_tracker = transaction_id_to_tracker_accessor.find(transaction_id);
if (query_tracker == transaction_id_to_tracker_accessor.end()) {
return;
}
query_tracker->tracker.SetActiveProc(procedure_id);
}
void QueriesMemoryControl::PauseProcedureTracking(uint64_t transaction_id) {
auto transaction_id_to_tracker_accessor = transaction_id_to_tracker.access();
auto query_tracker = transaction_id_to_tracker_accessor.find(transaction_id);
if (query_tracker == transaction_id_to_tracker_accessor.end()) {
return;
}
query_tracker->tracker.StopProcTracking();
}
#endif
void StartTrackingCurrentThreadTransaction(uint64_t transaction_id) {
@@ -137,4 +201,29 @@ void TryStopTrackingOnTransaction(uint64_t transaction_id) {
#endif
}
#if USE_JEMALLOC
bool IsTransactionTracked(uint64_t transaction_id) {
return GetQueriesMemoryControl().CheckTransactionIdTrackerExists(transaction_id);
}
#else
bool IsTransactionTracked(uint64_t /*transaction_id*/) { return false; }
#endif
void CreateOrContinueProcedureTracking(uint64_t transaction_id, int64_t procedure_id, size_t limit) {
#if USE_JEMALLOC
if (!GetQueriesMemoryControl().CheckTransactionIdTrackerExists(transaction_id)) {
LOG_FATAL("Memory tracker for transaction was not set");
}
GetQueriesMemoryControl().TryCreateTransactionProcTracker(transaction_id, procedure_id, limit);
GetQueriesMemoryControl().SetActiveProcIdTracker(transaction_id, procedure_id);
#endif
}
void PauseProcedureTracking(uint64_t transaction_id) {
#if USE_JEMALLOC
GetQueriesMemoryControl().PauseProcedureTracking(transaction_id);
#endif
}
} // namespace memgraph::memory

37
src/memory/query_memory_control.hpp
View File

@@ -16,10 +16,13 @@
#include <unordered_map>
#include "utils/memory_tracker.hpp"
#include "utils/query_memory_tracker.hpp"
#include "utils/skip_list.hpp"
namespace memgraph::memory {
static constexpr int64_t UNLIMITED_MEMORY{0};
#if USE_JEMALLOC
// Track memory allocations per query.
@@ -75,16 +78,21 @@ class QueriesMemoryControl {
// Important to reset if one thread gets reused for different transaction
void EraseThreadToTransactionId(const std::thread::id &, uint64_t);
// C-API functionality for thread to transaction mapping
void UpdateThreadToTransactionId(const char *, uint64_t);
// Find tracker for current thread if exists, track
// query allocation and procedure allocation if
// necessary
void TrackAllocOnCurrentThread(size_t size);
// C-API functionality for thread to transaction unmapping
void EraseThreadToTransactionId(const char *, uint64_t);
// Find tracker for current thread if exists, track
// query allocation and procedure allocation if
// necessary
void TrackFreeOnCurrentThread(size_t size);
// Get tracker to current thread if exists, otherwise return
// nullptr. This can happen only if tracker is still
// being constructed.
utils::MemoryTracker *GetTrackerCurrentThread();
void TryCreateTransactionProcTracker(uint64_t, int64_t, size_t);
void SetActiveProcIdTracker(uint64_t, int64_t);
void PauseProcedureTracking(uint64_t);
private:
std::unordered_map<unsigned, std::atomic<int>> arena_tracking;
@@ -102,7 +110,7 @@ class QueriesMemoryControl {
struct TransactionIdToTracker {
uint64_t transaction_id;
utils::MemoryTracker tracker;
utils::QueryMemoryTracker tracker;
bool operator<(const TransactionIdToTracker &other) const { return transaction_id < other.transaction_id; }
bool operator==(const TransactionIdToTracker &other) const { return transaction_id == other.transaction_id; }
@@ -138,4 +146,15 @@ void TryStartTrackingOnTransaction(uint64_t transaction_id, size_t limit);
// Does nothing if jemalloc is not enabled. Does nothing if tracker doesn't exist
void TryStopTrackingOnTransaction(uint64_t transaction_id);
// Is transaction with given id tracked in memory tracker
bool IsTransactionTracked(uint64_t transaction_id);
// Creates tracker on procedure if doesn't exist. Sets query tracker
// to track procedure with id.
void CreateOrContinueProcedureTracking(uint64_t transaction_id, int64_t procedure_id, size_t limit);
// Pauses procedure tracking. This enables to continue
// tracking on procedure once procedure execution resumes.
void PauseProcedureTracking(uint64_t transaction_id);
} // namespace memgraph::memory

51
src/query/plan/operator.cpp
View File

@@ -27,6 +27,7 @@
#include <cppitertools/chain.hpp>
#include <cppitertools/imap.hpp>
#include "memory/query_memory_control.hpp"
#include "query/common.hpp"
#include "spdlog/spdlog.h"
@@ -57,6 +58,7 @@
#include "utils/logging.hpp"
#include "utils/memory.hpp"
#include "utils/message.hpp"
#include "utils/on_scope_exit.hpp"
#include "utils/pmr/deque.hpp"
#include "utils/pmr/list.hpp"
#include "utils/pmr/unordered_map.hpp"
@@ -4617,7 +4619,7 @@ UniqueCursorPtr OutputTableStream::MakeCursor(utils::MemoryResource *mem) const
CallProcedure::CallProcedure(std::shared_ptr<LogicalOperator> input, std::string name, std::vector<Expression *> args,
std::vector<std::string> fields, std::vector<Symbol> symbols, Expression *memory_limit,
size_t memory_scale, bool is_write, bool void_procedure)
size_t memory_scale, bool is_write, int64_t procedure_id, bool void_procedure)
: input_(input ? input : std::make_shared<Once>()),
procedure_name_(name),
arguments_(args),
@@ -4626,6 +4628,7 @@ CallProcedure::CallProcedure(std::shared_ptr<LogicalOperator> input, std::string
memory_limit_(memory_limit),
memory_scale_(memory_scale),
is_write_(is_write),
procedure_id_(procedure_id),
void_procedure_(void_procedure) {}
ACCEPT_WITH_INPUT(CallProcedure);
@@ -4654,7 +4657,7 @@ namespace {
void CallCustomProcedure(const std::string_view fully_qualified_procedure_name, const mgp_proc &proc,
const std::vector<Expression *> &args, mgp_graph &graph, ExpressionEvaluator *evaluator,
utils::MemoryResource *memory, std::optional<size_t> memory_limit, mgp_result *result,
const bool call_initializer = false) {
int64_t procedure_id, uint64_t transaction_id, const bool call_initializer = false) {
static_assert(std::uses_allocator_v<mgp_value, utils::Allocator<mgp_value>>,
"Expected mgp_value to use custom allocator and makes STL "
"containers aware of that");
@@ -4686,12 +4689,46 @@ void CallCustomProcedure(const std::string_view fully_qualified_procedure_name,
if (memory_limit) {
SPDLOG_INFO("Running '{}' with memory limit of {}", fully_qualified_procedure_name,
utils::GetReadableSize(*memory_limit));
utils::LimitedMemoryResource limited_mem(memory, *memory_limit);
mgp_memory proc_memory{&limited_mem};
// Only allocations which can leak memory are
// our own mgp object allocations. Jemalloc can track
// memory correctly, but some memory may not be released
// immediately, so we want to give user info on leak still
// considering our allocations
utils::MemoryTrackingResource memory_tracking_resource{memory, *memory_limit};
// if we are already tracking, no harm no faul
// if we are not tracking, we need to start now, with unlimited memory
// for query, but limited for procedure
// check if transaction is tracked currently, so we
// can disable tracking on that arena if it is not
// once we are done with procedure tracking
bool is_transaction_tracked = memgraph::memory::IsTransactionTracked(transaction_id);
if (!is_transaction_tracked) {
// start tracking with unlimited limit on query
// which is same as not being tracked at all
memgraph::memory::TryStartTrackingOnTransaction(transaction_id, memgraph::memory::UNLIMITED_MEMORY);
}
memgraph::memory::StartTrackingCurrentThreadTransaction(transaction_id);
// due to mgp_batch_read_proc and mgp_batch_write_proc
// we can return to execution without exhausting whole
// memory. Here we need to update tracking
memgraph::memory::CreateOrContinueProcedureTracking(transaction_id, procedure_id, *memory_limit);
mgp_memory proc_memory{&memory_tracking_resource};
MG_ASSERT(result->signature == &proc.results);
utils::OnScopeExit on_scope_exit{[transaction_id = transaction_id]() {
memgraph::memory::StopTrackingCurrentThreadTransaction(transaction_id);
memgraph::memory::PauseProcedureTracking(transaction_id);
}};
// TODO: What about cross library boundary exceptions? OMG C++?!
proc.cb(&proc_args, &graph, result, &proc_memory);
size_t leaked_bytes = limited_mem.GetAllocatedBytes();
auto leaked_bytes = memory_tracking_resource.GetAllocatedBytes();
if (leaked_bytes > 0U) {
spdlog::warn("Query procedure '{}' leaked {} *tracked* bytes", fully_qualified_procedure_name, leaked_bytes);
}
@@ -4799,8 +4836,10 @@ class CallProcedureCursor : public Cursor {
auto *memory = self_->memory_resource;
auto memory_limit = EvaluateMemoryLimit(evaluator, self_->memory_limit_, self_->memory_scale_);
auto graph = mgp_graph::WritableGraph(*context.db_accessor, graph_view, context);
const auto transaction_id = context.db_accessor->GetTransactionId();
MG_ASSERT(transaction_id.has_value());
CallCustomProcedure(self_->procedure_name_, *proc, self_->arguments_, graph, &evaluator, memory, memory_limit,
result_, call_initializer);
result_, self_->procedure_id_, transaction_id.value(), call_initializer);
if (call_initializer) call_initializer = false;

4
src/query/plan/operator.hpp
View File

@@ -2361,7 +2361,7 @@ class CallProcedure : public memgraph::query::plan::LogicalOperator {
CallProcedure() = default;
CallProcedure(std::shared_ptr<LogicalOperator> input, std::string name, std::vector<Expression *> arguments,
std::vector<std::string> fields, std::vector<Symbol> symbols, Expression *memory_limit,
size_t memory_scale, bool is_write, bool void_procedure = false);
size_t memory_scale, bool is_write, int64_t procedure_id, bool void_procedure = false);
bool Accept(HierarchicalLogicalOperatorVisitor &visitor) override;
UniqueCursorPtr MakeCursor(utils::MemoryResource *) const override;
@@ -2383,6 +2383,7 @@ class CallProcedure : public memgraph::query::plan::LogicalOperator {
Expression *memory_limit_{nullptr};
size_t memory_scale_{1024U};
bool is_write_;
int64_t procedure_id_;
bool void_procedure_;
mutable utils::MonotonicBufferResource monotonic_memory{1024UL * 1024UL};
utils::MemoryResource *memory_resource = &monotonic_memory;
@@ -2405,6 +2406,7 @@ class CallProcedure : public memgraph::query::plan::LogicalOperator {
object->memory_limit_ = memory_limit_ ? memory_limit_->Clone(storage) : nullptr;
object->memory_scale_ = memory_scale_;
object->is_write_ = is_write_;
object->procedure_id_ = procedure_id_;
object->void_procedure_ = void_procedure_;
return object;
}

7
src/query/plan/rule_based_planner.hpp
View File

@@ -12,6 +12,7 @@
/// @file
#pragma once
#include <cstdint>
#include <optional>
#include <variant>
@@ -185,6 +186,10 @@ class RuleBasedPlanner {
uint64_t merge_id = 0;
uint64_t subquery_id = 0;
// procedures need to start from 1
// due to swapping mechanism of procedure
// tracking
uint64_t procedure_id = 1;
for (const auto &clause : single_query_part.remaining_clauses) {
MG_ASSERT(!utils::IsSubtype(*clause, Match::kType), "Unexpected Match in remaining clauses");
@@ -224,7 +229,7 @@ class RuleBasedPlanner {
input_op = std::make_unique<plan::CallProcedure>(
std::move(input_op), call_proc->procedure_name_, call_proc->arguments_, call_proc->result_fields_,
result_symbols, call_proc->memory_limit_, call_proc->memory_scale_, call_proc->is_write_,
call_proc->void_procedure_);
procedure_id++, call_proc->void_procedure_);
} else if (auto *load_csv = utils::Downcast<query::LoadCsv>(clause)) {
const auto &row_sym = context.symbol_table->at(*load_csv->row_var_);
context.bound_symbols.insert(row_sym);

3
src/utils/CMakeLists.txt
View File

@@ -14,7 +14,8 @@ set(utils_src_files
tsc.cpp
system_info.cpp
uuid.cpp
build_info.cpp)
build_info.cpp
query_memory_tracker.cpp)
find_package(Boost REQUIRED)
find_package(fmt REQUIRED)

6
src/utils/memory.hpp
View File

@@ -539,9 +539,9 @@ class SynchronizedPoolResource final : public MemoryResource {
bool DoIsEqual(const MemoryResource &other) const noexcept override { return this == &other; }
};
class LimitedMemoryResource final : public utils::MemoryResource {
class MemoryTrackingResource final : public utils::MemoryResource {
public:
explicit LimitedMemoryResource(utils::MemoryResource *memory, size_t max_allocated_bytes)
explicit MemoryTrackingResource(utils::MemoryResource *memory, size_t max_allocated_bytes)
: memory_(memory), max_allocated_bytes_(max_allocated_bytes) {}
size_t GetAllocatedBytes() const noexcept { return max_allocated_bytes_ - available_bytes_; }
@@ -552,13 +552,11 @@ class LimitedMemoryResource final : public utils::MemoryResource {
size_t available_bytes_{max_allocated_bytes_};
void *DoAllocate(size_t bytes, size_t alignment) override {
if (bytes > available_bytes_) throw utils::BadAlloc("Memory allocation limit exceeded!");
available_bytes_ -= bytes;
return memory_->Allocate(bytes, alignment);
}
void DoDeallocate(void *p, size_t bytes, size_t alignment) override {
MG_ASSERT(available_bytes_ + bytes > available_bytes_, "Failed deallocation");
available_bytes_ += bytes;
return memory_->Deallocate(p, bytes, alignment);
}

3
src/utils/memory_tracker.cpp
View File

@@ -111,7 +111,7 @@ void MemoryTracker::Alloc(const int64_t size) {
const auto current_hard_limit = hard_limit_.load(std::memory_order_relaxed);
if (UNLIKELY(current_hard_limit && will_be > current_hard_limit && MemoryTrackerCanThrow())) {
if (current_hard_limit && will_be > current_hard_limit && MemoryTrackerCanThrow()) [[unlikely]] {
MemoryTracker::OutOfMemoryExceptionBlocker exception_blocker;
amount_.fetch_sub(size, std::memory_order_relaxed);
@@ -121,7 +121,6 @@ void MemoryTracker::Alloc(const int64_t size) {
"use to {}, while the maximum allowed size for allocation is set to {}.",
GetReadableSize(size), GetReadableSize(will_be), GetReadableSize(current_hard_limit)));
}
UpdatePeak(will_be);
}

32
src/utils/memory_tracker.hpp
View File

@@ -25,17 +25,6 @@ class OutOfMemoryException : public utils::BasicException {
};
class MemoryTracker final {
private:
std::atomic<int64_t> amount_{0};
std::atomic<int64_t> peak_{0};
std::atomic<int64_t> hard_limit_{0};
// Maximum possible value of a hard limit. If it's set to 0, no upper bound on the hard limit is set.
int64_t maximum_hard_limit_{0};
void UpdatePeak(int64_t will_be);
static void LogMemoryUsage(int64_t current);
public:
void LogPeakMemoryUsage() const;
@@ -73,6 +62,13 @@ class MemoryTracker final {
void ResetTrackings();
bool IsProcedureTracked();
void SetProcTrackingLimit(size_t limit);
void StartProcTracking();
void StopProcTracking();
// 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 {
@@ -109,6 +105,20 @@ class MemoryTracker final {
private:
static thread_local uint64_t counter_;
};
private:
enum class TrackingMode { DEFAULT, ADDITIONAL_PROC };
TrackingMode tracking_mode_{TrackingMode::DEFAULT};
std::atomic<int64_t> amount_{0};
std::atomic<int64_t> peak_{0};
std::atomic<int64_t> hard_limit_{0};
// Maximum possible value of a hard limit. If it's set to 0, no upper bound on the hard limit is set.
int64_t maximum_hard_limit_{0};
void UpdatePeak(int64_t will_be);
static void LogMemoryUsage(int64_t current);
};
// Global memory tracker which tracks every allocation in the application.

78
src/utils/query_memory_tracker.cpp Normal file
View File

@@ -0,0 +1,78 @@
// 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 "utils/query_memory_tracker.hpp"
#include <atomic>
#include <optional>
#include "memory/query_memory_control.hpp"
#include "utils/memory_tracker.hpp"
namespace memgraph::utils {
void QueryMemoryTracker::TrackAlloc(size_t size) {
if (query_tracker_.has_value()) [[likely]] {
query_tracker_->Alloc(static_cast<int64_t>(size));
}
auto *proc_tracker = GetActiveProc();
if (proc_tracker == nullptr) {
return;
}
proc_tracker->Alloc(static_cast<int64_t>(size));
}
void QueryMemoryTracker::TrackFree(size_t size) {
if (query_tracker_.has_value()) [[likely]] {
query_tracker_->Free(static_cast<int64_t>(size));
}
auto *proc_tracker = GetActiveProc();
if (proc_tracker == nullptr) {
return;
}
proc_tracker->Free(static_cast<int64_t>(size));
}
void QueryMemoryTracker::SetQueryLimit(size_t size) {
if (size == memgraph::memory::UNLIMITED_MEMORY) {
return;
}
InitializeQueryTracker();
query_tracker_->SetMaximumHardLimit(static_cast<int64_t>(size));
query_tracker_->SetHardLimit(static_cast<int64_t>(size));
}
memgraph::utils::MemoryTracker *QueryMemoryTracker::GetActiveProc() {
if (active_proc_id == NO_PROCEDURE) [[likely]] {
return nullptr;
}
return &proc_memory_trackers_[active_proc_id];
}
void QueryMemoryTracker::SetActiveProc(int64_t new_active_proc) { active_proc_id = new_active_proc; }
void QueryMemoryTracker::StopProcTracking() { active_proc_id = QueryMemoryTracker::NO_PROCEDURE; }
void QueryMemoryTracker::TryCreateProcTracker(int64_t procedure_id, size_t limit) {
if (proc_memory_trackers_.contains(procedure_id)) {
return;
}
auto [it, inserted] = proc_memory_trackers_.emplace(procedure_id, utils::MemoryTracker{});
it->second.SetMaximumHardLimit(static_cast<int64_t>(limit));
it->second.SetHardLimit(static_cast<int64_t>(limit));
}
void QueryMemoryTracker::InitializeQueryTracker() { query_tracker_.emplace(MemoryTracker{}); }
} // namespace memgraph::utils

78
src/utils/query_memory_tracker.hpp Normal file
View File

@@ -0,0 +1,78 @@
// 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.
#pragma once
#include <atomic>
#include <optional>
#include <unordered_map>
#include <utility>
#include "utils/memory_tracker.hpp"
namespace memgraph::utils {
// Class which handles tracking of query memory
// together with procedure memory. Default
// active procedure id is set to -1, as
// procedures start from id 1.
// This class is meant to be used in environment
// where we don't execute multiple procedures in parallel
// but procedure can have multiple threads which are doing allocations
class QueryMemoryTracker {
public:
QueryMemoryTracker() = default;
QueryMemoryTracker(QueryMemoryTracker &&other) noexcept
: query_tracker_(std::move(other.query_tracker_)),
proc_memory_trackers_(std::move(other.proc_memory_trackers_)),
active_proc_id(other.active_proc_id) {
other.active_proc_id = NO_PROCEDURE;
}
QueryMemoryTracker(const QueryMemoryTracker &other) = delete;
QueryMemoryTracker &operator=(QueryMemoryTracker &&other) = delete;
QueryMemoryTracker &operator=(const QueryMemoryTracker &other) = delete;
~QueryMemoryTracker() = default;
// Track allocation on query and procedure if active
void TrackAlloc(size_t);
// Track Free on query and procedure if active
void TrackFree(size_t);
// Set query limit
void SetQueryLimit(size_t);
// Create proc tracker if doesn't exist
void TryCreateProcTracker(int64_t, size_t);
// Set currently active procedure
void SetActiveProc(int64_t);
// Stop procedure tracking
void StopProcTracking();
private:
static constexpr int64_t NO_PROCEDURE{-1};
void InitializeQueryTracker();
std::optional<memgraph::utils::MemoryTracker> query_tracker_{std::nullopt};
std::unordered_map<int64_t, memgraph::utils::MemoryTracker> proc_memory_trackers_;
// Procedure ids start from 1. Procedure id -1 means there is no procedure
// to track.
int64_t active_proc_id{NO_PROCEDURE};
memgraph::utils::MemoryTracker *GetActiveProc();
};
} // namespace memgraph::utils