diff --git a/mooncake-store/include/offset_allocator/offset_allocator.hpp b/mooncake-store/include/offset_allocator/offset_allocator.hpp index 4b3dacc4..832b14ae 100644 --- a/mooncake-store/include/offset_allocator/offset_allocator.hpp +++ b/mooncake-store/include/offset_allocator/offset_allocator.hpp @@ -26,8 +26,19 @@ static constexpr uint32 NUM_LEAF_BINS = NUM_TOP_BINS * BINS_PER_LEAF; struct OffsetAllocation { static constexpr uint32 NO_SPACE = 0xffffffff; + private: uint32 offset = NO_SPACE; NodeIndex metadata = NO_SPACE; // internal: node index + + public: + OffsetAllocation(uint32 offset_param, NodeIndex metadata_param) + : offset(offset_param), metadata(metadata_param) {} + // The real offset could be larger than uint32, so we need to cast it to + // uint64_t + uint64_t getOffset() const { return static_cast(offset); } + bool isNoSpace() const { return offset == NO_SPACE; } + + friend class __Allocator; }; struct OffsetAllocStorageReport { @@ -100,17 +111,22 @@ std::ostream& operator<<(std::ostream& os, // Wrapper class for __Allocator, it 1) supports thread-safe allocation and // deallocation, 2) supports creating a buffer or allocating a memory region // that is larger than the largest bin size (3.75GB). The __allocator class is -// also optimized to round up the allocated size to a bin size. This will -// a) slightly decrease the memory utilization ratio in general cases, b) makes -// no difference when the allocated size is equal to a bin size, c) largely -// improve the memory utilization ratio when the allocated size is mostly -// uniform and not equal to any bin size. +// also optimized to: +// 1) round up the allocated size to a bin size. This will a) slightly decrease +// the memory utilization ratio in general cases, b) makes no difference when +// the allocated size is equal to a bin size, c) largely improve the memory +// utilization ratio when the allocated size is mostly uniform and not equal to +// any bin size. +// 2) dynamically adjust the capacity of the allocator to the allocated size. +// This will a) reduce the memory consumption in general cases, b) auto +// increase the capacity in case there are a lot of small regions to be +// allocated. class OffsetAllocator : public std::enable_shared_from_this { public: // Factory method to create shared_ptr - static std::shared_ptr create(uint64_t base, size_t size, - uint32 maxAllocs = 128 * - 1024); + static std::shared_ptr create( + uint64_t base, size_t size, uint32 init_capacity = 128 * 1024, + uint32 max_capacity = (1 << 20)); // Disable copy constructor and copy assignment OffsetAllocator(const OffsetAllocator&) = delete; @@ -153,7 +169,7 @@ class OffsetAllocator : public std::enable_shared_from_this { const uint64_t m_base; // The real offset and size of the allocated memory need to be multiplied by // m_multiplier - const uint64_t m_multiplier; + const uint64_t m_multiplier_bits; const uint64_t m_capacity; mutable Mutex m_mutex; @@ -162,12 +178,13 @@ class OffsetAllocator : public std::enable_shared_from_this { uint64_t m_allocated_num GUARDED_BY(m_mutex) = 0; // Private constructor - use create() factory method instead - OffsetAllocator(uint64_t base, size_t size, uint32 maxAllocs = 128 * 1024); + OffsetAllocator(uint64_t base, size_t size, uint32 init_capacity, + uint32 max_capacity); }; class __Allocator { public: - __Allocator(uint32 size, uint32 maxAllocs = 128 * 1024); + __Allocator(uint32 size, uint32 init_capacity, uint32 max_capacity); __Allocator(__Allocator&& other); ~__Allocator(); void reset(); @@ -196,7 +213,8 @@ class __Allocator { }; uint32 m_size; - uint32 m_maxAllocs; + uint32 m_current_capacity; + uint32 m_max_capacity; uint32 m_freeStorage; uint32 m_usedBinsTop; diff --git a/mooncake-store/src/allocator.cpp b/mooncake-store/src/allocator.cpp index 8b75ca4a..d73c7091 100644 --- a/mooncake-store/src/allocator.cpp +++ b/mooncake-store/src/allocator.cpp @@ -118,14 +118,21 @@ OffsetBufferAllocator::OffsetBufferAllocator(std::string segment_name, << " size=" << size; try { - uint32_t max_allocs = size < (1ull << 32) - ? size / 4096 - : 1024 * 1024; // min(size / 4K, 1M) - max_allocs = - std::max(max_allocs, 1024u * 64u); // at least 64K allocations + // 1k <= init_capacity <= 64k + uint64_t init_capacity = size / 4096; + init_capacity = std::max(init_capacity, static_cast(1024)); + init_capacity = + std::min(init_capacity, static_cast(64 * 1024)); + // 1M <= max_capacity <= 64G / 1K = 64M + uint64_t max_capacity = size / 1024; + max_capacity = + std::max(max_capacity, static_cast(1024 * 1024)); + max_capacity = + std::min(max_capacity, static_cast(64 * 1024 * 1024)); // Create the offset allocator - offset_allocator_ = - offset_allocator::OffsetAllocator::create(base, size, max_allocs); + offset_allocator_ = offset_allocator::OffsetAllocator::create( + base, size, static_cast(init_capacity), + static_cast(max_capacity)); if (!offset_allocator_) { LOG(ERROR) << "status=failed_to_create_offset_allocator"; throw std::runtime_error("Failed to create offset allocator"); diff --git a/mooncake-store/src/offset_allocator.cpp b/mooncake-store/src/offset_allocator.cpp index 17f99ba1..ddf04dc0 100644 --- a/mooncake-store/src/offset_allocator.cpp +++ b/mooncake-store/src/offset_allocator.cpp @@ -126,20 +126,22 @@ uint32 findLowestSetBitAfter(uint32 bitMask, uint32 startBitIndex) { } // __Allocator... -__Allocator::__Allocator(uint32 size, uint32 maxAllocs) +__Allocator::__Allocator(uint32 size, uint32 init_capacity, uint32 max_capacity) : m_size(size), - m_maxAllocs(maxAllocs), + m_current_capacity(init_capacity), + m_max_capacity(std::max(init_capacity, max_capacity)), m_nodes(nullptr), m_freeNodes(nullptr) { if (sizeof(NodeIndex) == 2) { - ASSERT(maxAllocs <= 65536); + ASSERT(m_max_capacity <= 65536); } reset(); } __Allocator::__Allocator(__Allocator&& other) : m_size(other.m_size), - m_maxAllocs(other.m_maxAllocs), + m_current_capacity(other.m_current_capacity), + m_max_capacity(other.m_max_capacity), m_freeStorage(other.m_freeStorage), m_usedBinsTop(other.m_usedBinsTop), m_nodes(other.m_nodes), @@ -151,14 +153,15 @@ __Allocator::__Allocator(__Allocator&& other) other.m_nodes = nullptr; other.m_freeNodes = nullptr; other.m_freeOffset = 0; - other.m_maxAllocs = 0; + other.m_current_capacity = 0; + other.m_max_capacity = 0; other.m_usedBinsTop = 0; } void __Allocator::reset() { m_freeStorage = 0; m_usedBinsTop = 0; - m_freeOffset = m_maxAllocs - 1; + m_freeOffset = 0; for (uint32 i = 0; i < NUM_TOP_BINS; i++) m_usedBins[i] = 0; @@ -167,12 +170,12 @@ void __Allocator::reset() { if (m_nodes) delete[] m_nodes; if (m_freeNodes) delete[] m_freeNodes; - m_nodes = new Node[m_maxAllocs]; - m_freeNodes = new NodeIndex[m_maxAllocs]; + m_nodes = new Node[m_max_capacity]; + m_freeNodes = new NodeIndex[m_max_capacity]; // Freelist is a stack. Nodes in inverse order so that [0] pops first. - for (uint32 i = 0; i < m_maxAllocs; i++) { - m_freeNodes[i] = m_maxAllocs - i - 1; + for (uint32 i = 0; i < m_current_capacity; i++) { + m_freeNodes[i] = i; } // Start state: Whole storage as one big node @@ -187,9 +190,13 @@ __Allocator::~__Allocator() { OffsetAllocation __Allocator::allocate(uint32 size) { // Out of allocations? - if (m_freeOffset == 0) { - return {.offset = OffsetAllocation::NO_SPACE, - .metadata = OffsetAllocation::NO_SPACE}; + if (m_freeOffset == m_max_capacity) { + return OffsetAllocation(OffsetAllocation::NO_SPACE, + OffsetAllocation::NO_SPACE); + } + if (m_freeOffset == m_current_capacity) { + m_freeNodes[m_current_capacity] = m_current_capacity; + m_current_capacity++; } // Round up to bin index to ensure that alloc >= bin @@ -214,8 +221,8 @@ OffsetAllocation __Allocator::allocate(uint32 size) { // Out of space? if (topBinIndex == OffsetAllocation::NO_SPACE) { - return {.offset = OffsetAllocation::NO_SPACE, - .metadata = OffsetAllocation::NO_SPACE}; + return OffsetAllocation(OffsetAllocation::NO_SPACE, + OffsetAllocation::NO_SPACE); } // All leaf bins here fit the alloc, since the top bin was rounded up. @@ -277,7 +284,7 @@ OffsetAllocation __Allocator::allocate(uint32 size) { node.neighborNext = newNodeIndex; } - return {.offset = node.dataOffset, .metadata = nodeIndex}; + return OffsetAllocation(node.dataOffset, nodeIndex); } void __Allocator::free(OffsetAllocation allocation) { @@ -328,9 +335,9 @@ void __Allocator::free(OffsetAllocation allocation) { // Insert the removed node to freelist #ifdef DEBUG_VERBOSE printf("Putting node %u into freelist[%u] (free)\n", nodeIndex, - m_freeOffset + 1); + m_freeOffset - 1); #endif - m_freeNodes[++m_freeOffset] = nodeIndex; + m_freeNodes[--m_freeOffset] = nodeIndex; // Insert the (combined) free node to bin uint32 combinedNodeIndex = insertNodeIntoBin(size, offset); @@ -363,9 +370,9 @@ uint32 __Allocator::insertNodeIntoBin(uint32 size, uint32 dataOffset) { // Take a freelist node and insert on top of the bin linked list (next = old // top) uint32 topNodeIndex = m_binIndices[binIndex]; - uint32 nodeIndex = m_freeNodes[m_freeOffset--]; + uint32 nodeIndex = m_freeNodes[m_freeOffset++]; #ifdef DEBUG_VERBOSE - printf("Getting node %u from freelist[%u]\n", nodeIndex, m_freeOffset + 1); + printf("Getting node %u from freelist[%u]\n", nodeIndex, m_freeOffset - 1); #endif m_nodes[nodeIndex] = {.dataOffset = dataOffset, .dataSize = size, @@ -420,9 +427,9 @@ void __Allocator::removeNodeFromBin(uint32 nodeIndex) { // Insert the node to freelist #ifdef DEBUG_VERBOSE printf("Putting node %u into freelist[%u] (removeNodeFromBin)\n", nodeIndex, - m_freeOffset + 1); + m_freeOffset - 1); #endif - m_freeNodes[++m_freeOffset] = nodeIndex; + m_freeNodes[--m_freeOffset] = nodeIndex; m_freeStorage -= node.dataSize; #ifdef DEBUG_VERBOSE @@ -443,7 +450,7 @@ OffsetAllocStorageReport __Allocator::storageReport() const { uint32 freeStorage = 0; // Out of allocations? -> Zero free space - if (m_freeOffset > 0) { + if (m_freeOffset < m_max_capacity) { freeStorage = m_freeStorage; if (m_usedBinsTop) { uint32 topBinIndex = 31 - lzcnt_nonzero(m_usedBinsTop); @@ -527,24 +534,30 @@ OffsetAllocationHandle::~OffsetAllocationHandle() { // Helper function to calculate the multiplier static uint64_t calculateMultiplier(size_t size) { - uint64_t multiplier = 1; - for (; SmallFloat::MAX_BIN_SIZE < size / multiplier; multiplier *= 2) { + uint64_t multiplier_bits = 0; + for (; SmallFloat::MAX_BIN_SIZE < (size >> multiplier_bits); + multiplier_bits++) { } - return multiplier; + return multiplier_bits; } // Thread-safe OffsetAllocator implementation std::shared_ptr OffsetAllocator::create(uint64_t base, size_t size, - uint32 maxAllocs) { + uint32 init_capacity, + uint32 max_capacity) { // Use a custom deleter to allow private constructor return std::shared_ptr( - new OffsetAllocator(base, size, maxAllocs)); + new OffsetAllocator(base, size, init_capacity, max_capacity)); } -OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, uint32 maxAllocs) - : m_base(base), m_multiplier(calculateMultiplier(size)), m_capacity(size) { - m_allocator = std::make_unique<__Allocator>(size / m_multiplier, maxAllocs); +OffsetAllocator::OffsetAllocator(uint64_t base, size_t size, + uint32 init_capacity, uint32 max_capacity) + : m_base(base), + m_multiplier_bits(calculateMultiplier(size)), + m_capacity(size) { + m_allocator = std::make_unique<__Allocator>(size >> m_multiplier_bits, + init_capacity, max_capacity); } std::optional OffsetAllocator::allocate(size_t size) { @@ -558,14 +571,17 @@ std::optional OffsetAllocator::allocate(size_t size) { } size_t fake_size = - m_multiplier > 1 ? (size + m_multiplier - 1) / m_multiplier : size; + m_multiplier_bits > 0 + ? ((size + (static_cast(1) << m_multiplier_bits) - 1u) >> + m_multiplier_bits) + : size; if (fake_size > SmallFloat::MAX_BIN_SIZE) { return std::nullopt; } OffsetAllocation allocation = m_allocator->allocate(fake_size); - if (allocation.offset == OffsetAllocation::NO_SPACE) { + if (allocation.isNoSpace()) { // Log metrics to help understand why allocation failed // Note: We're already holding m_mutex, so use internal method OffsetAllocatorMetrics metrics = get_metrics_internal(); @@ -579,9 +595,9 @@ std::optional OffsetAllocator::allocate(size_t size) { m_allocated_num++; // Use shared_from_this to get a shared_ptr to this OffsetAllocator - return OffsetAllocationHandle(shared_from_this(), allocation, - m_base + allocation.offset * m_multiplier, - size); + return OffsetAllocationHandle( + shared_from_this(), allocation, + m_base + (allocation.getOffset() << m_multiplier_bits), size); } OffsetAllocStorageReport OffsetAllocator::storageReport() const { @@ -590,8 +606,8 @@ OffsetAllocStorageReport OffsetAllocator::storageReport() const { return {0, 0}; } OffsetAllocStorageReport report = m_allocator->storageReport(); - return {report.totalFreeSpace * m_multiplier, - report.largestFreeRegion * m_multiplier}; + return {report.totalFreeSpace << m_multiplier_bits, + report.largestFreeRegion << m_multiplier_bits}; } OffsetAllocStorageReportFull OffsetAllocator::storageReportFull() const { @@ -603,7 +619,7 @@ OffsetAllocStorageReportFull OffsetAllocator::storageReportFull() const { OffsetAllocStorageReportFull report = m_allocator->storageReportFull(); for (uint32 i = 0; i < NUM_LEAF_BINS; i++) { report.freeRegions[i] = { - .size = report.freeRegions[i].size * m_multiplier, + .size = report.freeRegions[i].size << m_multiplier_bits, .count = report.freeRegions[i].count}; } return report; @@ -617,11 +633,12 @@ OffsetAllocatorMetrics OffsetAllocator::get_metrics_internal() const { // Get basic storage report OffsetAllocStorageReport basic_report = m_allocator->storageReport(); return { - m_allocated_size, // allocated_size_ - m_allocated_num, // allocated_num_ - basic_report.largestFreeRegion * m_multiplier, // largest_free_region_ - basic_report.totalFreeSpace * m_multiplier, // total_free_space_ - m_capacity, // capacity + m_allocated_size, // allocated_size_ + m_allocated_num, // allocated_num_ + basic_report.largestFreeRegion + << m_multiplier_bits, // largest_free_region_ + basic_report.totalFreeSpace << m_multiplier_bits, // total_free_space_ + m_capacity, // capacity }; } @@ -653,8 +670,7 @@ std::ostream& operator<<(std::ostream& os, << ", allocs=" << metrics.allocated_num_ << ", capacity=" << mooncake::byte_size_to_string(metrics.capacity) << ", utilization=" << std::fixed << std::setprecision(1) << utilization - << "%" - << ", free_space=" + << "%" << ", free_space=" << mooncake::byte_size_to_string(metrics.total_free_space_) << ", largest_free=" << mooncake::byte_size_to_string(metrics.largest_free_region_) << "}"; diff --git a/mooncake-store/tests/offset_allocator_test.cpp b/mooncake-store/tests/offset_allocator_test.cpp index 82d4df2e..f265c332 100644 --- a/mooncake-store/tests/offset_allocator_test.cpp +++ b/mooncake-store/tests/offset_allocator_test.cpp @@ -104,7 +104,8 @@ class AllocatorWrapper : public std::enable_shared_from_this { public: // Constructor AllocatorWrapper(uint64_t base, size_t size, uint32 maxAllocs = 128 * 1024) - : m_allocator(OffsetAllocator::create(base, size, maxAllocs)), + : m_allocator( + OffsetAllocator::create(base, size, maxAllocs / 2, maxAllocs)), m_base(base), m_buffer_size(size) { // The allocator is created with the specified base and size @@ -408,7 +409,7 @@ TEST_F(OffsetAllocatorTest, RepeatedLargeSizeAllocation) { } } -// Can only allocate MAX_ALLOCS - 2 times. +// Can only allocate MAX_ALLOCS - 1 times. TEST_F(OffsetAllocatorTest, MaxNumAllocations) { constexpr uint32 ALLOCATOR_SIZE = 1024 * 1024 * 1024; constexpr uint32 MAX_ALLOCS = 1000; @@ -416,7 +417,7 @@ TEST_F(OffsetAllocatorTest, MaxNumAllocations) { std::make_shared(0, ALLOCATOR_SIZE, MAX_ALLOCS); std::vector handles; - for (uint32 i = 0; i < MAX_ALLOCS - 2; ++i) { + for (uint32 i = 0; i < MAX_ALLOCS - 1; ++i) { auto handle = allocator->allocate(1024); ASSERT_TRUE(handle.has_value()) << "Failed to allocate size: " << 1024 << " at iteration: " << i; @@ -731,7 +732,7 @@ TEST_F(OffsetAllocatorTest, MaxAllocationCountEdgeCase) { std::vector handles; // Allocate up to the limit - for (uint32 i = 0; i < MAX_ALLOCS - 2; ++i) { + for (uint32 i = 0; i < MAX_ALLOCS - 1; ++i) { auto handle = allocator->allocate(1024); ASSERT_TRUE(handle.has_value()) << "Failed at iteration " << i; handles.push_back(std::move(*handle));