fix: process-wide spinlock in userspace heap
This commit is contained in:
@@ -37,6 +37,22 @@ namespace heap_detail {
|
|||||||
inline FreeNode g_overflow{0, nullptr};
|
inline FreeNode g_overflow{0, nullptr};
|
||||||
inline bool g_initialized = false;
|
inline bool g_initialized = false;
|
||||||
|
|
||||||
|
// Process-wide heap lock. Userspace threads share the heap, so the
|
||||||
|
// public malloc/mfree/realloc entry points must serialize access to
|
||||||
|
// g_buckets/g_overflow. Kept inline here (not in thread.h) because
|
||||||
|
// thread.h depends on heap.h, and the internal helpers below are not
|
||||||
|
// reentrant into the public API, so a plain spinlock suffices.
|
||||||
|
inline volatile uint32_t g_heap_lock = 0;
|
||||||
|
|
||||||
|
static inline void heap_lock_acquire() {
|
||||||
|
while (__atomic_exchange_n(&g_heap_lock, 1, __ATOMIC_ACQUIRE) != 0) {
|
||||||
|
syscall0(Montauk::SYS_YIELD);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
static inline void heap_lock_release() {
|
||||||
|
__atomic_store_n(&g_heap_lock, 0, __ATOMIC_RELEASE);
|
||||||
|
}
|
||||||
|
|
||||||
static inline Header* get_header(void* block) {
|
static inline Header* get_header(void* block) {
|
||||||
return (Header*)((uint8_t*)block - sizeof(Header));
|
return (Header*)((uint8_t*)block - sizeof(Header));
|
||||||
}
|
}
|
||||||
@@ -146,15 +162,17 @@ namespace heap_detail {
|
|||||||
inline void* malloc(uint64_t size) {
|
inline void* malloc(uint64_t size) {
|
||||||
using namespace heap_detail;
|
using namespace heap_detail;
|
||||||
|
|
||||||
|
// Guard against overflow: size + Header must not wrap
|
||||||
|
if (size > UINT64_MAX - sizeof(Header) - 15)
|
||||||
|
return nullptr;
|
||||||
|
|
||||||
|
heap_lock_acquire();
|
||||||
|
|
||||||
if (!g_initialized) {
|
if (!g_initialized) {
|
||||||
grow(16 * 0x1000); // seed with 64 KiB
|
grow(16 * 0x1000); // seed with 64 KiB
|
||||||
g_initialized = true;
|
g_initialized = true;
|
||||||
}
|
}
|
||||||
|
|
||||||
// Guard against overflow: size + Header must not wrap
|
|
||||||
if (size > UINT64_MAX - sizeof(Header) - 15)
|
|
||||||
return nullptr;
|
|
||||||
|
|
||||||
uint64_t needed = size + sizeof(Header);
|
uint64_t needed = size + sizeof(Header);
|
||||||
needed = (needed + 15) & ~15ULL;
|
needed = (needed + 15) & ~15ULL;
|
||||||
|
|
||||||
@@ -162,8 +180,10 @@ namespace heap_detail {
|
|||||||
|
|
||||||
if (idx >= 0) {
|
if (idx >= 0) {
|
||||||
// Small allocation — use segregated bucket (O(1))
|
// Small allocation — use segregated bucket (O(1))
|
||||||
if (g_buckets[idx] == nullptr && !refill_bucket(idx))
|
if (g_buckets[idx] == nullptr && !refill_bucket(idx)) {
|
||||||
|
heap_lock_release();
|
||||||
return nullptr;
|
return nullptr;
|
||||||
|
}
|
||||||
|
|
||||||
FreeNode* node = g_buckets[idx];
|
FreeNode* node = g_buckets[idx];
|
||||||
g_buckets[idx] = node->next;
|
g_buckets[idx] = node->next;
|
||||||
@@ -171,20 +191,22 @@ namespace heap_detail {
|
|||||||
Header* header = (Header*)node;
|
Header* header = (Header*)node;
|
||||||
header->magic = HEADER_MAGIC;
|
header->magic = HEADER_MAGIC;
|
||||||
header->size = size;
|
header->size = size;
|
||||||
|
heap_lock_release();
|
||||||
return (void*)((uint8_t*)header + sizeof(Header));
|
return (void*)((uint8_t*)header + sizeof(Header));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Large allocation — search overflow list
|
// Large allocation — search overflow list
|
||||||
void* block = take_from_overflow(needed);
|
void* block = take_from_overflow(needed);
|
||||||
if (block == nullptr) {
|
if (block == nullptr) {
|
||||||
if (!grow(needed)) return nullptr;
|
if (!grow(needed)) { heap_lock_release(); return nullptr; }
|
||||||
block = take_from_overflow(needed);
|
block = take_from_overflow(needed);
|
||||||
if (block == nullptr) return nullptr;
|
if (block == nullptr) { heap_lock_release(); return nullptr; }
|
||||||
}
|
}
|
||||||
|
|
||||||
Header* header = (Header*)block;
|
Header* header = (Header*)block;
|
||||||
header->magic = HEADER_MAGIC;
|
header->magic = HEADER_MAGIC;
|
||||||
header->size = size;
|
header->size = size;
|
||||||
|
heap_lock_release();
|
||||||
return (void*)((uint8_t*)header + sizeof(Header));
|
return (void*)((uint8_t*)header + sizeof(Header));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -195,8 +217,10 @@ namespace heap_detail {
|
|||||||
|
|
||||||
Header* header = get_header(ptr);
|
Header* header = get_header(ptr);
|
||||||
|
|
||||||
if (header->magic == FREED_MAGIC) return; // double-free
|
heap_lock_acquire();
|
||||||
if (header->magic != HEADER_MAGIC) return; // corrupt
|
|
||||||
|
if (header->magic == FREED_MAGIC) { heap_lock_release(); return; } // double-free
|
||||||
|
if (header->magic != HEADER_MAGIC) { heap_lock_release(); return; } // corrupt
|
||||||
header->magic = FREED_MAGIC;
|
header->magic = FREED_MAGIC;
|
||||||
|
|
||||||
uint64_t blockSize = header->size + sizeof(Header);
|
uint64_t blockSize = header->size + sizeof(Header);
|
||||||
@@ -214,15 +238,18 @@ namespace heap_detail {
|
|||||||
// Large block — sorted insert with coalescing
|
// Large block — sorted insert with coalescing
|
||||||
insert_overflow((void*)header, blockSize);
|
insert_overflow((void*)header, blockSize);
|
||||||
}
|
}
|
||||||
|
heap_lock_release();
|
||||||
}
|
}
|
||||||
|
|
||||||
inline void* realloc(void* ptr, uint64_t size) {
|
inline void* realloc(void* ptr, uint64_t size) {
|
||||||
if (ptr == nullptr) return malloc(size);
|
if (ptr == nullptr) return malloc(size);
|
||||||
|
|
||||||
|
// Read old size under the lock to avoid racing with another
|
||||||
|
// thread that might be freeing/recycling this header.
|
||||||
|
heap_detail::heap_lock_acquire();
|
||||||
auto* header = heap_detail::get_header(ptr);
|
auto* header = heap_detail::get_header(ptr);
|
||||||
uint64_t old = header->size;
|
uint64_t old = header->size;
|
||||||
|
|
||||||
// Compute actual block size (accounting for bucket rounding)
|
|
||||||
uint64_t oldBlock = (old + sizeof(heap_detail::Header) + 15) & ~15ULL;
|
uint64_t oldBlock = (old + sizeof(heap_detail::Header) + 15) & ~15ULL;
|
||||||
int idx = heap_detail::bucket_index(oldBlock);
|
int idx = heap_detail::bucket_index(oldBlock);
|
||||||
if (idx >= 0) oldBlock = heap_detail::BUCKET_SIZES[idx];
|
if (idx >= 0) oldBlock = heap_detail::BUCKET_SIZES[idx];
|
||||||
@@ -230,8 +257,10 @@ namespace heap_detail {
|
|||||||
uint64_t newNeed = (size + sizeof(heap_detail::Header) + 15) & ~15ULL;
|
uint64_t newNeed = (size + sizeof(heap_detail::Header) + 15) & ~15ULL;
|
||||||
if (newNeed <= oldBlock) {
|
if (newNeed <= oldBlock) {
|
||||||
header->size = size;
|
header->size = size;
|
||||||
|
heap_detail::heap_lock_release();
|
||||||
return ptr;
|
return ptr;
|
||||||
}
|
}
|
||||||
|
heap_detail::heap_lock_release();
|
||||||
|
|
||||||
void* newBlock = malloc(size);
|
void* newBlock = malloc(size);
|
||||||
if (newBlock == nullptr) return nullptr;
|
if (newBlock == nullptr) return nullptr;
|
||||||
|
|||||||
Reference in New Issue
Block a user