forked from mooncake-track/Mooncake
[Store] Optimize Offset Allocator (#706)
* Allow offset allocator to dynamically increase the capacity * Change multiplier to multiplier_bits for fast bit-wise operation
This commit is contained in:
ykwd
GitHub
parent
fc1eb07da9
commit
975eef12bf
|
|
@ -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<uint64_t>(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<OffsetAllocator> {
|
||||
public:
|
||||
// Factory method to create shared_ptr<OffsetAllocator>
|
||||
static std::shared_ptr<OffsetAllocator> create(uint64_t base, size_t size,
|
||||
uint32 maxAllocs = 128 *
|
||||
1024);
|
||||
static std::shared_ptr<OffsetAllocator> 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<OffsetAllocator> {
|
|||
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<OffsetAllocator> {
|
|||
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;
|
||||
|
|
|
|||
|
|
@ -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<uint64_t>(1024));
|
||||
init_capacity =
|
||||
std::min(init_capacity, static_cast<uint64_t>(64 * 1024));
|
||||
// 1M <= max_capacity <= 64G / 1K = 64M
|
||||
uint64_t max_capacity = size / 1024;
|
||||
max_capacity =
|
||||
std::max(max_capacity, static_cast<uint64_t>(1024 * 1024));
|
||||
max_capacity =
|
||||
std::min(max_capacity, static_cast<uint64_t>(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<uint32_t>(init_capacity),
|
||||
static_cast<uint32_t>(max_capacity));
|
||||
if (!offset_allocator_) {
|
||||
LOG(ERROR) << "status=failed_to_create_offset_allocator";
|
||||
throw std::runtime_error("Failed to create offset allocator");
|
||||
|
|
|
|||
|
|
@ -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> 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<OffsetAllocator>(
|
||||
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<OffsetAllocationHandle> OffsetAllocator::allocate(size_t size) {
|
||||
|
|
@ -558,14 +571,17 @@ std::optional<OffsetAllocationHandle> 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<uint64_t>(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<OffsetAllocationHandle> 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_) << "}";
|
||||
|
|
|
|||
|
|
@ -104,7 +104,8 @@ class AllocatorWrapper : public std::enable_shared_from_this<AllocatorWrapper> {
|
|||
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<AllocatorWrapper>(0, ALLOCATOR_SIZE, MAX_ALLOCS);
|
||||
|
||||
std::vector<AllocationHandleWrapper> 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<AllocationHandleWrapper> 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));
|
||||
|
|
|
|||
Loading…
Reference in New Issue