Skip to content

Commit

Permalink
Merge pull request #527 from Coldwings/refact_pooledalloc
Browse files Browse the repository at this point in the history
Make PooledStackAllocator implemention more clear, and fix
  • Loading branch information
lihuiba authored Aug 1, 2024
2 parents 8c41f96 + 0208f49 commit 1fd201d
Showing 1 changed file with 23 additions and 22 deletions.
45 changes: 23 additions & 22 deletions thread/stack-allocator.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,11 @@ class PooledStackAllocator {
__builtin_ffsl(MAX_ALLOCATION_SIZE / MIN_ALLOCATION_SIZE);

public:
PooledStackAllocator() {}
PooledStackAllocator() {
for (size_t i = 0; i < N_SLOTS; i++) {
slots[i].slotsize = MIN_ALLOCATION_SIZE << i;
}
}

protected:
size_t in_pool_size = 0;
Expand All @@ -62,24 +66,26 @@ class PooledStackAllocator {
}

struct Slot {
std::vector<std::pair<void*, size_t>> pool;
std::vector<void*> pool;
uint32_t slotsize;

~Slot() {
for (auto pt : pool) {
__dealloc(pt.first, pt.second);
__dealloc(pt, slotsize);
}
}
std::pair<void*, size_t> get() {
void* get() {
if (!pool.empty()) {
auto ret = pool.back();
pool.pop_back();
return ret;
}
return {nullptr, 0};
return __alloc(slotsize);
}
void put(void* ptr, size_t size) { pool.emplace_back(ptr, size); }
void put(void* ptr) { pool.emplace_back(ptr); }
};

// get_slot(length) returns first slot that larger or equal to length
static inline uint32_t get_slot(uint32_t length) {
static auto base = __builtin_clz(MIN_ALLOCATION_SIZE - 1);
auto index = __builtin_clz(length - 1);
Expand All @@ -91,31 +97,26 @@ class PooledStackAllocator {
public:
void* alloc(size_t size) {
auto idx = get_slot(size);
if (unlikely(idx > N_SLOTS)) {
if (unlikely(idx >= N_SLOTS)) {
// larger than biggest slot
return __alloc(size);
}
auto ptr = slots[idx].get();
if (unlikely(!ptr.first)) {
// slots[idx] empty
auto aligned_size = MIN_ALLOCATION_SIZE << idx;
return __alloc(aligned_size);
}
// got from pool
in_pool_size -= ptr.second;
return ptr.first;
in_pool_size -= slots[idx].slotsize;
return ptr;
}
int dealloc(void* ptr, size_t size) {
auto idx = get_slot(size);
if (unlikely(idx > N_SLOTS ||
(in_pool_size + size >= trim_threshold))) {
if (unlikely(idx >= N_SLOTS ||
(in_pool_size + slots[idx].slotsize >= trim_threshold))) {
// big block or in-pool buffers reaches to threshold
__dealloc(ptr, size);
__dealloc(ptr, idx >= N_SLOTS ? size : slots[idx].slotsize);
return 0;
}
// Collect into pool
in_pool_size += size;
slots[idx].put(ptr, size);
in_pool_size += slots[idx].slotsize;
slots[idx].put(ptr);
return 0;
}
size_t trim(size_t keep_size) {
Expand All @@ -124,9 +125,9 @@ class PooledStackAllocator {
if (!slots[i].pool.empty()) {
auto ptr = slots[i].pool.back();
slots[i].pool.pop_back();
in_pool_size -= ptr.second;
count += ptr.second;
__dealloc(ptr.first, ptr.second);
in_pool_size -= slots[i].slotsize;
count += slots[i].slotsize;
__dealloc(ptr, slots[i].slotsize);
}
}
return count;
Expand Down

0 comments on commit 1fd201d

Please sign in to comment.