Support all power of 2 alignments in MonotonicBufferResource

Reviewers: mtomic, mferencevic

Reviewed By: mtomic

Subscribers: pullbot

Differential Revision: https://phabricator.memgraph.io/D2188

src/utils/math.hpp

@@ -15,7 +15,7 @@ static_assert(
 /// This function computes the log2 function on integer types. It is faster than
 /// the cmath `log2` function because it doesn't use floating point values for
 /// calculation.
-inline uint64_t Log2(uint64_t val) {
+constexpr inline uint64_t Log2(uint64_t val) {
   // The `clz` function is undefined when the passed value is 0 and the value of
   // `log` is `-inf` so we special case it here.
   if (val == 0) return 0;
@@ -26,14 +26,14 @@ inline uint64_t Log2(uint64_t val) {
 }
 
 /// Return `true` if `val` is a power of 2.
-inline bool IsPow2(uint64_t val) noexcept {
+constexpr inline bool IsPow2(uint64_t val) noexcept {
   return val != 0ULL && (val & (val - 1ULL)) == 0ULL;
 }
 
 /// Return `val` if it is power of 2, otherwise get the next power of 2 value.
 /// If `val` is sufficiently large, the next power of 2 value may not fit into
 /// the result type and you will get a wrapped value to 1ULL.
-inline uint64_t Ceil2(uint64_t val) noexcept {
+constexpr inline uint64_t Ceil2(uint64_t val) noexcept {
   if (val == 0ULL || val == 1ULL) return 1ULL;
   return 1ULL << (Log2(val - 1ULL) + 1ULL);
 }

src/utils/memory.cpp

@@ -69,9 +69,11 @@ MonotonicBufferResource &MonotonicBufferResource::operator=(
 void MonotonicBufferResource::Release() {
   for (auto *b = current_buffer_; b;) {
     auto *next = b->next;
-    auto capacity = b->capacity;
+    auto *ptr = b->data();
+    auto alloc_size = b->size();
+    auto alignment = b->alignment;
     b->~Buffer();
-    memory_->Deallocate(b, sizeof(*b) + capacity);
+    memory_->Deallocate(ptr, alloc_size, alignment);
     b = next;
   }
   current_buffer_ = nullptr;
@@ -80,22 +82,25 @@ void MonotonicBufferResource::Release() {
 
 void *MonotonicBufferResource::DoAllocate(size_t bytes, size_t alignment) {
   static_assert(std::is_same_v<size_t, uintptr_t>);
-  if (alignment > alignof(std::max_align_t))
-    throw BadAlloc(
-        "Alignment greater than alignof(std::max_align_t) is unsupported");
-
-  auto push_current_buffer = [this, bytes](size_t next_size) {
-    // Set capacity so that the bytes fit.
-    size_t capacity = next_size > bytes ? next_size : bytes;
+  auto push_current_buffer = [this, bytes, alignment](size_t next_size) {
+    // Set size so that the bytes fit.
+    size_t size = next_size > bytes ? next_size : bytes;
     // Handle the case when we need to align `Buffer::data` to a greater
-    // `alignment`. We will simply always allocate with
-    // alignof(std::max_align_t), and `sizeof(Buffer)` needs to be a multiple of
-    // that to keep `data` correctly aligned.
-    static_assert(sizeof(Buffer) % alignof(std::max_align_t) == 0);
-    size_t alloc_size = sizeof(Buffer) + capacity;
-    if (alloc_size <= capacity) throw BadAlloc("Allocation size overflow");
-    void *ptr = memory_->Allocate(alloc_size);
-    current_buffer_ = new (ptr) Buffer{current_buffer_, capacity};
+    // `alignment`. We will simply always allocate Ceil2(required_size), so that
+    // we can use the end of the allocated bytes for Buffer instance.
+    static_assert(IsPow2(alignof(Buffer)),
+                  "Buffer should not be a packed struct in order to be placed "
+                  "at the end of an allocation request");
+    size_t bytes = sizeof(Buffer) + size;
+    if (bytes < size) throw BadAlloc("Allocation size overflow");
+    size_t alloc_size = Ceil2(bytes);
+    if (alloc_size < bytes) throw BadAlloc("Allocation size overflow");
+    size_t alloc_align = std::max(alignment, alignof(std::max_align_t));
+    void *ptr = memory_->Allocate(alloc_size, alloc_align);
+    // Instantiate the Buffer at the end of the allocated block.
+    current_buffer_ =
+        new (reinterpret_cast<char *>(ptr) + alloc_size - sizeof(Buffer))
+            Buffer{current_buffer_, alloc_size - sizeof(Buffer), alloc_align};
     allocated_ = 0;
   };
 

src/utils/memory.hpp

@@ -283,14 +283,18 @@ inline MemoryResource *NewDeleteResource() noexcept {
 ///
 /// MonotonicBufferResource is not thread-safe!
 ///
-/// MonotonicBufferResource cannot handle alignment requests greater than
-/// `alignof(std::max_align_t)`!
-///
 /// It's meant to be used for very fast allocations in situations where memory
 /// is used to build objects and release them all at once. The class is
 /// constructed with initial buffer size for storing allocated objects. When the
 /// buffer is exhausted, a new one is requested from the upstream memory
 /// resource.
+///
+/// Note that each buffer of memory is actually a block of `Ceil2(size +
+/// sizeof(Buffer))` due to bookkeeping `Buffer` object being appended at the end.
+/// This means that if you use an `initial_size` of 1024 bytes, you will
+/// actually allocate `Ceil2(1024 + sizeof(Buffer))` which will be 2048 bytes.
+/// Therefore you will have `2048 - sizeof(Buffer)` bytes available before a new
+/// buffer will need to be allocated.
 class MonotonicBufferResource final : public MemoryResource {
  public:
   /// Construct the resource with the buffer size of at least `initial_size`.
@@ -333,7 +337,11 @@ class MonotonicBufferResource final : public MemoryResource {
   struct Buffer {
     Buffer *next;
     size_t capacity;
-    char *data() { return reinterpret_cast<char *>(this) + sizeof(Buffer); }
+    size_t alignment;
+    /// Get total allocated size.
+    size_t size() const { return sizeof(*this) + capacity; }
+    /// Get the pointer to data which is before the Buffer instance itself.
+    char *data() { return reinterpret_cast<char *>(this) - capacity; }
   };
 
   MemoryResource *memory_{NewDeleteResource()};

tests/unit/utils_memory.cpp

@@ -106,10 +106,10 @@ TEST(MonotonicBufferResource, AllocationOverInitialSize) {
 TEST(MonotonicBufferResource, AllocationOverCapacity) {
   TestMemory test_mem;
   {
-    utils::MonotonicBufferResource mem(1024, &test_mem);
+    utils::MonotonicBufferResource mem(1000, &test_mem);
     CheckAllocation(&mem, 24, 1);
     EXPECT_EQ(test_mem.new_count_, 1);
-    CheckAllocation(&mem, 1001);
+    CheckAllocation(&mem, 976);
     EXPECT_EQ(test_mem.new_count_, 2);
     EXPECT_EQ(test_mem.delete_count_, 0);
     mem.Release();
@@ -139,9 +139,9 @@ TEST(MonotonicBufferResource, AllocationWithSize0) {
 
 // NOLINTNEXTLINE(hicpp-special-member-functions)
 TEST(MonotonicBufferResource, AllocationWithAlignmentGreaterThanMaxAlign) {
-  utils::MonotonicBufferResource mem(1024);
-  EXPECT_THROW(mem.Allocate(24, 2U * alignof(std::max_align_t)),
-               std::bad_alloc);
+  TestMemory test_mem;
+  utils::MonotonicBufferResource mem(1024, &test_mem);
+  CheckAllocation(&mem, 24, 2U * alignof(std::max_align_t));
 }
 
 TEST(MonotonicBufferResource, AllocationWithSizeOverflow) {