Add memory tracker definition (#93)

* Allow size 0 in MemoryTracker
* Block only exception throwing
* Subtract unsuccessfully allocated size in memory tracker
* Add oom exception enablers and blockers

src/memgraph.cpp

@@ -35,6 +35,7 @@
 #include "utils/file.hpp"
 #include "utils/flag_validation.hpp"
 #include "utils/logging.hpp"
+#include "utils/memory_tracker.hpp"
 #include "utils/signals.hpp"
 #include "utils/string.hpp"
 #include "utils/sysinfo/memory.hpp"
@@ -1024,5 +1025,7 @@ int main(int argc, char **argv) {
   // Shutdown Python
   Py_Finalize();
   PyMem_RawFree(program_name);
+
+  utils::total_memory_tracker.LogPeakMemoryUsage();
   return 0;
 }

src/utils/CMakeLists.txt

@@ -4,6 +4,7 @@ set(utils_src_files
     file.cpp
     file_locker.cpp
     memory.cpp
+    memory_tracker.cpp
     signals.cpp
     thread.cpp
     thread_pool.cpp
@@ -13,4 +14,4 @@ add_library(mg-utils STATIC ${utils_src_files})
 target_link_libraries(mg-utils stdc++fs Threads::Threads spdlog fmt gflags uuid)
 
 add_library(mg-new-delete STATIC new_delete.cpp)
-target_link_libraries(mg-new-delete jemalloc)
+target_link_libraries(mg-new-delete jemalloc fmt)

src/utils/memory_tracker.cpp

@@ -0,0 +1,103 @@
+#include "utils/memory_tracker.hpp"
+
+#include <atomic>
+#include <exception>
+#include <stdexcept>
+
+#include "utils/likely.hpp"
+#include "utils/logging.hpp"
+#include "utils/on_scope_exit.hpp"
+
+namespace utils {
+
+namespace {
+
+// Prevent memory tracker for throwing during the stack unwinding
+bool MemoryTrackerCanThrow() {
+  return !std::uncaught_exceptions() && MemoryTracker::OutOfMemoryExceptionEnabler::CanThrow() &&
+         !MemoryTracker::OutOfMemoryExceptionBlocker::IsBlocked();
+}
+
+std::string GetReadableSize(double size) {
+  // TODO (antonio2368): Add support for base 1000 (KB, GB, TB...)
+  constexpr std::array units = {"B", "KiB", "MiB", "GiB", "TiB"};
+  constexpr double delimiter = 1024;
+
+  size_t i = 0;
+  for (; i + 1 < units.size() && size >= delimiter; ++i) {
+    size /= delimiter;
+  }
+
+  // bytes don't need decimals
+  if (i == 0) {
+    return fmt::format("{:.0f}{}", size, units[i]);
+  }
+
+  return fmt::format("{:.2f}{}", size, units[i]);
+}
+
+}  // namespace
+
+thread_local uint64_t MemoryTracker::OutOfMemoryExceptionEnabler::counter_ = 0;
+MemoryTracker::OutOfMemoryExceptionEnabler::OutOfMemoryExceptionEnabler() { ++counter_; }
+MemoryTracker::OutOfMemoryExceptionEnabler::~OutOfMemoryExceptionEnabler() { --counter_; }
+bool MemoryTracker::OutOfMemoryExceptionEnabler::CanThrow() { return counter_ > 0; }
+
+thread_local uint64_t MemoryTracker::OutOfMemoryExceptionBlocker::counter_ = 0;
+MemoryTracker::OutOfMemoryExceptionBlocker::OutOfMemoryExceptionBlocker() { ++counter_; }
+MemoryTracker::OutOfMemoryExceptionBlocker::~OutOfMemoryExceptionBlocker() { --counter_; }
+bool MemoryTracker::OutOfMemoryExceptionBlocker::IsBlocked() { return counter_ > 0; }
+
+MemoryTracker total_memory_tracker;
+
+// TODO (antonio2368): Define how should the peak memory be logged.
+// Logging every time the peak changes is too much so some kind of distribution
+// should be used.
+void MemoryTracker::LogPeakMemoryUsage() const { spdlog::info("Peak memory usage: {}", GetReadableSize(peak_)); }
+
+// TODO (antonio2368): Define how should the memory be logged.
+// Logging on each allocation is too much so some kind of distribution
+// should be used.
+void MemoryTracker::LogMemoryUsage(const int64_t current) {
+  spdlog::info("Current memory usage: {}", GetReadableSize(current));
+}
+
+void MemoryTracker::UpdatePeak(const int64_t will_be) {
+  auto peak_old = peak_.load(std::memory_order_relaxed);
+  if (will_be > peak_old) {
+    peak_.store(will_be, std::memory_order_relaxed);
+  }
+}
+
+void MemoryTracker::SetHardLimit(const int64_t limit) { hard_limit_.store(limit, std::memory_order_relaxed); }
+
+void MemoryTracker::SetOrRaiseHardLimit(const int64_t limit) {
+  int64_t old_limit = hard_limit_.load(std::memory_order_relaxed);
+  while (old_limit < limit && !hard_limit_.compare_exchange_weak(old_limit, limit))
+    ;
+}
+
+void MemoryTracker::Alloc(const int64_t size) {
+  MG_ASSERT(size >= 0, "Negative size passed to the MemoryTracker.");
+
+  const int64_t will_be = size + amount_.fetch_add(size, std::memory_order_relaxed);
+
+  const auto current_hard_limit = hard_limit_.load(std::memory_order_relaxed);
+
+  if (UNLIKELY(current_hard_limit && will_be > current_hard_limit && MemoryTrackerCanThrow())) {
+    MemoryTracker::OutOfMemoryExceptionBlocker exception_blocker;
+
+    amount_.fetch_sub(size, std::memory_order_relaxed);
+
+    throw OutOfMemoryException(
+        fmt::format("Memory limit exceeded! Atempting to allocate a chunk of {} which would put the current "
+                    "use to {}, while the maximum allowed size for allocation is set to {}.",
+                    GetReadableSize(size), GetReadableSize(will_be), GetReadableSize(current_hard_limit)));
+  }
+
+  UpdatePeak(will_be);
+}
+
+void MemoryTracker::Free(const int64_t size) { amount_.fetch_sub(size, std::memory_order_relaxed); }
+
+}  // namespace utils

src/utils/memory_tracker.hpp

@@ -0,0 +1,85 @@
+#pragma once
+
+#include <atomic>
+
+#include "utils/exceptions.hpp"
+
+namespace utils {
+
+class OutOfMemoryException : public utils::BasicException {
+ public:
+  explicit OutOfMemoryException(const std::string &msg) : utils::BasicException(msg) {}
+};
+
+class MemoryTracker final {
+ private:
+  std::atomic<int64_t> amount_{0};
+  std::atomic<int64_t> peak_{0};
+  std::atomic<int64_t> hard_limit_{0};
+
+  void UpdatePeak(int64_t will_be);
+
+  static void LogMemoryUsage(int64_t current);
+
+ public:
+  void LogPeakMemoryUsage() const;
+
+  MemoryTracker() = default;
+  ~MemoryTracker() = default;
+
+  MemoryTracker(const MemoryTracker &) = delete;
+  MemoryTracker &operator=(const MemoryTracker &) = delete;
+  MemoryTracker(MemoryTracker &&) = delete;
+  MemoryTracker &operator=(MemoryTracker &&) = delete;
+
+  void Alloc(int64_t size);
+  void Free(int64_t size);
+
+  auto Amount() const { return amount_.load(std::memory_order_relaxed); }
+
+  auto Peak() const { return peak_.load(std::memory_order_relaxed); }
+
+  void SetHardLimit(int64_t limit);
+  void SetOrRaiseHardLimit(int64_t limit);
+
+  // 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 {
+   public:
+    OutOfMemoryExceptionEnabler(const OutOfMemoryExceptionEnabler &) = delete;
+    OutOfMemoryExceptionEnabler &operator=(const OutOfMemoryExceptionEnabler &) = delete;
+    OutOfMemoryExceptionEnabler(OutOfMemoryExceptionEnabler &&) = delete;
+    OutOfMemoryExceptionEnabler &operator=(OutOfMemoryExceptionEnabler &&) = delete;
+
+    OutOfMemoryExceptionEnabler();
+    ~OutOfMemoryExceptionEnabler();
+
+    static bool CanThrow();
+
+   private:
+    static thread_local uint64_t counter_;
+  };
+
+  // By creating an object of this class, we negate the effect of every OutOfMemoryExceptionEnabler
+  // object. We need this object so we can guard only the smaller parts of code from exceptions while
+  // allowing the exception in the other parts if the OutOfMemoryExceptionEnabler is defined.
+  class OutOfMemoryExceptionBlocker final {
+   public:
+    OutOfMemoryExceptionBlocker(const OutOfMemoryExceptionBlocker &) = delete;
+    OutOfMemoryExceptionBlocker &operator=(const OutOfMemoryExceptionBlocker &) = delete;
+    OutOfMemoryExceptionBlocker(OutOfMemoryExceptionBlocker &&) = delete;
+    OutOfMemoryExceptionBlocker &operator=(OutOfMemoryExceptionBlocker &&) = delete;
+
+    OutOfMemoryExceptionBlocker();
+    ~OutOfMemoryExceptionBlocker();
+
+    static bool IsBlocked();
+
+   private:
+    static thread_local uint64_t counter_;
+  };
+};
+
+// Global memory tracker which tracks every allocation in the application.
+extern MemoryTracker total_memory_tracker;
+}  // namespace utils

src/utils/new_delete.cpp

@@ -1,4 +1,3 @@
-#include <iostream>
 #include <new>
 
 #if USE_JEMALLOC
@@ -8,6 +7,7 @@
 #endif
 
 #include "utils/likely.hpp"
+#include "utils/memory_tracker.hpp"
 
 namespace {
 void *newImpl(std::size_t size) {
@@ -39,20 +39,87 @@ void deleteSized(void *ptr, const std::size_t /*unused*/) noexcept { free(ptr);
 
 #endif
 
+void TrackMemory(const size_t size) {
+  size_t actual_size = size;
+
+#if USE_JEMALLOC
+  if (LIKELY(size != 0)) {
+    actual_size = nallocx(size, 0);
+  }
+#endif
+  utils::total_memory_tracker.Alloc(actual_size);
+}
+
+bool TrackMemoryNoExcept(const size_t size) {
+  try {
+    TrackMemory(size);
+  } catch (...) {
+    return false;
+  }
+
+  return true;
+}
+
+void UntrackMemory([[maybe_unused]] void *ptr, [[maybe_unused]] size_t size = 0) noexcept {
+  try {
+#if USE_JEMALLOC
+    if (LIKELY(ptr != nullptr)) {
+      utils::total_memory_tracker.Free(sallocx(ptr, 0));
+    }
+#else
+    if (size) {
+      utils::total_memory_tracker.Free(size);
+    } else {
+      // Innaccurate because malloc_usable_size() result is greater or equal to allocated size.
+      utils::total_memory_tracker.Free(malloc_usable_size(ptr));
+    }
+#endif
+  } catch (...) {
+  }
+}
+
 }  // namespace
 
-void *operator new(std::size_t size) { return newImpl(size); }
+void *operator new(std::size_t size) {
+  TrackMemory(size);
+  return newImpl(size);
+}
 
-void *operator new[](std::size_t size) { return newImpl(size); }
+void *operator new[](std::size_t size) {
+  TrackMemory(size);
+  return newImpl(size);
+}
 
-void *operator new(std::size_t size, const std::nothrow_t & /*unused*/) noexcept { return newNoExcept(size); }
+void *operator new(std::size_t size, const std::nothrow_t & /*unused*/) noexcept {
+  if (LIKELY(TrackMemoryNoExcept(size))) {
+    return newNoExcept(size);
+  }
+  return nullptr;
+}
 
-void *operator new[](std::size_t size, const std::nothrow_t & /*unused*/) noexcept { return newNoExcept(size); }
+void *operator new[](std::size_t size, const std::nothrow_t & /*unused*/) noexcept {
+  if (LIKELY(TrackMemoryNoExcept(size))) {
+    return newNoExcept(size);
+  }
+  return nullptr;
+}
 
-void operator delete(void *ptr) noexcept { deleteImpl(ptr); }
+void operator delete(void *ptr) noexcept {
+  UntrackMemory(ptr);
+  deleteImpl(ptr);
+}
 
-void operator delete[](void *ptr) noexcept { deleteImpl(ptr); }
+void operator delete[](void *ptr) noexcept {
+  UntrackMemory(ptr);
+  deleteImpl(ptr);
+}
 
-void operator delete(void *ptr, std::size_t size) noexcept { deleteSized(ptr, size); }
+void operator delete(void *ptr, std::size_t size) noexcept {
+  UntrackMemory(ptr, size);
+  deleteSized(ptr, size);
+}
 
-void operator delete[](void *ptr, std::size_t size) noexcept { deleteSized(ptr, size); }
+void operator delete[](void *ptr, std::size_t size) noexcept {
+  UntrackMemory(ptr, size);
+  deleteSized(ptr, size);
+}

tests/unit/CMakeLists.txt

@@ -193,6 +193,9 @@ target_link_libraries(${test_prefix}utils_math mg-utils)
 add_unit_test(utils_memory.cpp)
 target_link_libraries(${test_prefix}utils_memory mg-utils)
 
+add_unit_test(utils_memory_tracker.cpp)
+target_link_libraries(${test_prefix}utils_memory_tracker mg-utils)
+
 add_unit_test(utils_on_scope_exit.cpp)
 target_link_libraries(${test_prefix}utils_on_scope_exit mg-utils)
 

tests/unit/utils_memory_tracker.cpp

@@ -0,0 +1,54 @@
+#include <thread>
+
+#include <gtest/gtest.h>
+
+#include <utils/memory_tracker.hpp>
+
+TEST(MemoryTrackerTest, ExceptionEnabler) {
+  utils::MemoryTracker memory_tracker;
+
+  constexpr size_t hard_limit = 10;
+  memory_tracker.SetHardLimit(hard_limit);
+
+  std::atomic<bool> can_continue{false};
+  std::atomic<bool> enabler_created{false};
+  std::thread t1{[&] {
+    // wait until the second thread creates exception enabler
+    while (!enabler_created)
+      ;
+    ASSERT_NO_THROW(memory_tracker.Alloc(hard_limit + 1));
+    ASSERT_EQ(memory_tracker.Amount(), hard_limit + 1);
+
+    // tell the second thread it can finish its test
+    can_continue = true;
+  }};
+
+  std::thread t2{[&] {
+    utils::MemoryTracker::OutOfMemoryExceptionEnabler exception_enabler;
+    enabler_created = true;
+    ASSERT_THROW(memory_tracker.Alloc(hard_limit + 1), utils::OutOfMemoryException);
+
+    // hold the enabler until the first thread finishes
+    while (!can_continue)
+      ;
+  }};
+
+  t1.join();
+  t2.join();
+}
+
+TEST(MemoryTrackerTest, ExceptionBlocker) {
+  utils::MemoryTracker memory_tracker;
+
+  constexpr size_t hard_limit = 10;
+  memory_tracker.SetHardLimit(hard_limit);
+
+  utils::MemoryTracker::OutOfMemoryExceptionEnabler exception_enabler;
+  {
+    utils::MemoryTracker::OutOfMemoryExceptionBlocker exception_blocker;
+
+    ASSERT_NO_THROW(memory_tracker.Alloc(hard_limit + 1));
+    ASSERT_EQ(memory_tracker.Amount(), hard_limit + 1);
+  }
+  ASSERT_THROW(memory_tracker.Alloc(hard_limit + 1), utils::OutOfMemoryException);
+}