diff --git a/programs/include/montauk/heap.h b/programs/include/montauk/heap.h index 283eeac..87ca143 100644 --- a/programs/include/montauk/heap.h +++ b/programs/include/montauk/heap.h @@ -37,6 +37,22 @@ namespace heap_detail { inline FreeNode g_overflow{0, nullptr}; 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) { return (Header*)((uint8_t*)block - sizeof(Header)); } @@ -146,15 +162,17 @@ namespace heap_detail { inline void* malloc(uint64_t size) { 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) { grow(16 * 0x1000); // seed with 64 KiB 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); needed = (needed + 15) & ~15ULL; @@ -162,8 +180,10 @@ namespace heap_detail { if (idx >= 0) { // 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; + } FreeNode* node = g_buckets[idx]; g_buckets[idx] = node->next; @@ -171,20 +191,22 @@ namespace heap_detail { Header* header = (Header*)node; header->magic = HEADER_MAGIC; header->size = size; + heap_lock_release(); return (void*)((uint8_t*)header + sizeof(Header)); } // Large allocation — search overflow list void* block = take_from_overflow(needed); if (block == nullptr) { - if (!grow(needed)) return nullptr; + if (!grow(needed)) { heap_lock_release(); return nullptr; } block = take_from_overflow(needed); - if (block == nullptr) return nullptr; + if (block == nullptr) { heap_lock_release(); return nullptr; } } Header* header = (Header*)block; header->magic = HEADER_MAGIC; header->size = size; + heap_lock_release(); return (void*)((uint8_t*)header + sizeof(Header)); } @@ -195,8 +217,10 @@ namespace heap_detail { Header* header = get_header(ptr); - if (header->magic == FREED_MAGIC) return; // double-free - if (header->magic != HEADER_MAGIC) return; // corrupt + heap_lock_acquire(); + + 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; uint64_t blockSize = header->size + sizeof(Header); @@ -214,15 +238,18 @@ namespace heap_detail { // Large block — sorted insert with coalescing insert_overflow((void*)header, blockSize); } + heap_lock_release(); } inline void* realloc(void* ptr, uint64_t 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); uint64_t old = header->size; - // Compute actual block size (accounting for bucket rounding) uint64_t oldBlock = (old + sizeof(heap_detail::Header) + 15) & ~15ULL; int idx = heap_detail::bucket_index(oldBlock); 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; if (newNeed <= oldBlock) { header->size = size; + heap_detail::heap_lock_release(); return ptr; } + heap_detail::heap_lock_release(); void* newBlock = malloc(size); if (newBlock == nullptr) return nullptr;