From 524acb17a191b404c988efd42b2eb010876076d2 Mon Sep 17 00:00:00 2001
From: Kostas Kyrimis <kostaskyrim@gmail.com>
Date: Mon, 7 Jun 2021 15:45:05 +0300
Subject: [PATCH] Add global version allocators for C in query modules (#162)
---
include/mg_procedure.h | 133 ++++++++----------
src/query/procedure/mg_procedure_impl.cpp | 37 +++--
src/query/procedure/module.cpp | 3 +
src/query/procedure/module.hpp | 4 +
tests/e2e/memory/CMakeLists.txt | 9 ++
.../e2e/memory/memory_limit_global_alloc.cpp | 26 ++++
.../memory/memory_limit_global_alloc_proc.cpp | 31 ++++
tests/e2e/memory/procedures/CMakeLists.txt | 5 +
.../memory/procedures/global_memory_limit.c | 36 +++++
.../procedures/global_memory_limit_proc.c | 63 +++++++++
tests/e2e/memory/workloads.yaml | 10 ++
tests/e2e/runner.py | 4 +
12 files changed, 279 insertions(+), 82 deletions(-)
create mode 100644 tests/e2e/memory/memory_limit_global_alloc.cpp
create mode 100644 tests/e2e/memory/memory_limit_global_alloc_proc.cpp
create mode 100644 tests/e2e/memory/procedures/CMakeLists.txt
create mode 100644 tests/e2e/memory/procedures/global_memory_limit.c
create mode 100644 tests/e2e/memory/procedures/global_memory_limit_proc.c
diff --git a/include/mg_procedure.h b/include/mg_procedure.h
index 2a2567aba..3b0168bcb 100644
--- a/include/mg_procedure.h
+++ b/include/mg_procedure.h
@@ -17,10 +17,11 @@ extern "C" {
/// addition to efficiency, Memgraph can set the limit on allowed allocations
/// thus providing some safety with regards to memory usage. The allocated
/// memory is only valid during the execution of mgp_main. You must not allocate
-/// global resources with these functions. None of the functions are
+/// global resources with these functions and none of the functions are
/// thread-safe, because we provide a single thread of execution when invoking a
-/// custom procedure. This allows Memgraph to be more efficient as stated
-/// before.
+/// custom procedure. For allocating global resources, you can use the _global
+/// variations of the aforementioned allocators. This allows Memgraph to be
+/// more efficient as explained before.
///@{
/// Provides memory managament access and state.
@@ -39,8 +40,7 @@ void *mgp_alloc(struct mgp_memory *memory, size_t size_in_bytes);
/// `alignment` must be a power of 2 value.
/// The returned pointer must be freed with mgp_free.
/// NULL is returned if unable to serve the requested allocation.
-void *mgp_aligned_alloc(struct mgp_memory *memory, size_t size_in_bytes,
- size_t alignment);
+void *mgp_aligned_alloc(struct mgp_memory *memory, size_t size_in_bytes, size_t alignment);
/// Deallocate an allocation from mgp_alloc or mgp_aligned_alloc.
/// Unlike free, this function is not thread-safe.
@@ -48,6 +48,26 @@ void *mgp_aligned_alloc(struct mgp_memory *memory, size_t size_in_bytes,
/// The behavior is undefined if `ptr` is not a value returned from a prior
/// mgp_alloc or mgp_aligned_alloc call with the corresponding `memory`.
void mgp_free(struct mgp_memory *memory, void *ptr);
+
+/// Allocate a global block of memory with given size in bytes.
+/// This function can be used to allocate global memory that persists
+/// beyond a single invocation of mgp_main.
+/// The returned pointer must be freed with mgp_global_free.
+/// NULL is returned if unable to serve the requested allocation.
+void *mgp_global_alloc(size_t size_in_bytes);
+
+/// Allocate an aligned global block of memory with given size in bytes.
+/// This function can be used to allocate global memory that persists
+/// beyond a single invocation of mgp_main.
+/// The returned pointer must be freed with mgp_global_free.
+/// NULL is returned if unable to serve the requested allocation.
+void *mgp_global_aligned_alloc(size_t size_in_bytes, size_t alignment);
+
+/// Deallocate an allocation from mgp_global_alloc or mgp_global_aligned_alloc.
+/// If `ptr` is NULL, this function does nothing.
+/// 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);
///@}
/// @name Operations on mgp_value
@@ -119,8 +139,7 @@ struct mgp_value *mgp_value_make_double(double val, struct mgp_memory *memory);
/// Construct a character string value from a NULL terminated string.
/// You need to free the instance through mgp_value_destroy.
/// NULL is returned if unable to allocate a mgp_value.
-struct mgp_value *mgp_value_make_string(const char *val,
- struct mgp_memory *memory);
+struct mgp_value *mgp_value_make_string(const char *val, struct mgp_memory *memory);
/// Create a mgp_value storing a mgp_list.
/// You need to free the instance through mgp_value_destroy. The ownership of
@@ -238,8 +257,7 @@ const struct mgp_path *mgp_value_get_path(const struct mgp_value *val);
/// of mgp_value, but it will not contain any elements. Therefore,
/// mgp_list_size will return 0.
/// NULL is returned if unable to allocate a new list.
-struct mgp_list *mgp_list_make_empty(size_t capacity,
- struct mgp_memory *memory);
+struct mgp_list *mgp_list_make_empty(size_t capacity, struct mgp_memory *memory);
/// Free the memory used by the given mgp_list and contained elements.
void mgp_list_destroy(struct mgp_list *list);
@@ -288,8 +306,7 @@ void mgp_map_destroy(struct mgp_map *map);
/// you still need to free their memory explicitly.
/// Return non-zero on success, or 0 if there's no memory to insert a new
/// mapping or a previous mapping already exists.
-int mgp_map_insert(struct mgp_map *map, const char *key,
- const struct mgp_value *value);
+int mgp_map_insert(struct mgp_map *map, const char *key, const struct mgp_value *value);
/// Return the number of items stored in mgp_map.
size_t mgp_map_size(const struct mgp_map *map);
@@ -314,8 +331,7 @@ struct mgp_map_items_iterator;
/// The returned mgp_map_items_iterator needs to be deallocated with
/// mgp_map_items_iterator_destroy.
/// NULL is returned if unable to allocate a new iterator.
-struct mgp_map_items_iterator *mgp_map_iter_items(const struct mgp_map *map,
- struct mgp_memory *memory);
+struct mgp_map_items_iterator *mgp_map_iter_items(const struct mgp_map *map, struct mgp_memory *memory);
/// Deallocate memory used by mgp_map_items_iterator.
void mgp_map_items_iterator_destroy(struct mgp_map_items_iterator *it);
@@ -328,27 +344,23 @@ void mgp_map_items_iterator_destroy(struct mgp_map_items_iterator *it);
/// as the value before, and use them after invoking
/// mgp_map_items_iterator_next.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_map_item *mgp_map_items_iterator_get(
- const struct mgp_map_items_iterator *it);
+const struct mgp_map_item *mgp_map_items_iterator_get(const struct mgp_map_items_iterator *it);
/// Advance the iterator to the next item stored in map and return it.
/// The previous pointer obtained through mgp_map_items_iterator_get will
/// be invalidated, but the pointers to key and value will remain valid.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_map_item *mgp_map_items_iterator_next(
- struct mgp_map_items_iterator *it);
+const struct mgp_map_item *mgp_map_items_iterator_next(struct mgp_map_items_iterator *it);
/// Create a path with the copy of the given starting vertex.
/// You need to free the created instance with mgp_path_destroy.
/// NULL is returned if unable to allocate a path.
-struct mgp_path *mgp_path_make_with_start(const struct mgp_vertex *vertex,
- struct mgp_memory *memory);
+struct mgp_path *mgp_path_make_with_start(const struct mgp_vertex *vertex, struct mgp_memory *memory);
/// Copy a mgp_path.
/// Returned pointer must be freed with mgp_path_destroy.
/// NULL is returned if unable to allocate a mgp_path.
-struct mgp_path *mgp_path_copy(const struct mgp_path *path,
- struct mgp_memory *memory);
+struct mgp_path *mgp_path_copy(const struct mgp_path *path, struct mgp_memory *memory);
/// Free the memory used by the given mgp_path and contained vertices and edges.
void mgp_path_destroy(struct mgp_path *path);
@@ -370,14 +382,12 @@ size_t mgp_path_size(const struct mgp_path *path);
/// Return the vertex from a path at given index.
/// The valid index range is [0, mgp_path_size].
/// NULL is returned if index is out of range.
-const struct mgp_vertex *mgp_path_vertex_at(const struct mgp_path *path,
- size_t index);
+const struct mgp_vertex *mgp_path_vertex_at(const struct mgp_path *path, size_t index);
/// Return the edge from a path at given index.
/// The valid index range is [0, mgp_path_size - 1].
/// NULL is returned if index is out of range.
-const struct mgp_edge *mgp_path_edge_at(const struct mgp_path *path,
- size_t index);
+const struct mgp_edge *mgp_path_edge_at(const struct mgp_path *path, size_t index);
/// Return non-zero if given paths are equal, otherwise 0.
int mgp_path_equal(const struct mgp_path *p1, const struct mgp_path *p2);
@@ -408,9 +418,7 @@ struct mgp_result_record *mgp_result_new_record(struct mgp_result *res);
/// Return 0 if there's no memory to copy the mgp_value to mgp_result_record or
/// if the combination of `field_name` and `val` does not satisfy the
/// procedure's result signature.
-int mgp_result_record_insert(struct mgp_result_record *record,
- const char *field_name,
- const struct mgp_value *val);
+int mgp_result_record_insert(struct mgp_result_record *record, const char *field_name, const struct mgp_value *val);
///@}
/// @name Graph Constructs
@@ -446,15 +454,13 @@ struct mgp_property {
/// When the mgp_properties_iterator_next is invoked, the previous
/// mgp_property is invalidated and its value must not be used.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_property *mgp_properties_iterator_get(
- const struct mgp_properties_iterator *it);
+const struct mgp_property *mgp_properties_iterator_get(const struct mgp_properties_iterator *it);
/// Advance the iterator to the next property and return it.
/// The previous mgp_property obtained through mgp_properties_iterator_get
/// will be invalidated, and you must not use its value.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_property *mgp_properties_iterator_next(
- struct mgp_properties_iterator *it);
+const struct mgp_property *mgp_properties_iterator_next(struct mgp_properties_iterator *it);
/// Iterator over edges of a vertex.
struct mgp_edges_iterator;
@@ -475,8 +481,7 @@ struct mgp_vertex_id mgp_vertex_get_id(const struct mgp_vertex *v);
/// Copy a mgp_vertex.
/// Returned pointer must be freed with mgp_vertex_destroy.
/// NULL is returned if unable to allocate a mgp_vertex.
-struct mgp_vertex *mgp_vertex_copy(const struct mgp_vertex *v,
- struct mgp_memory *memory);
+struct mgp_vertex *mgp_vertex_copy(const struct mgp_vertex *v, struct mgp_memory *memory);
/// Free the memory used by a mgp_vertex.
void mgp_vertex_destroy(struct mgp_vertex *v);
@@ -495,43 +500,37 @@ struct mgp_label mgp_vertex_label_at(const struct mgp_vertex *v, size_t index);
int mgp_vertex_has_label(const struct mgp_vertex *v, struct mgp_label label);
/// Return non-zero if the given vertex has a label with given name.
-int mgp_vertex_has_label_named(const struct mgp_vertex *v,
- const char *label_name);
+int mgp_vertex_has_label_named(const struct mgp_vertex *v, const char *label_name);
/// Get a copy of a vertex property mapped to a given name.
/// Returned value must be freed with mgp_value_destroy.
/// NULL is returned if unable to allocate a mgp_value.
-struct mgp_value *mgp_vertex_get_property(const struct mgp_vertex *v,
- const char *property_name,
+struct mgp_value *mgp_vertex_get_property(const struct mgp_vertex *v, const char *property_name,
struct mgp_memory *memory);
/// Start iterating over properties stored in the given vertex.
/// The returned mgp_properties_iterator needs to be deallocated with
/// mgp_properties_iterator_destroy.
/// NULL is returned if unable to allocate a new iterator.
-struct mgp_properties_iterator *mgp_vertex_iter_properties(
- const struct mgp_vertex *v, struct mgp_memory *memory);
+struct mgp_properties_iterator *mgp_vertex_iter_properties(const struct mgp_vertex *v, struct mgp_memory *memory);
/// Start iterating over inbound edges of the given vertex.
/// The returned mgp_edges_iterator needs to be deallocated with
/// mgp_edges_iterator_destroy.
/// NULL is returned if unable to allocate a new iterator.
-struct mgp_edges_iterator *mgp_vertex_iter_in_edges(const struct mgp_vertex *v,
- struct mgp_memory *memory);
+struct mgp_edges_iterator *mgp_vertex_iter_in_edges(const struct mgp_vertex *v, struct mgp_memory *memory);
/// Start iterating over outbound edges of the given vertex.
/// The returned mgp_edges_iterator needs to be deallocated with
/// mgp_edges_iterator_destroy.
/// NULL is returned if unable to allocate a new iterator.
-struct mgp_edges_iterator *mgp_vertex_iter_out_edges(const struct mgp_vertex *v,
- struct mgp_memory *memory);
+struct mgp_edges_iterator *mgp_vertex_iter_out_edges(const struct mgp_vertex *v, struct mgp_memory *memory);
/// Get the current edge pointed to by the iterator.
/// When the mgp_edges_iterator_next is invoked, the previous
/// mgp_edge is invalidated and its value must not be used.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_edge *mgp_edges_iterator_get(
- const struct mgp_edges_iterator *it);
+const struct mgp_edge *mgp_edges_iterator_get(const struct mgp_edges_iterator *it);
/// Advance the iterator to the next edge and return it.
/// The previous mgp_edge obtained through mgp_edges_iterator_get
@@ -552,8 +551,7 @@ struct mgp_edge_id mgp_edge_get_id(const struct mgp_edge *e);
/// Copy a mgp_edge.
/// Returned pointer must be freed with mgp_edge_destroy.
/// NULL is returned if unable to allocate a mgp_edge.
-struct mgp_edge *mgp_edge_copy(const struct mgp_edge *e,
- struct mgp_memory *memory);
+struct mgp_edge *mgp_edge_copy(const struct mgp_edge *e, struct mgp_memory *memory);
/// Free the memory used by a mgp_edge.
void mgp_edge_destroy(struct mgp_edge *e);
@@ -573,16 +571,13 @@ const struct mgp_vertex *mgp_edge_get_to(const struct mgp_edge *e);
/// Get a copy of a edge property mapped to a given name.
/// Returned value must be freed with mgp_value_destroy.
/// NULL is returned if unable to allocate a mgp_value.
-struct mgp_value *mgp_edge_get_property(const struct mgp_edge *e,
- const char *property_name,
- struct mgp_memory *memory);
+struct mgp_value *mgp_edge_get_property(const struct mgp_edge *e, const char *property_name, struct mgp_memory *memory);
/// Start iterating over properties stored in the given edge.
/// The returned mgp_properties_iterator needs to be deallocated with
/// mgp_properties_iterator_destroy.
/// NULL is returned if unable to allocate a new iterator.
-struct mgp_properties_iterator *mgp_edge_iter_properties(
- const struct mgp_edge *e, struct mgp_memory *memory);
+struct mgp_properties_iterator *mgp_edge_iter_properties(const struct mgp_edge *e, struct mgp_memory *memory);
/// State of the graph database.
struct mgp_graph;
@@ -590,8 +585,7 @@ struct mgp_graph;
/// Return the vertex corresponding to given ID.
/// The returned vertex must be freed using mgp_vertex_destroy.
/// NULL is returned if unable to allocate the vertex or if ID is not valid.
-struct mgp_vertex *mgp_graph_get_vertex_by_id(const struct mgp_graph *g,
- struct mgp_vertex_id id,
+struct mgp_vertex *mgp_graph_get_vertex_by_id(const struct mgp_graph *g, struct mgp_vertex_id id,
struct mgp_memory *memory);
/// Iterator over vertices.
@@ -604,22 +598,19 @@ void mgp_vertices_iterator_destroy(struct mgp_vertices_iterator *it);
/// The returned mgp_vertices_iterator needs to be deallocated with
/// mgp_vertices_iterator_destroy.
/// NULL is returned if unable to allocate a new iterator.
-struct mgp_vertices_iterator *mgp_graph_iter_vertices(
- const struct mgp_graph *g, struct mgp_memory *memory);
+struct mgp_vertices_iterator *mgp_graph_iter_vertices(const struct mgp_graph *g, struct mgp_memory *memory);
/// Get the current vertex pointed to by the iterator.
/// When the mgp_vertices_iterator_next is invoked, the previous
/// mgp_vertex is invalidated and its value must not be used.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_vertex *mgp_vertices_iterator_get(
- const struct mgp_vertices_iterator *it);
+const struct mgp_vertex *mgp_vertices_iterator_get(const struct mgp_vertices_iterator *it);
/// Advance the iterator to the next vertex and return it.
/// The previous mgp_vertex obtained through mgp_vertices_iterator_get
/// will be invalidated, and you must not use its value.
/// NULL is returned if the end of the iteration has been reached.
-const struct mgp_vertex *mgp_vertices_iterator_next(
- struct mgp_vertices_iterator *it);
+const struct mgp_vertex *mgp_vertices_iterator_next(struct mgp_vertices_iterator *it);
///@}
/// @name Type System
@@ -718,8 +709,8 @@ struct mgp_proc;
/// Passed in arguments will not live longer than the callback's execution.
/// Therefore, you must not store them globally or use the passed in mgp_memory
/// to allocate global resources.
-typedef void (*mgp_proc_cb)(const struct mgp_list *, const struct mgp_graph *,
- struct mgp_result *, struct mgp_memory *);
+typedef void (*mgp_proc_cb)(const struct mgp_list *, const struct mgp_graph *, struct mgp_result *,
+ struct mgp_memory *);
/// Register a read-only procedure with a module.
///
@@ -730,9 +721,7 @@ typedef void (*mgp_proc_cb)(const struct mgp_list *, const struct mgp_graph *,
///
/// NULL is returned if unable to allocate memory for mgp_proc; if `name` is
/// not valid or a procedure with the same name was already registered.
-struct mgp_proc *mgp_module_add_read_procedure(struct mgp_module *module,
- const char *name,
- mgp_proc_cb cb);
+struct mgp_proc *mgp_module_add_read_procedure(struct mgp_module *module, const char *name, mgp_proc_cb cb);
/// Add a required argument to a procedure.
///
@@ -748,8 +737,7 @@ struct mgp_proc *mgp_module_add_read_procedure(struct mgp_module *module,
/// 0 is returned if unable to allocate memory for an argument; if invoking this
/// function after setting an optional argument or if `name` is not valid.
/// Non-zero is returned on success.
-int mgp_proc_add_arg(struct mgp_proc *proc, const char *name,
- const struct mgp_type *type);
+int mgp_proc_add_arg(struct mgp_proc *proc, const char *name, const struct mgp_type *type);
/// Add an optional argument with a default value to a procedure.
///
@@ -772,8 +760,7 @@ int mgp_proc_add_arg(struct mgp_proc *proc, const char *name,
/// 0 is returned if unable to allocate memory for an argument; if `name` is
/// not valid or `default_value` does not satisfy `type`. Non-zero is returned
/// on success.
-int mgp_proc_add_opt_arg(struct mgp_proc *proc, const char *name,
- const struct mgp_type *type,
+int mgp_proc_add_opt_arg(struct mgp_proc *proc, const char *name, const struct mgp_type *type,
const struct mgp_value *default_value);
/// Add a result field to a procedure.
@@ -787,15 +774,13 @@ int mgp_proc_add_opt_arg(struct mgp_proc *proc, const char *name,
/// 0 is returned if unable to allocate memory for a result field; if
/// `name` is not valid or if a result field with the same name was already
/// added. Non-zero is returned on success.
-int mgp_proc_add_result(struct mgp_proc *proc, const char *name,
- const struct mgp_type *type);
+int mgp_proc_add_result(struct mgp_proc *proc, const char *name, const struct mgp_type *type);
/// Add a result field to a procedure and mark it as deprecated.
///
/// This is the same as mgp_proc_add_result, but the result field will be marked
/// as deprecated.
-int mgp_proc_add_deprecated_result(struct mgp_proc *proc, const char *name,
- const struct mgp_type *type);
+int mgp_proc_add_deprecated_result(struct mgp_proc *proc, const char *name, const struct mgp_type *type);
///@}
/// @name Execution
diff --git a/src/query/procedure/mg_procedure_impl.cpp b/src/query/procedure/mg_procedure_impl.cpp
index 98e7d66aa..e65f9437a 100644
--- a/src/query/procedure/mg_procedure_impl.cpp
+++ b/src/query/procedure/mg_procedure_impl.cpp
@@ -6,21 +6,20 @@
#include <regex>
#include <type_traits>
+#include "module.hpp"
#include "utils/algorithm.hpp"
#include "utils/logging.hpp"
#include "utils/math.hpp"
+#include "utils/memory.hpp"
#include "utils/string.hpp"
-
// This file contains implementation of top level C API functions, but this is
// all actually part of query::procedure. So use that namespace for simplicity.
// NOLINTNEXTLINE(google-build-using-namespace)
using namespace query::procedure;
-void *mgp_alloc(mgp_memory *memory, size_t size_in_bytes) {
- return mgp_aligned_alloc(memory, size_in_bytes, alignof(std::max_align_t));
-}
+namespace {
-void *mgp_aligned_alloc(mgp_memory *memory, const size_t size_in_bytes, const size_t alignment) {
+void *MgpAlignedAllocImpl(utils::MemoryResource &memory, const size_t size_in_bytes, const size_t alignment) {
if (size_in_bytes == 0U || !utils::IsPow2(alignment)) return nullptr;
// Simplify alignment by always using values greater or equal to max_align.
const size_t alloc_align = std::max(alignment, alignof(std::max_align_t));
@@ -37,7 +36,7 @@ void *mgp_aligned_alloc(mgp_memory *memory, const size_t size_in_bytes, const si
const size_t alloc_size = bytes_for_header + size_in_bytes;
if (alloc_size < size_in_bytes) return nullptr;
try {
- void *ptr = memory->impl->Allocate(alloc_size, alloc_align);
+ void *ptr = memory.Allocate(alloc_size, alloc_align);
char *data = reinterpret_cast<char *>(ptr) + bytes_for_header;
std::memcpy(data - sizeof(size_in_bytes), &size_in_bytes, sizeof(size_in_bytes));
std::memcpy(data - sizeof(size_in_bytes) - sizeof(alloc_align), &alloc_align, sizeof(alloc_align));
@@ -47,7 +46,7 @@ void *mgp_aligned_alloc(mgp_memory *memory, const size_t size_in_bytes, const si
}
}
-void mgp_free(mgp_memory *memory, void *const p) {
+void MgpFreeImpl(utils::MemoryResource &memory, void *const p) {
if (!p) return;
char *const data = reinterpret_cast<char *>(p);
// Read the header containing size & alignment info.
@@ -63,9 +62,31 @@ void mgp_free(mgp_memory *memory, void *const p) {
const size_t alloc_size = bytes_for_header + size_in_bytes;
// Get the original ptr we allocated.
void *const original_ptr = data - bytes_for_header;
- memory->impl->Deallocate(original_ptr, alloc_size, alloc_align);
+ memory.Deallocate(original_ptr, alloc_size, alloc_align);
}
+} // namespace
+
+void *mgp_alloc(mgp_memory *memory, size_t size_in_bytes) {
+ return mgp_aligned_alloc(memory, size_in_bytes, alignof(std::max_align_t));
+}
+
+void *mgp_aligned_alloc(mgp_memory *memory, const size_t size_in_bytes, const size_t alignment) {
+ return MgpAlignedAllocImpl(*memory->impl, size_in_bytes, alignment);
+}
+
+void mgp_free(mgp_memory *memory, void *const p) { MgpFreeImpl(*memory->impl, p); }
+
+void *mgp_global_alloc(size_t size_in_bytes) {
+ return mgp_global_aligned_alloc(size_in_bytes, alignof(std::max_align_t));
+}
+
+void *mgp_global_aligned_alloc(size_t size_in_bytes, size_t alignment) {
+ return MgpAlignedAllocImpl(gModuleRegistry.GetSharedMemoryResource(), size_in_bytes, alignment);
+}
+
+void mgp_global_free(void *const p) { MgpFreeImpl(gModuleRegistry.GetSharedMemoryResource(), p); }
+
namespace {
// May throw whatever the constructor of U throws. `std::bad_alloc` is handled
diff --git a/src/query/procedure/module.cpp b/src/query/procedure/module.cpp
index 99bc54dc9..a0c9a4103 100644
--- a/src/query/procedure/module.cpp
+++ b/src/query/procedure/module.cpp
@@ -1,4 +1,5 @@
#include "query/procedure/module.hpp"
+#include "utils/memory.hpp"
extern "C" {
#include <dlfcn.h>
@@ -478,6 +479,8 @@ void ModuleRegistry::UnloadAllModules() {
DoUnloadAllModules();
}
+utils::MemoryResource &ModuleRegistry::GetSharedMemoryResource() { return *shared_; }
+
std::optional<std::pair<procedure::ModulePtr, const mgp_proc *>> FindProcedure(
const ModuleRegistry &module_registry, const std::string_view &fully_qualified_procedure_name,
utils::MemoryResource *memory) {
diff --git a/src/query/procedure/module.hpp b/src/query/procedure/module.hpp
index cdae588c8..0a96c7d83 100644
--- a/src/query/procedure/module.hpp
+++ b/src/query/procedure/module.hpp
@@ -52,6 +52,7 @@ class ModulePtr final {
class ModuleRegistry final {
std::map<std::string, std::unique_ptr<Module>, std::less<>> modules_;
mutable utils::RWLock lock_{utils::RWLock::Priority::WRITE};
+ std::unique_ptr<utils::MemoryResource> shared_{std::make_unique<utils::ResourceWithOutOfMemoryException>()};
bool RegisterModule(const std::string_view &name, std::unique_ptr<Module> module);
@@ -96,6 +97,9 @@ class ModuleRegistry final {
/// Takes a write lock.
void UnloadAllModules();
+ /// Returns the shared memory allocator used by modules
+ utils::MemoryResource &GetSharedMemoryResource();
+
private:
std::vector<std::filesystem::path> modules_dirs_;
};
diff --git a/tests/e2e/memory/CMakeLists.txt b/tests/e2e/memory/CMakeLists.txt
index 4e258f61a..95f3a145c 100644
--- a/tests/e2e/memory/CMakeLists.txt
+++ b/tests/e2e/memory/CMakeLists.txt
@@ -1,2 +1,11 @@
+add_subdirectory(procedures)
+
add_executable(memgraph__e2e__memory__control memory_control.cpp)
target_link_libraries(memgraph__e2e__memory__control gflags mgclient mg-utils mg-io Threads::Threads)
+
+add_executable(memgraph__e2e__memory__limit_global_alloc memory_limit_global_alloc.cpp)
+target_link_libraries(memgraph__e2e__memory__limit_global_alloc gflags mgclient mg-utils mg-io Threads::Threads)
+
+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)
+
diff --git a/tests/e2e/memory/memory_limit_global_alloc.cpp b/tests/e2e/memory/memory_limit_global_alloc.cpp
new file mode 100644
index 000000000..0d2094fea
--- /dev/null
+++ b/tests/e2e/memory/memory_limit_global_alloc.cpp
@@ -0,0 +1,26 @@
+#include <gflags/gflags.h>
+#include <mgclient.hpp>
+
+#include "utils/logging.hpp"
+#include "utils/timer.hpp"
+
+DEFINE_uint64(bolt_port, 7687, "Bolt port");
+DEFINE_uint64(timeout, 120, "Timeout seconds");
+
+int main(int argc, char **argv) {
+ google::SetUsageMessage("Memgraph E2E Memory Limit For Global Allocators");
+ gflags::ParseCommandLineFlags(&argc, &argv, true);
+ 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!");
+ }
+
+ bool result = client->Execute("CALL libglobal_memory_limit.procedure() YIELD *");
+ MG_ASSERT(result == false);
+ return 0;
+}
diff --git a/tests/e2e/memory/memory_limit_global_alloc_proc.cpp b/tests/e2e/memory/memory_limit_global_alloc_proc.cpp
new file mode 100644
index 000000000..2e57da9d6
--- /dev/null
+++ b/tests/e2e/memory/memory_limit_global_alloc_proc.cpp
@@ -0,0 +1,31 @@
+#include <gflags/gflags.h>
+#include <mgclient.hpp>
+#include <algorithm>
+
+#include "utils/logging.hpp"
+#include "utils/timer.hpp"
+
+DEFINE_uint64(bolt_port, 7687, "Bolt port");
+DEFINE_uint64(timeout, 120, "Timeout seconds");
+
+int main(int argc, char **argv) {
+ google::SetUsageMessage("Memgraph E2E Memory Limit For Global Allocators");
+ gflags::ParseCommandLineFlags(&argc, &argv, true);
+ 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!");
+ }
+ bool result = client->Execute("CALL libglobal_memory_limit_proc.error() YIELD *");
+ auto result1 = client->FetchAll();
+ MG_ASSERT(result1 != std::nullopt && result1->size() == 0);
+
+ result = client->Execute("CALL libglobal_memory_limit_proc.success() YIELD *");
+ auto result2 = client->FetchAll();
+ MG_ASSERT(result2 != std::nullopt && result2->size() > 0);
+ return 0;
+}
diff --git a/tests/e2e/memory/procedures/CMakeLists.txt b/tests/e2e/memory/procedures/CMakeLists.txt
new file mode 100644
index 000000000..21201e59b
--- /dev/null
+++ b/tests/e2e/memory/procedures/CMakeLists.txt
@@ -0,0 +1,5 @@
+add_library(global_memory_limit SHARED global_memory_limit.c)
+target_include_directories(global_memory_limit PRIVATE ${CMAKE_SOURCE_DIR}/include)
+
+add_library(global_memory_limit_proc SHARED global_memory_limit_proc.c)
+target_include_directories(global_memory_limit_proc PRIVATE ${CMAKE_SOURCE_DIR}/include)
diff --git a/tests/e2e/memory/procedures/global_memory_limit.c b/tests/e2e/memory/procedures/global_memory_limit.c
new file mode 100644
index 000000000..a2003cb59
--- /dev/null
+++ b/tests/e2e/memory/procedures/global_memory_limit.c
@@ -0,0 +1,36 @@
+#include "mg_procedure.h"
+
+int *gVal = NULL;
+
+void set_error(struct mgp_result *result) { mgp_result_set_error_msg(result, "Something went wrong"); }
+
+static void procedure(const struct mgp_list *args, const struct mgp_graph *graph, struct mgp_result *result,
+ struct mgp_memory *memory) {
+ struct mgp_result_record *record = mgp_result_new_record(result);
+ if (record == NULL) return set_error(result);
+
+ struct mgp_value *result_msg = mgp_value_make_string("mgp_init_module allocation works", memory);
+ if (result_msg == NULL) return set_error(result);
+
+ int result_inserted = mgp_result_record_insert(record, "result", result_msg);
+ mgp_value_destroy(result_msg);
+ if (!result_inserted) return set_error(result);
+}
+
+int mgp_init_module(struct mgp_module *module, struct mgp_memory *memory) {
+ const size_t one_gb = 1 << 30;
+ gVal = mgp_global_alloc(one_gb);
+ if (!gVal) return 1;
+
+ struct mgp_proc *proc = mgp_module_add_read_procedure(module, "procedure", procedure);
+ if (!proc) return 1;
+
+ if (!mgp_proc_add_result(proc, "result", mgp_type_string())) return 1;
+
+ return 0;
+}
+
+int mgp_shutdown_module() {
+ if (gVal) mgp_global_free(gVal);
+ return 0;
+}
diff --git a/tests/e2e/memory/procedures/global_memory_limit_proc.c b/tests/e2e/memory/procedures/global_memory_limit_proc.c
new file mode 100644
index 000000000..519f11f05
--- /dev/null
+++ b/tests/e2e/memory/procedures/global_memory_limit_proc.c
@@ -0,0 +1,63 @@
+#include "mg_procedure.h"
+
+int *gVal = NULL;
+
+void set_error(struct mgp_result *result) { mgp_result_set_error_msg(result, "Something went wrong"); }
+
+void set_out_of_memory_error(struct mgp_result *result) { mgp_result_set_error_msg(result, "Out of memory"); }
+
+static void error(const struct mgp_list *args, const struct mgp_graph *graph, struct mgp_result *result,
+ struct mgp_memory *memory) {
+ const size_t one_gb = 1 << 30;
+ if (gVal) {
+ mgp_global_free(gVal);
+ gVal = NULL;
+ }
+ if (!gVal) {
+ gVal = mgp_global_alloc(one_gb);
+ if (!gVal) return set_out_of_memory_error(result);
+ }
+ struct mgp_result_record *record = mgp_result_new_record(result);
+ if (record == NULL) return set_error(result);
+ struct mgp_value *error_value = mgp_value_make_string("ERROR", memory);
+ if (error_value == NULL) return set_error(result);
+ int result_inserted = mgp_result_record_insert(record, "error_result", error_value);
+ mgp_value_destroy(error_value);
+ if (!result_inserted) return set_error(result);
+}
+
+static void success(const struct mgp_list *args, const struct mgp_graph *graph, struct mgp_result *result,
+ struct mgp_memory *memory) {
+ const size_t bytes = 1024;
+ if (!gVal) {
+ gVal = mgp_global_alloc(bytes);
+ if (!gVal) set_out_of_memory_error(result);
+ }
+
+ struct mgp_result_record *record = mgp_result_new_record(result);
+ if (record == NULL) return set_error(result);
+ struct mgp_value *success_value = mgp_value_make_string("sucess", memory);
+ if (success_value == NULL) return set_error(result);
+ int result_inserted = mgp_result_record_insert(record, "success_result", success_value);
+ mgp_value_destroy(success_value);
+ if (!result_inserted) return set_error(result);
+}
+
+int mgp_init_module(struct mgp_module *module, struct mgp_memory *memory) {
+ struct mgp_proc *error_proc = mgp_module_add_read_procedure(module, "error", error);
+ if (!error_proc) return 1;
+
+ if (!mgp_proc_add_result(error_proc, "error_result", mgp_type_string())) return 1;
+
+ struct mgp_proc *succ_proc = mgp_module_add_read_procedure(module, "success", success);
+ if (!succ_proc) return 1;
+
+ if (!mgp_proc_add_result(succ_proc, "success_result", mgp_type_string())) return 1;
+
+ return 0;
+}
+
+int mgp_shutdown_module() {
+ if (gVal) mgp_global_free(gVal);
+ return 0;
+}
diff --git a/tests/e2e/memory/workloads.yaml b/tests/e2e/memory/workloads.yaml
index adec01260..bf7ba373e 100644
--- a/tests/e2e/memory/workloads.yaml
+++ b/tests/e2e/memory/workloads.yaml
@@ -13,4 +13,14 @@ workloads:
args: ["--bolt-port", *bolt_port, "--timeout", "180"]
<<: *template_cluster
+ - name: "Memory limit for modules upon loading"
+ binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc"
+ args: ["--bolt-port", *bolt_port, "--timeout", "180"]
+ proc: "tests/e2e/memory/procedures/"
+ <<: *template_cluster
+ - name: "Memory limit for modules inside a procedure"
+ binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc_proc"
+ args: ["--bolt-port", *bolt_port, "--timeout", "180"]
+ proc: "tests/e2e/memory/procedures/"
+ <<: *template_cluster
diff --git a/tests/e2e/runner.py b/tests/e2e/runner.py
index 36a7a34de..74b69955e 100755
--- a/tests/e2e/runner.py
+++ b/tests/e2e/runner.py
@@ -51,6 +51,10 @@ def run(args):
mg_instances[name] = mg_instance
log_file_path = os.path.join(BUILD_DIR, 'logs', config['log_file'])
binary_args = config['args'] + ["--log-file", log_file_path]
+ if 'proc' in workload:
+ procdir = "--query-modules-directory=" + os.path.join(BUILD_DIR, workload['proc'])
+ binary_args.append(procdir)
+
mg_instance.start(args=binary_args)
for query in config['setup_queries']:
mg_instance.query(query)