Merge pull request #8 from arangodb/bugfix/iresearch-memory-management

IResearch memory management implementation

Author: kvahed

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,69 @@ struct ResourceManagementOptions {
   IResourceManager* cached_columns{&IResourceManager::kNoop};
 };
 
+//  Memory manager for IResearch.
+//  This is a singleton class
+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 == _memoryLimit) {
+      // since we have no limit, we can simply use fetch-add for the increment
+      _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 = _current.load(std::memory_order_relaxed);
+      std::uint64_t next;
+      do {
+        next = cur + value;
+        if (IRS_UNLIKELY(next > _memoryLimit.load(std::memory_order_relaxed))) {
+          throw std::bad_alloc();
+        }
+      } while (!_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);
+  }
+
+  //  NOTE: IResearchFeature owns and manages this memory limit.
+  //  That is why this method should only be used by IResearchFeature.
+  virtual void SetMemoryLimit(uint64_t memoryLimit) {
+    _memoryLimit.store(memoryLimit);
+  }
+
+private:
+  //  This limit should be set exclusively by IResearchFeature.
+  //  During IResearchFeature::validateOptions() this limit is set to a
+  //  percentage of either the total available physical memory or the value
+  //  of ARANGODB_OVERRIDE_DETECTED_TOTAL_MEMORY envvar if specified.
+  std::atomic<std::uint64_t> _memoryLimit = { 0 };
+  std::atomic<std::uint64_t> _current = { 0 };
+
+  //  Singleton
+  static inline std::shared_ptr<IResearchMemoryManager> _instance;
+
+public:
+  static std::shared_ptr<IResearchMemoryManager> GetInstance() {
+    if (!_instance.get())
+      _instance.reset(new IResearchMemoryManager());
+
+    return _instance;
+  }
+};
+
 template<typename T>
 struct ManagedTypedAllocator
   : ManagedAllocator<std::allocator<T>, IResourceManager> {
@@ -77,7 +146,7 @@ struct ManagedTypedAllocator
 #if !defined(_MSC_VER) && defined(IRESEARCH_DEBUG)
         IResourceManager::kForbidden
 #else
-        IResourceManager::kNoop
+        *IResearchMemoryManager::GetInstance()
 #endif
       ) {
   }

tests/CMakeLists.txt

@@ -121,6 +121,7 @@ set(IReSearch_tests_sources
   ./utils/fst_utils_test.cpp
   ./utils/fst_string_weight_test.cpp
   ./utils/ngram_match_utils_tests.cpp
+  ./memory/IResearchMemoryManager_tests.cpp
   ./tests_param.cpp
   ./tests_main.cpp
   )

tests/memory/IResearchMemoryManager_tests.cpp

@@ -0,0 +1,92 @@
+////////////////////////////////////////////////////////////////////////////////
+/// DISCLAIMER
+///
+/// Copyright 2014-2025 ArangoDB GmbH, Cologne, Germany
+/// Copyright 2004-2014 triAGENS GmbH, Cologne, Germany
+///
+/// Licensed under the Apache License, Version 2.0 (the "License");
+/// you may not use this file except in compliance with the License.
+/// You may obtain a copy of the License at
+///
+///     http://www.apache.org/licenses/LICENSE-2.0
+///
+/// Unless required by applicable law or agreed to in writing, software
+/// distributed under the License is distributed on an "AS IS" BASIS,
+/// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+/// See the License for the specific language governing permissions and
+/// limitations under the License.
+///
+/// Copyright holder is ArangoDB GmbH, Cologne, Germany
+///
+/// @author Koushal Kawade
+////////////////////////////////////////////////////////////////////////////////
+
+#include <iostream>
+#include <memory>
+#include "tests_shared.hpp"
+#include "utils/managed_allocator.hpp"
+#include "resource_manager.hpp"
+
+using namespace irs;
+
+struct LargeData {
+    char data[1024];
+};
+
+TEST(IResearchMemoryLimitTest, no_limit_set_allows_infinite_allocations) {
+
+    //  The expectation is to allow potentially infinite allocations
+    //  even if we're testing with just 1MB.
+
+    //  allocate vector
+    ManagedVector<LargeData> vec;
+
+    for (size_t i = 0; i < 1024; i++) {
+        ASSERT_NO_THROW(vec.push_back(LargeData()));
+    }
+}
+
+TEST(IResearchMemoryLimitTest, zero_limit_allow_infinite_allocations) {
+
+    auto memoryMgr = IResearchMemoryManager::GetInstance();
+    memoryMgr->SetMemoryLimit(0);
+
+    //  allocate vector
+    ManagedVector<LargeData> vec;
+
+    for (size_t i = 0; i < 1024; i++) {
+        ASSERT_NO_THROW(vec.push_back(LargeData()));
+    }
+}
+
+TEST(IResearchMemoryLimitTest, managed_allocator_smoke_test) {
+
+    auto memoryMgr = IResearchMemoryManager::GetInstance();
+    memoryMgr->SetMemoryLimit(3);
+
+    using type = char;
+    std::vector<type, ManagedTypedAllocator<type>> arr(*memoryMgr);
+
+    arr.push_back('a');
+    arr.push_back('b');
+    ASSERT_THROW(arr.push_back('c'), std::bad_alloc);
+}
+
+TEST(IResearchMemoryLimitTest, memory_manager_smoke_test) {
+
+    //  set limit
+    auto memoryMgr = IResearchMemoryManager::GetInstance();
+    memoryMgr->SetMemoryLimit(47);
+
+    //  allocate vector
+    ManagedVector<uint64_t> vec;
+    vec.push_back(10);  //  Allocate 8 bytes
+    vec.push_back(11);  //  Allocate 16, Free previous 8, Copy both elements in the new array.
+
+    ASSERT_THROW(vec.push_back(12), std::bad_alloc);    //  Allocate 32 while holding previous 16, total 48 (bad_alloc)
+
+    //  Increase memory limit to accommodate the 3rd element.
+    memoryMgr->SetMemoryLimit(48);
+
+    ASSERT_NO_THROW(vec.push_back(12));
+}

tests/utils/utf8_path_tests.cpp

@@ -26,6 +26,7 @@
 #include <filesystem>
 #include <fstream>
 #include <thread>
+#include <algorithm>
 
 #include "tests_shared.hpp"
 #include "utils/file_utils.hpp"