Implement MonotonicBufferResource

Reviewers: mtomic, mferencevic, msantl

Reviewed By: mtomic

Subscribers: pullbot

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

src/utils/CMakeLists.txt

@@ -1,6 +1,7 @@
 set(utils_src_files
     demangle.cpp
     file.cpp
+    memory.cpp
     signals.cpp
     thread.cpp
     thread/sync.cpp

src/utils/memory.cpp

@@ -0,0 +1,115 @@
+#include "utils/memory.hpp"
+
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <new>
+#include <type_traits>
+
+namespace utils {
+
+// MonotonicBufferResource
+
+namespace {
+
+// NOTE: std::bad_alloc has no constructor accepting a message, so we wrap our
+// exceptions in this class.
+class BadAlloc final : public std::bad_alloc {
+  std::string msg_;
+
+ public:
+  explicit BadAlloc(const std::string &msg) : msg_(msg) {}
+
+  const char *what() const noexcept override { return msg_.c_str(); }
+};
+
+size_t GrowMonotonicBuffer(size_t current_size, size_t max_size) {
+  double next_size = current_size * 1.34;
+  if (next_size >= static_cast<double>(max_size)) {
+    // Overflow, clamp to max_size
+    return max_size;
+  }
+  return std::ceil(next_size);
+}
+
+}  // namespace
+
+MonotonicBufferResource::MonotonicBufferResource(size_t initial_size)
+    : initial_size_(initial_size) {}
+
+MonotonicBufferResource::MonotonicBufferResource(size_t initial_size,
+                                                 MemoryResource *memory)
+    : memory_(memory), initial_size_(initial_size) {}
+
+MonotonicBufferResource::MonotonicBufferResource(
+    MonotonicBufferResource &&other) noexcept
+    : memory_(other.memory_),
+      current_buffer_(other.current_buffer_),
+      initial_size_(other.initial_size_),
+      allocated_(other.allocated_) {
+  other.current_buffer_ = nullptr;
+}
+
+MonotonicBufferResource &MonotonicBufferResource::operator=(
+    MonotonicBufferResource &&other) noexcept {
+  MonotonicBufferResource tmp(std::move(other));
+  std::swap(*this, tmp);
+  return *this;
+}
+
+void MonotonicBufferResource::Release() {
+  for (auto *b = current_buffer_; b;) {
+    auto *next = b->next;
+    b->~Buffer();
+    memory_->Deallocate(b, sizeof(b) + b->capacity);
+    b = next;
+  }
+  current_buffer_ = nullptr;
+  allocated_ = 0U;
+}
+
+void *MonotonicBufferResource::DoAllocate(size_t bytes, size_t alignment) {
+  static_assert(std::is_same_v<size_t, uintptr_t>);
+  const bool is_pow2 = alignment != 0U && (alignment & (alignment - 1U)) == 0U;
+  if (bytes == 0U || !is_pow2) throw BadAlloc("Invalid allocation request");
+  if (alignment > alignof(std::max_align_t))
+    alignment = alignof(std::max_align_t);
+
+  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;
+    // 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};
+    allocated_ = 0;
+  };
+
+  if (!current_buffer_) push_current_buffer(initial_size_);
+  char *buffer_head = current_buffer_->data() + allocated_;
+  void *aligned_ptr = buffer_head;
+  size_t available = current_buffer_->capacity - allocated_;
+  if (!std::align(alignment, bytes, aligned_ptr, available)) {
+    // Not enough memory, so allocate a new block with aligned data.
+    push_current_buffer(next_buffer_size_);
+    aligned_ptr = buffer_head = current_buffer_->data();
+    next_buffer_size_ = GrowMonotonicBuffer(
+        next_buffer_size_, std::numeric_limits<size_t>::max() - sizeof(Buffer));
+  }
+  if (reinterpret_cast<char *>(aligned_ptr) < buffer_head)
+    throw BadAlloc("Allocation alignment overflow");
+  if (reinterpret_cast<char *>(aligned_ptr) + bytes <= aligned_ptr)
+    throw BadAlloc("Allocation size overflow");
+  allocated_ =
+      reinterpret_cast<char *>(aligned_ptr) - current_buffer_->data() + bytes;
+  return aligned_ptr;
+}
+
+// MonotonicBufferResource END
+
+}  // namespace utils

src/utils/memory.hpp

@@ -5,7 +5,7 @@
 #pragma once
 
 #include <cstddef>
-// Although <memory_resource> is in C++17, gcc libstdc++ still need to
+// Although <memory_resource> is in C++17, gcc libstdc++ still needs to
 // implement it fully. It should be available in the next major release
 // version, i.e. gcc 9.x.
 #include <experimental/memory_resource>
@@ -17,6 +17,18 @@ class MemoryResource {
  public:
   virtual ~MemoryResource() {}
 
+  /// Allocate storage with a size of at least `bytes` bytes.
+  ///
+  /// The returned storage is aligned to to `alignment` clamped to
+  /// `alignof(std::max_align_t)`. This means that it is valid to request larger
+  /// alignment, but the storage will actually be aligned to
+  /// `alignof(std::max_align_t)`.
+  //
+  /// Additionaly, `alignment` must be a power of 2, if it is not
+  /// `std::bad_alloc` is thrown.
+  ///
+  /// @throw std::bad_alloc if the requested storage and alignment combination
+  ///     cannot be obtained.
   void *Allocate(size_t bytes, size_t alignment = alignof(std::max_align_t)) {
     return DoAllocate(bytes, alignment);
   }
@@ -123,4 +135,64 @@ inline MemoryResource *NewDeleteResource() noexcept {
   return &memory;
 }
 
+/// MemoryResource which releases the memory only when the resource is
+/// destroyed.
+///
+/// MonotonicBufferResource is not thread-safe!
+///
+/// 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.
+class MonotonicBufferResource final : public MemoryResource {
+ public:
+  /// Construct the resource with the buffer size of at least `initial_size`.
+  ///
+  /// Uses the default NewDeleteResource for requesting buffer memory.
+  explicit MonotonicBufferResource(size_t initial_size);
+
+  /// Construct the resource with the buffer size of at least `initial_size` and
+  /// use the `memory` as the upstream resource.
+  MonotonicBufferResource(size_t initial_size, MemoryResource *memory);
+
+  MonotonicBufferResource(const MonotonicBufferResource &) = delete;
+  MonotonicBufferResource &operator=(const MonotonicBufferResource &) = delete;
+
+  MonotonicBufferResource(MonotonicBufferResource &&other) noexcept;
+  MonotonicBufferResource &operator=(MonotonicBufferResource &&other) noexcept;
+
+  ~MonotonicBufferResource() override { Release(); }
+
+  /// Release all allocated memory by calling Deallocate on upstream
+  /// MemoryResource.
+  ///
+  /// All memory is released even though Deallocate may have not been called on
+  /// this instance for some of the allocated blocks.
+  void Release();
+
+  MemoryResource *GetUpstreamResource() const { return memory_; }
+
+ private:
+  struct Buffer {
+    Buffer *next;
+    size_t capacity;
+    char *data() { return reinterpret_cast<char *>(this) + sizeof(Buffer); }
+  };
+
+  MemoryResource *memory_{NewDeleteResource()};
+  Buffer *current_buffer_{nullptr};
+  size_t initial_size_{0U};
+  size_t next_buffer_size_{initial_size_};
+  size_t allocated_{0U};
+
+  void *DoAllocate(size_t bytes, size_t alignment) override;
+
+  void DoDeallocate(void *, size_t, size_t) override {}
+
+  bool DoIsEqual(const MemoryResource &other) const noexcept override {
+    return this == &other;
+  }
+};
+
 }  // namespace utils

tests/unit/CMakeLists.txt

@@ -351,6 +351,9 @@ target_link_libraries(${test_prefix}utils_file mg-utils)
 add_unit_test(utils_math.cpp)
 target_link_libraries(${test_prefix}utils_math mg-utils)
 
+add_unit_test(utils_memory.cpp)
+target_link_libraries(${test_prefix}utils_memory mg-utils)
+
 add_unit_test(utils_on_scope_exit.cpp)
 target_link_libraries(${test_prefix}utils_on_scope_exit mg-utils)
 

tests/unit/utils_memory.cpp

@@ -0,0 +1,148 @@
+#include <cstdint>
+#include <cstring>
+#include <limits>
+
+#include <gtest/gtest.h>
+
+#include "utils/memory.hpp"
+
+class TestMemory final : public utils::MemoryResource {
+ public:
+  size_t new_count_{0};
+  size_t delete_count_{0};
+
+ private:
+  void *DoAllocate(size_t bytes, size_t alignment) override {
+    new_count_++;
+    EXPECT_TRUE(alignment != 0U && (alignment & (alignment - 1U)) == 0U)
+        << "Alignment must be power of 2";
+    EXPECT_TRUE(alignment <= alignof(std::max_align_t));
+    EXPECT_NE(bytes, 0);
+    const size_t pad_size = 32;
+    EXPECT_TRUE(bytes + pad_size > bytes) << "TestMemory size overflow";
+    EXPECT_TRUE(bytes + pad_size + alignment > bytes + alignment)
+        << "TestMemory size overflow";
+    EXPECT_TRUE(2U * alignment > alignment) << "TestMemory alignment overflow";
+    // Allocate a block containing extra alignment and pad_size bytes, but
+    // aligned to 2 * alignment. Then we can offset the ptr so that it's never
+    // aligned to 2 * alignment. This ought to make allocator alignment issues
+    // more obvious.
+    void *ptr = utils::NewDeleteResource()->Allocate(
+        alignment + bytes + pad_size, 2U * alignment);
+    // Clear allocated memory to 0xFF, marking the invalid region.
+    memset(ptr, 0xFF, alignment + bytes + pad_size);
+    // Offset the ptr so it's not aligned to 2 * alignment, but still aligned to
+    // alignment.
+    ptr = static_cast<char *>(ptr) + alignment;
+    // Clear the valid region to 0x00, so that we can more easily test that the
+    // allocator is doing the right thing.
+    memset(ptr, 0, bytes);
+    return ptr;
+  }
+
+  void DoDeallocate(void *ptr, size_t bytes, size_t alignment) override {
+    delete_count_++;
+    // Dealloate the original ptr, before alignment adjustment.
+    return utils::NewDeleteResource()->Deallocate(
+        static_cast<char *>(ptr) - alignment, bytes, alignment);
+  }
+
+  bool DoIsEqual(const utils::MemoryResource &other) const noexcept override {
+    return this == &other;
+  }
+};
+
+void *CheckAllocation(utils::MemoryResource *mem, size_t bytes,
+                      size_t alignment = alignof(std::max_align_t)) {
+  void *ptr = mem->Allocate(bytes, alignment);
+  if (alignment > alignof(std::max_align_t))
+    alignment = alignof(std::max_align_t);
+  EXPECT_TRUE(ptr);
+  EXPECT_EQ(reinterpret_cast<uintptr_t>(ptr) % alignment, 0)
+      << "Allocated misaligned pointer!";
+  // There should be no 0xFF bytes because they are either padded at the end of
+  // the allocated block or are found in already checked allocations.
+  EXPECT_FALSE(memchr(ptr, 0xFF, bytes)) << "Invalid memory region!";
+  // Mark the checked allocation with 0xFF bytes.
+  memset(ptr, 0xFF, bytes);
+  return ptr;
+}
+
+TEST(MonotonicBufferResource, AllocationWithinInitialSize) {
+  TestMemory test_mem;
+  {
+    utils::MonotonicBufferResource mem(1024, &test_mem);
+    void *fst_ptr = CheckAllocation(&mem, 24, 1);
+    void *snd_ptr = CheckAllocation(&mem, 1000, 1);
+    EXPECT_EQ(test_mem.new_count_, 1);
+    EXPECT_EQ(test_mem.delete_count_, 0);
+    mem.Deallocate(snd_ptr, 1000, 1);
+    mem.Deallocate(fst_ptr, 24, 1);
+    EXPECT_EQ(test_mem.delete_count_, 0);
+    mem.Release();
+    EXPECT_EQ(test_mem.delete_count_, 1);
+    CheckAllocation(&mem, 1024);
+    EXPECT_EQ(test_mem.new_count_, 2);
+    EXPECT_EQ(test_mem.delete_count_, 1);
+  }
+  EXPECT_EQ(test_mem.delete_count_, 2);
+}
+
+TEST(MonotonicBufferResource, AllocationOverInitialSize) {
+  TestMemory test_mem;
+  {
+    utils::MonotonicBufferResource mem(1024, &test_mem);
+    CheckAllocation(&mem, 1025, 1);
+    EXPECT_EQ(test_mem.new_count_, 1);
+  }
+  EXPECT_EQ(test_mem.delete_count_, 1);
+  {
+    utils::MonotonicBufferResource mem(1024, &test_mem);
+    // Test with large alignment
+    CheckAllocation(&mem, 1024, 1024);
+    EXPECT_EQ(test_mem.new_count_, 2);
+  }
+  EXPECT_EQ(test_mem.delete_count_, 2);
+}
+
+TEST(MonotonicBufferResource, AllocationOverCapacity) {
+  TestMemory test_mem;
+  {
+    utils::MonotonicBufferResource mem(1024, &test_mem);
+    CheckAllocation(&mem, 24, 1);
+    EXPECT_EQ(test_mem.new_count_, 1);
+    CheckAllocation(&mem, 1000, 64);
+    EXPECT_EQ(test_mem.new_count_, 2);
+    EXPECT_EQ(test_mem.delete_count_, 0);
+    mem.Release();
+    EXPECT_EQ(test_mem.new_count_, 2);
+    EXPECT_EQ(test_mem.delete_count_, 2);
+    CheckAllocation(&mem, 1025, 1);
+    EXPECT_EQ(test_mem.new_count_, 3);
+    CheckAllocation(&mem, 1023, 1);
+    // MonotonicBufferResource state after Release is called may or may not
+    // allocate a larger block right from the start (i.e. tracked buffer sizes
+    // before Release may be retained).
+    EXPECT_TRUE(test_mem.new_count_ >= 3);
+  }
+  EXPECT_TRUE(test_mem.delete_count_ >= 3);
+}
+
+TEST(MonotonicBufferResource, AllocationWithAlignmentNotPowerOf2) {
+  utils::MonotonicBufferResource mem(1024);
+  EXPECT_THROW(mem.Allocate(24, 3), std::bad_alloc);
+  EXPECT_THROW(mem.Allocate(24, 0), std::bad_alloc);
+}
+
+TEST(MonotonicBufferResource, AllocationWithSize0) {
+  utils::MonotonicBufferResource mem(1024);
+  EXPECT_THROW(mem.Allocate(0), std::bad_alloc);
+}
+
+TEST(MonotonicBufferResource, AllocationWithSizeOverflow) {
+  size_t max_size = std::numeric_limits<size_t>::max();
+  utils::MonotonicBufferResource mem(1024);
+  // Setup so that the next allocation aligning max_size causes overflow.
+  mem.Allocate(1, 1);
+  EXPECT_THROW(mem.Allocate(max_size, 4), std::bad_alloc);
+}