feat(): aggregating CountingAllocator

tests/counting_allocator.cpp

@@ -17,25 +17,28 @@
 #include <gtest/gtest.h>
 #include <atomic>
 #include <cstddef>
+#include <cstdint>
 #include <memory>
 #include <thread>
 
+#include <boost/smart_ptr/allocate_shared_array.hpp>
+
 #include "countingallocator.h"
-#include "rowgroup.h"
 
 using namespace allocators;
 
 // Example class to be managed by the allocator
 struct TestClass
 {
-  int value;
+  int value[1024];
 
   TestClass(int val) : value(val)
   {
   }
 };
 
-static const constexpr int64_t MemoryAllowance = 10 * 1024 * 1024;
+static const constexpr int64_t MemoryAllowance = 1 * 1024 * 1024;
+static const constexpr int64_t MemoryLimitStep = MemoryAllowance / 100;
 
 // Test Fixture for AtomicCounterAllocator
 class CountingAllocatorTest : public ::testing::Test
@@ -49,7 +52,8 @@ class CountingAllocatorTest : public ::testing::Test
 
   // Constructor
   CountingAllocatorTest()
-   : allocatedMemory(MemoryAllowance), allocator(&allocatedMemory, MemoryAllowance / 100)
+   : allocatedMemory(MemoryAllowance)
+   , allocator(&allocatedMemory, MemoryAllowance / 100, MemoryAllowance / 100)
   {
   }
 
@@ -63,7 +67,15 @@ TEST_F(CountingAllocatorTest, Allocation)
   const std::size_t numObjects = 5;
   TestClass* ptr = allocator.allocate(numObjects);
   EXPECT_NE(ptr, nullptr);
-  EXPECT_EQ(allocatedMemory.load(), MemoryAllowance - numObjects * static_cast<int64_t>(sizeof(TestClass)));
+  if (MemoryLimitStep > numObjects * static_cast<int64_t>(sizeof(TestClass)))
+  {
+    EXPECT_EQ(allocatedMemory.load() - allocator.getCurrentLocalMemoryUsage(),
+              MemoryAllowance - numObjects * static_cast<int64_t>(sizeof(TestClass)));
+  }
+  else
+  {
+    EXPECT_EQ(allocatedMemory.load(), MemoryAllowance - numObjects * static_cast<int64_t>(sizeof(TestClass)));
+  }
   allocator.deallocate(ptr, numObjects);
 }
 
@@ -72,8 +84,15 @@ TEST_F(CountingAllocatorTest, Deallocation)
 {
   const std::size_t numObjects = 3;
   TestClass* ptr = allocator.allocate(numObjects);
-  EXPECT_EQ(allocatedMemory.load(), MemoryAllowance - numObjects * static_cast<int64_t>(sizeof(TestClass)));
-
+  if (MemoryLimitStep > numObjects * static_cast<int64_t>(sizeof(TestClass)))
+  {
+    EXPECT_EQ(allocatedMemory.load() - allocator.getCurrentLocalMemoryUsage(),
+              MemoryAllowance - numObjects * static_cast<int64_t>(sizeof(TestClass)));
+  }
+  else
+  {
+    EXPECT_EQ(allocatedMemory.load(), MemoryAllowance - numObjects * static_cast<int64_t>(sizeof(TestClass)));
+  }
   allocator.deallocate(ptr, numObjects);
   EXPECT_EQ(allocatedMemory.load(), MemoryAllowance);
 }
@@ -94,34 +113,42 @@ TEST_F(CountingAllocatorTest, AllocatorEquality)
 TEST_F(CountingAllocatorTest, AllocateSharedUsesAllocator)
 {
   // Create a shared_ptr using allocate_shared with the custom allocator
-  std::shared_ptr<TestClass> ptr = std::allocate_shared<TestClass>(allocator, 100);
+  CountingAllocator<TestClass> allocatorSmallerStep(&allocatedMemory, MemoryAllowance / 100,
+                                                    MemoryAllowance / 1000);
+  std::shared_ptr<TestClass> ptr1 = std::allocate_shared<TestClass>(allocatorSmallerStep, 100);
+  std::shared_ptr<TestClass> ptr2 = std::allocate_shared<TestClass>(allocatorSmallerStep, 100);
+  std::shared_ptr<TestClass> ptr3 = std::allocate_shared<TestClass>(allocatorSmallerStep, 100);
 
   // Check that the counter has increased by the size of TestClass plus control block
   // Exact size depends on the implementation, so we verify it's at least sizeof(TestClass)
-  EXPECT_LE(allocatedMemory.load(), MemoryAllowance - static_cast<int64_t>(sizeof(TestClass)));
+  EXPECT_LE(allocatedMemory.load(), MemoryAllowance - 3 * static_cast<int64_t>(sizeof(TestClass)));
 
   // Reset the shared_ptr and check that the counter decreases appropriately
-  ptr.reset();
+  ptr1.reset();
+  ptr2.reset();
+  ptr3.reset();
   // After deallocation, the counter should return to zero
   EXPECT_EQ(allocatedMemory.load(), MemoryAllowance);
 
-  auto deleter = [this](TestClass* ptr) { this->allocator.deallocate(ptr, 1); };
-  ptr.reset(allocator.allocate(1), deleter);
-  EXPECT_LE(allocatedMemory.load(), MemoryAllowance - static_cast<int64_t>(sizeof(TestClass)));
+  // auto deleter = [&allocatorSmallerStep](TestClass* ptr) { allocatorSmallerStep.deallocate(ptr, 1); };
+  // ptr1.reset(allocatorSmallerStep.allocate(3), deleter);
+  // EXPECT_LE(allocatedMemory.load(), MemoryAllowance - static_cast<int64_t>(sizeof(TestClass)));
 
-  ptr.reset();
+  ptr1.reset();
   EXPECT_EQ(allocatedMemory.load(), MemoryAllowance);
 
-  size_t allocSize = 16ULL * rowgroup::rgCommonSize;
-  auto buf = boost::allocate_shared<rowgroup::RGDataBufType>(allocator, allocSize);
+  using RGDataBufType = uint8_t[];
+  size_t allocSize = 16ULL * 8192;
+  auto buf = boost::allocate_shared<RGDataBufType>(allocator, allocSize);
+
   EXPECT_LE(allocatedMemory.load(), MemoryAllowance - allocSize);
 
   buf.reset();
   EXPECT_EQ(allocatedMemory.load(), MemoryAllowance);
 
-  CountingAllocator<rowgroup::RGDataBufType> allocator1(&allocatedMemory, MemoryAllowance / 100);
-  std::optional<CountingAllocator<rowgroup::RGDataBufType>> allocator2(allocator1);
-  auto buf1 = boost::allocate_shared<rowgroup::RGDataBufType>(*allocator2, allocSize);
+  CountingAllocator<RGDataBufType> allocator1(&allocatedMemory, MemoryAllowance / 100, MemoryAllowance / 100);
+  std::optional<CountingAllocator<RGDataBufType>> allocator2(allocator1);
+  auto buf1 = boost::allocate_shared<RGDataBufType>(*allocator2, allocSize);
   EXPECT_LE(allocatedMemory.load(), MemoryAllowance - allocSize);
 
   buf1.reset();
@@ -136,11 +163,27 @@ TEST_F(CountingAllocatorTest, ThreadSafety)
 
   auto worker = [this]()
   {
+    std::vector<TestClass*> ptrs;
+    CountingAllocator<TestClass> allocatorLocal(&allocatedMemory, MemoryAllowance / 100,
+                                                MemoryAllowance / 1000);
     for (std::size_t i = 0; i < allocationsPerThread; ++i)
     {
-      TestClass* ptr = allocator.allocate(1);
-      allocator.deallocate(ptr, 1);
+      ptrs.push_back(allocatorLocal.allocate(1));
     }
+
+    int64_t usedMemory = allocationsPerThread * sizeof(TestClass);
+    EXPECT_EQ(allocatorLocal.getCurrentLocalMemoryUsage(), allocationsPerThread * sizeof(TestClass));
+    EXPECT_GE(usedMemory - allocatorLocal.getlastMemoryLimitCheckpoint(), 0LL);
+    EXPECT_LE(allocatedMemory.load(), MemoryAllowance - allocatorLocal.getlastMemoryLimitCheckpoint());
+
+    for (auto* ptr : ptrs)
+    {
+      allocatorLocal.deallocate(ptr, 1);
+    }
+
+    EXPECT_EQ(allocatorLocal.getCurrentLocalMemoryUsage(), 0);
+    EXPECT_EQ(allocatorLocal.getlastMemoryLimitCheckpoint(), 0);
+    EXPECT_GE(allocatedMemory.load(), allocationsPerThread * sizeof(TestClass));
   };
 
   std::vector<std::thread> threads;
@@ -156,7 +199,7 @@ TEST_F(CountingAllocatorTest, ThreadSafety)
     th.join();
   }
 
-  // After all allocations and deallocations, the counter should be zero
+  // After all allocations and deallocations, the counter should be zero minus the remainder
   EXPECT_EQ(allocatedMemory.load(), MemoryAllowance);
 }
 
@@ -172,8 +215,8 @@ TEST_F(CountingAllocatorTest, AllocateZeroObjects)
 
 TEST_F(CountingAllocatorTest, CopyAssignable)
 {
-    CountingAllocator<TestClass> allocator1(&allocatedMemory);
-    CountingAllocator<TestClass> allocator2(&allocatedMemory);
-    allocator1 = allocator2;
-    EXPECT_EQ(allocator1, allocator2);
+  CountingAllocator<TestClass> allocator1(&allocatedMemory);
+  CountingAllocator<TestClass> allocator2(&allocatedMemory);
+  allocator1 = allocator2;
+  EXPECT_EQ(allocator1, allocator2);
 }