From 5567e6ec9688976e0b4152444145fd735c9200cf Mon Sep 17 00:00:00 2001
From: Teon Banek <teon.banek@memgraph.io>
Date: Tue, 5 Nov 2019 10:54:06 +0100
Subject: [PATCH] Add utils::LimitedMemoryResource

Reviewers: llugovic, mferencevic, ipaljak

Reviewed By: mferencevic

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D2534
---
 src/utils/memory.hpp | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/src/utils/memory.hpp b/src/utils/memory.hpp
index 75ff7a69f..82af8b420 100644
--- a/src/utils/memory.hpp
+++ b/src/utils/memory.hpp
@@ -569,4 +569,37 @@ class SynchronizedPoolResource final : public MemoryResource {
   }
 };
 
+class LimitedMemoryResource final : public utils::MemoryResource {
+ public:
+  explicit LimitedMemoryResource(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_;
+  }
+
+ private:
+  utils::MemoryResource *memory_;
+  size_t max_allocated_bytes_;
+  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 {
+    CHECK(available_bytes_ + bytes > available_bytes_);
+    available_bytes_ += bytes;
+    return memory_->Deallocate(p, bytes, alignment);
+  }
+
+  bool DoIsEqual(const MemoryResource &other) const noexcept override {
+    return this == &other;
+  }
+};
+
 }  // namespace utils