add limits

include/mg_procedure.h

@@ -111,6 +111,13 @@ 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;
+
+enum mgp_error mgp_track_thread_allocations(struct mgp_graph *graph, const char *thread_id);
+
+enum mgp_error mgp_untrack_thread_allocations(struct mgp_graph *graph, const char *);
 ///@}
 
 /// @name Operations on mgp_value
@@ -851,9 +858,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.

src/memory/memory_control.cpp

@@ -267,7 +267,12 @@ void UpdateThreadToTransactionId(const std::thread::id &thread_id, uint64_t tran
   std::ostringstream oss;
   oss << thread_id;
   thread_id_to_transaction_id[oss.str()] = transaction_id;
-  std::cout << "set:" << oss.str() << std::endl;
+  std::cout << "set tracking for thread:" << oss.str() << ", on transaction: " << transaction_id << std::endl;
+}
+
+void UpdateThreadToTransactionId(const char *thread_id, uint64_t transaction_id) {
+  thread_id_to_transaction_id[std::string(thread_id)] = transaction_id;
+  std::cout << "set tracking:" << std::string(thread_id) << ", on transaction: " << transaction_id << std::endl;
 }
 
 void ResetThreadToTransactionId(const std::thread::id &thread_id) {
@@ -276,6 +281,8 @@ void ResetThreadToTransactionId(const std::thread::id &thread_id) {
   thread_id_to_transaction_id.erase(oss.str());
 }
 
+void ResetThreadToTransactionId(const char *thread_id) { thread_id_to_transaction_id.erase(std::string(thread_id)); }
+
 void PurgeUnusedMemory() {
 #if USE_JEMALLOC
   mallctl("arena." STRINGIFY(MALLCTL_ARENAS_ALL) ".purge", nullptr, nullptr, nullptr, 0);

src/memory/memory_control.hpp

@@ -27,6 +27,8 @@ 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 *);
 
 // TODO(AF) : Do we need arena to transaction id and transaction_id to tracker?
 inline std::unordered_map<std::string, uint64_t> thread_id_to_transaction_id;

src/query/db_accessor.cpp

@@ -21,6 +21,14 @@ 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::UntrackThreadAllocations(const char *thread_id) {
+  return db_accessor_.UntrackThreadAllocations(thread_id);
+}
+
 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,18 @@ 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 UntrackThreadAllocations(const char *thread_id) {
+    memgraph::memory::ResetThreadToTransactionId(thread_id);
+    auto arena = memgraph::memory::GetArenaForThread();
+    memgraph::memory::RemoveTrackingOnArena(arena);
+  }
+
   std::optional<uint64_t> GetTransactionId() { return accessor_->GetTransactionId(); }
 
   VerticesIterable Vertices(storage::View view) { return VerticesIterable(accessor_->Vertices(view)); }
@@ -645,6 +658,10 @@ class SubgraphDbAccessor final {
 
   static SubgraphDbAccessor *MakeSubgraphDbAccessor(DbAccessor *db_accessor, Graph *graph);
 
+  void TrackThreadAllocations(const char *thread_id);
+
+  void UntrackThreadAllocations(const char *thread_id);
+
   storage::PropertyId NameToProperty(std::string_view name);
 
   storage::LabelId NameToLabel(std::string_view name);

src/query/procedure/mg_procedure_impl.cpp

@@ -3540,3 +3540,16 @@ 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_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);
+  });
+}

tests/e2e/memory/procedures/CMakeLists.txt

@@ -3,3 +3,7 @@ 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) #todo link with utils not to use onscopeexit

tests/e2e/memory/procedures/query_memory_limit_proc_multi_thread.cpp

@@ -0,0 +1,110 @@
+// 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 <exception>
+#include <functional>
+#include <mgp.hpp>
+#include <mutex>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <utility>
+#include <vector>
+#include "mg_procedure.h"
+
+template <typename Callable>
+class [[nodiscard]] OnScopeExit {
+ public:
+  explicit OnScopeExit(Callable &&function) : function_{std::forward<Callable>(function)}, doCall_{true} {}
+  OnScopeExit(OnScopeExit const &) = delete;
+  OnScopeExit(OnScopeExit &&) = delete;
+  OnScopeExit &operator=(OnScopeExit const &) = delete;
+  OnScopeExit &operator=(OnScopeExit &&) = delete;
+  ~OnScopeExit() {
+    if (doCall_) function_();
+  }
+
+  void Disable() { doCall_ = false; }
+
+ private:
+  std::function<void()> function_;
+  bool doCall_;
+};
+template <typename Callable>
+OnScopeExit(Callable &&) -> OnScopeExit<Callable>;
+
+// TODO(af) fix this hack
+std::atomic<int> num_allocations{0};
+
+void AllocFunc(mgp_graph *graph) {
+  std::ostringstream oss;
+  oss << std::this_thread::get_id();
+  std::string thread_id = oss.str();
+  [[maybe_unused]] const enum mgp_error tracking_error = mgp_track_thread_allocations(graph, thread_id.c_str());
+  const size_t two_sixty_eight_mb = 1 << 28;
+  void *ptr = nullptr;
+
+  OnScopeExit<std::function<void(void)>> cleanup{[&ptr]() {
+    if (nullptr == ptr) {
+      return;
+    }
+    mgp_global_free(ptr);
+  }};
+
+  const enum mgp_error alloc_err = mgp_global_alloc(two_sixty_eight_mb, (void **)(&ptr));
+  if (alloc_err != mgp_error::MGP_ERROR_UNABLE_TO_ALLOCATE) {
+    num_allocations.fetch_add(1, std::memory_order_relaxed);
+  }
+
+  [[maybe_unused]] const enum mgp_error untracking_error = mgp_untrack_thread_allocations(graph, thread_id.c_str());
+}
+
+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);
+
+  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();
+    }
+
+    auto new_record = record_factory.NewRecord();
+
+    new_record.Insert("allocated_all", num_allocations.load(std::memory_order_relaxed) == 2);
+    // ASSERT only one exception should occur
+  } 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);
+
+  } 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,65 @@
+// 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 <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();
+      std::cout << row[0].ValueBool() << std::endl;
+    }
+
+  } catch (const mg::ClientException &e) {
+    std::cout << e.what() << std::endl;
+    error = true;
+  }
+
+  MG_ASSERT(error, "Error should have happend");
+
+  return 0;
+}

tests/e2e/memory/workloads.yaml

@@ -23,50 +23,69 @@ 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"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__control"
-    args: ["--bolt-port", *bolt_port, "--timeout", "180"]
-    <<: *in_memory_cluster
+  # - name: "Memory control"
+  #   binary: "tests/e2e/memory/memgraph__e2e__memory__control"
+  #   args: ["--bolt-port", *bolt_port, "--timeout", "180"]
+  #   <<: *in_memory_cluster
 
-  - name: "Memory control multi database"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__control"
-    args: ["--bolt-port", *bolt_port, "--timeout", "180", "--multi-db", "true"]
-    <<: *in_memory_cluster
+  # - name: "Memory control multi database"
+  #   binary: "tests/e2e/memory/memgraph__e2e__memory__control"
+  #   args: ["--bolt-port", *bolt_port, "--timeout", "180", "--multi-db", "true"]
+  #   <<: *in_memory_cluster
 
-  - name: "Memory limit for modules upon loading"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc"
+  # - 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/"
+  #   <<: *in_memory_cluster
+
+  # - name: "Memory limit for modules upon loading multi database"
+  #   binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc"
+  #   args: ["--bolt-port", *bolt_port, "--timeout", "180", "--multi-db", "true"]
+  #   proc: "tests/e2e/memory/procedures/"
+  #   <<: *in_memory_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/"
+  #   <<: *in_memory_cluster
+
+  # - name: "Memory limit for modules inside a procedure multi database"
+  #   binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc_proc"
+  #   args: ["--bolt-port", *bolt_port, "--timeout", "180", "--multi-db", "true"]
+  #   proc: "tests/e2e/memory/procedures/"
+  #   <<: *in_memory_cluster
+
+  # - name: "Memory limit for modules upon loading for on-disk storage"
+  #   binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc"
+  #   args: ["--bolt-port", *bolt_port, "--timeout", "180"]
+  #   proc: "tests/e2e/memory/procedures/"
+  #   <<: *disk_cluster
+
+  # - name: "Memory limit for modules inside a procedure for on-disk storage"
+  #   binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc_proc"
+  #   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_cluster
-
-  - name: "Memory limit for modules upon loading multi database"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc"
-    args: ["--bolt-port", *bolt_port, "--timeout", "180", "--multi-db", "true"]
-    proc: "tests/e2e/memory/procedures/"
-    <<: *in_memory_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/"
-    <<: *in_memory_cluster
-
-  - name: "Memory limit for modules inside a procedure multi database"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc_proc"
-    args: ["--bolt-port", *bolt_port, "--timeout", "180", "--multi-db", "true"]
-    proc: "tests/e2e/memory/procedures/"
-    <<: *in_memory_cluster
-
-  - name: "Memory limit for modules upon loading for on-disk storage"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc"
-    args: ["--bolt-port", *bolt_port, "--timeout", "180"]
-    proc: "tests/e2e/memory/procedures/"
-    <<: *disk_cluster
-
-  - name: "Memory limit for modules inside a procedure for on-disk storage"
-    binary: "tests/e2e/memory/memgraph__e2e__memory__limit_global_alloc_proc"
-    args: ["--bolt-port", *bolt_port, "--timeout", "180"]
-    proc: "tests/e2e/memory/procedures/"
-    <<: *disk_cluster
+    <<: *in_memory_query_limit_cluster