Implemented memory manager to use with allocators

Author: k0ushal

core/resource_manager.hpp

@@ -22,11 +22,17 @@
 
 #pragma once
 
+#include <atomic>
 #include <vector>
 
 #include "shared.hpp"
 #include "utils/managed_allocator.hpp"
 
+namespace arangodb::iresearch
+{
+  class IResearchFeature;
+};
+
 namespace irs {
 
 struct IResourceManager {
@@ -68,6 +74,65 @@ struct ResourceManagementOptions {
   IResourceManager* cached_columns{&IResourceManager::kNoop};
 };
 
+struct IResearchMemoryManager : public IResourceManager {
+protected:
+  IResearchMemoryManager() = default;
+
+public:
+  virtual ~IResearchMemoryManager() = default;
+
+  virtual void Increase([[maybe_unused]] uint64_t value) override {
+
+    IRS_ASSERT(this != &kForbidden);
+    IRS_ASSERT(value >= 0);
+
+    if (0 == IResearchMemoryManager::_memory_limit) {
+      // since we have no limit, we can simply use fetch-add for the increment
+      IResearchMemoryManager::_current.fetch_add(value, std::memory_order_relaxed);
+    } else {
+      // we only want to perform the update if we don't exceed the limit!
+      std::uint64_t cur = IResearchMemoryManager::_current.load(std::memory_order_relaxed);
+      std::uint64_t next;
+      do {
+        next = cur + value;
+        if (IRS_UNLIKELY(next > IResearchMemoryManager::_memory_limit.load(std::memory_order_relaxed))) {
+          throw std::bad_alloc();
+        }
+      } while (!IResearchMemoryManager::_current.compare_exchange_weak(
+        cur, next, std::memory_order_relaxed));
+    }
+  }
+
+  virtual void Decrease([[maybe_unused]] uint64_t value) noexcept override {
+    IRS_ASSERT(this != &kForbidden);
+    IRS_ASSERT(value >= 0);
+    _current.fetch_sub(value, std::memory_order_relaxed);
+  }
+
+protected:
+  //  This limit can be set only by IResearchFeature.
+  //  During IResearchFeature::prepare() it is set to a pre-defined
+  //  percentage of the total available physical memory or to the value
+  //  of ARANGODB_OVERRIDE_DETECTED_TOTAL_MEMORY option if specified.
+  static inline std::atomic<std::uint64_t> _memory_limit = { 0 };
+
+  //  Singleton
+  static inline std::shared_ptr<IResourceManager> _instance;
+
+private:
+  static inline std::atomic<std::uint64_t> _current = { 0 };
+
+public:
+  static std::shared_ptr<IResourceManager> GetInstance() {
+    if (!_instance.get())
+      _instance.reset(new IResearchMemoryManager());
+
+    return _instance;
+  }
+
+  friend class arangodb::iresearch::IResearchFeature;
+};
+
 template<typename T>
 struct ManagedTypedAllocator
   : ManagedAllocator<std::allocator<T>, IResourceManager> {
@@ -77,7 +142,7 @@ struct ManagedTypedAllocator
 #if !defined(_MSC_VER) && defined(IRESEARCH_DEBUG)
         IResourceManager::kForbidden
 #else
-        IResourceManager::kNoop
+        *IResearchMemoryManager::GetInstance()
 #endif
       ) {
   }