daniel/MontaukOS
1
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity

feat: overhaul userspace heap and virtual memory

Browse Source
This commit is contained in:
daniel
2026-08-10 12:25:27 +02:00
parent 4ecec32c5c
commit 9b23d99082
21 changed files with 905 additions and 451 deletions
Download Patch File Download Diff File Expand all files Collapse all files
programs/include/montauk/heap.h
+20 -271
View File
@@ -1,290 +1,39 @@
/*
* heap.h
* Userspace heap allocator for MontaukOS programs
* Copyright (c) 2025 Daniel Hammer
* Unified userspace heap API for MontaukOS programs
* Copyright (c) 2025-2026 Daniel Hammer
*/
#pragma once
#include <montauk/syscall.h>
#include <montauk/string.h>
#include <cstddef>
#include <cstdint>
// The allocator lives in libc. Keeping these declarations here lets
// freestanding C++ programs use the Montauk API without pulling in all of
// <stdlib.h>, while ensuring C, C++, and libraries share one heap.
extern "C" {
void* malloc(std::size_t size);
void free(void* ptr);
void* realloc(void* ptr, std::size_t size);
void* calloc(std::size_t count, std::size_t size);
}
namespace montauk {
namespace heap_detail {
static constexpr uint64_t HEADER_MAGIC = 0x5A484541; // "ZHEA"
static constexpr uint64_t FREED_MAGIC = 0xDEADFEEE;
struct Header {
uint64_t magic;
uint64_t size; // user-requested size
} __attribute__((packed));
struct FreeNode {
uint64_t size; // total size of this free block (including node)
FreeNode* next;
};
// Segregated free lists: power-of-2 size classes for blocks <= 4096 bytes.
// Blocks larger than 4096 go to the overflow list.
static constexpr int NUM_BUCKETS = 8;
static constexpr uint64_t BUCKET_SIZES[NUM_BUCKETS] = {
32, 64, 128, 256, 512, 1024, 2048, 4096
};
// Per-process heap state — must be `inline` (not `static`) so that all
// translation units in a multi-TU program share a single heap.
inline FreeNode* g_buckets[NUM_BUCKETS] = {};
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::abi::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));
}
// Determine which bucket a block size belongs to, or -1 for overflow
static inline int bucket_index(uint64_t blockSize) {
if (blockSize <= 32) return 0;
if (blockSize <= 64) return 1;
if (blockSize <= 128) return 2;
if (blockSize <= 256) return 3;
if (blockSize <= 512) return 4;
if (blockSize <= 1024) return 5;
if (blockSize <= 2048) return 6;
if (blockSize <= 4096) return 7;
return -1;
}
// Insert into overflow list (sorted by address, with adjacent-block coalescing)
static inline void insert_overflow(void* ptr, uint64_t size) {
auto* node = (FreeNode*)ptr;
node->size = size;
FreeNode* prev = &g_overflow;
FreeNode* cur = g_overflow.next;
while (cur != nullptr && cur < node) {
prev = cur;
cur = cur->next;
}
bool merged_prev = false;
if (prev != &g_overflow &&
(uint8_t*)prev + prev->size == (uint8_t*)node) {
prev->size += size;
node = prev;
merged_prev = true;
}
if (cur != nullptr &&
(uint8_t*)node + node->size == (uint8_t*)cur) {
node->size += cur->size;
node->next = cur->next;
if (!merged_prev) prev->next = node;
} else if (!merged_prev) {
node->next = cur;
prev->next = node;
}
}
// Take a block of at least `needed` bytes from the overflow list.
// Splits remainder back into overflow if worthwhile.
static inline void* take_from_overflow(uint64_t needed) {
FreeNode* prev = &g_overflow;
FreeNode* cur = g_overflow.next;
while (cur != nullptr) {
if (cur->size >= needed) {
uint64_t blockSize = cur->size;
prev->next = cur->next;
if (blockSize > needed + sizeof(FreeNode) + 16) {
insert_overflow((uint8_t*)cur + needed, blockSize - needed);
}
return (void*)cur;
}
prev = cur;
cur = cur->next;
}
return nullptr;
}
// Next slab size for heap growth. The kernel tracks a finite number
// of SYS_ALLOC records per process (MaxHeapAllocs), so growing once
// per large allocation exhausts them under allocation-heavy loads
// (the native ld ran out mid-link). Doubling slabs keep the syscall
// count logarithmic in total heap size.
inline uint64_t g_grow_slab = 16 * 0x1000;
static inline bool grow(uint64_t bytes) {
uint64_t want = (bytes + 0xFFF) & ~0xFFFULL;
if (want < 0x4000) want = 0x4000;
uint64_t slab = (want > g_grow_slab) ? want : g_grow_slab;
if (g_grow_slab < 4 * 1024 * 1024) g_grow_slab *= 2;
void* mem = montauk::alloc(slab);
if (mem == nullptr && slab > want) {
// Big slab refused (low memory): retry with the exact need.
slab = want;
mem = montauk::alloc(slab);
}
if (mem == nullptr) return false;
insert_overflow(mem, slab);
return true;
}
// Refill a small-block bucket by carving a page-sized chunk from overflow
static inline bool refill_bucket(int idx) {
uint64_t bsize = BUCKET_SIZES[idx];
uint64_t chunk = (bsize < 4096) ? 4096 : bsize;
void* block = take_from_overflow(chunk);
if (block == nullptr) {
if (!grow(chunk)) return false;
block = take_from_overflow(chunk);
if (block == nullptr) return false;
}
uint64_t count = chunk / bsize;
for (uint64_t i = 0; i < count; i++) {
auto* node = (FreeNode*)((uint8_t*)block + i * bsize);
node->size = bsize;
node->next = g_buckets[idx];
g_buckets[idx] = node;
}
return true;
}
} // namespace heap_detail
// ---- Public API ----
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;
}
uint64_t needed = size + sizeof(Header);
needed = (needed + 15) & ~15ULL;
int idx = bucket_index(needed);
if (idx >= 0) {
// Small allocation — use segregated bucket (O(1))
if (g_buckets[idx] == nullptr && !refill_bucket(idx)) {
heap_lock_release();
return nullptr;
}
FreeNode* node = g_buckets[idx];
g_buckets[idx] = node->next;
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)) { heap_lock_release(); return nullptr; }
block = take_from_overflow(needed);
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));
return ::malloc((std::size_t)size);
}
inline void mfree(void* ptr) {
using namespace heap_detail;
if (ptr == nullptr) return;
Header* header = get_header(ptr);
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);
blockSize = (blockSize + 15) & ~15ULL;
int idx = bucket_index(blockSize);
if (idx >= 0) {
// Small block — push onto bucket (O(1))
auto* node = (FreeNode*)header;
node->size = BUCKET_SIZES[idx];
node->next = g_buckets[idx];
g_buckets[idx] = node;
} else {
// Large block — sorted insert with coalescing
insert_overflow((void*)header, blockSize);
}
heap_lock_release();
::free(ptr);
}
inline void* realloc(void* ptr, uint64_t size) {
if (ptr == nullptr) return malloc(size);
return ::realloc(ptr, (std::size_t)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;
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];
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;
uint64_t copySize = (old < size) ? old : size;
memcpy(newBlock, ptr, copySize);
mfree(ptr);
return newBlock;
inline void* calloc(uint64_t count, uint64_t size) {
return ::calloc((std::size_t)count, (std::size_t)size);
}
} // namespace montauk
programs/include/montauk/thread.h
+47 -9
View File
@@ -44,8 +44,22 @@ namespace montauk {
ThreadEntry user_entry;
void* user_arg;
void* stack_base;
int tid;
ThreadCtx* next;
};
inline ThreadCtx* g_thread_records = nullptr;
inline volatile uint32_t g_thread_records_lock = 0;
inline void records_lock() {
while (__atomic_exchange_n(&g_thread_records_lock, 1, __ATOMIC_ACQUIRE) != 0)
montauk::yield();
}
inline void records_unlock() {
__atomic_store_n(&g_thread_records_lock, 0, __ATOMIC_RELEASE);
}
// Userspace trampoline: bridges from the raw entry the kernel jumps
// to into the typed entry, then funnels into SYS_THREAD_EXIT. We
// route the exit through libc rather than relying on a kernel-side
@@ -53,21 +67,24 @@ namespace montauk {
// memory on this path.
//
// The thread's stack itself is intentionally not freed here: we are
// still running on it. It is reclaimed when the process exits, or
// the joiner may free it explicitly after thread_join.
// still running on it. It is reclaimed by a successful thread_join,
// or as part of whole-process teardown if the thread is never joined.
[[noreturn]] inline void thread_trampoline(detail::ThreadCtx* ctx) {
// A sibling CPU can start the thread before thread_spawn has
// returned its TID. Wait until the parent has published the record
// needed by thread_join to reclaim this stack.
while (__atomic_load_n(&ctx->tid, __ATOMIC_ACQUIRE) == 0)
montauk::yield();
int code = ctx->user_entry(ctx->user_arg);
montauk::mfree(ctx);
thread_exit(code);
}
}
// Spawn a new thread that begins executing `entry(arg)`. Returns the
// new TID on success, or -1 on failure. The thread's stack is
// allocated from the user heap; it is leaked on thread exit (the
// thread itself cannot free the stack it is running on). The kernel
// reclaims it on process exit. Callers that need to spawn many short-
// lived threads should pool stacks themselves.
// allocated from the user heap. The exiting thread cannot free the stack
// it is running on, so thread_join reclaims both it and the trampoline
// context after the kernel has reaped the sibling.
inline int thread_spawn(ThreadEntry entry, void* arg,
uint64_t stack_bytes = 0) {
if (entry == nullptr) return -1;
@@ -84,6 +101,8 @@ namespace montauk {
ctx->user_entry = entry;
ctx->user_arg = arg;
ctx->stack_base = stack;
ctx->tid = 0;
ctx->next = nullptr;
uint64_t stack_top = ((uint64_t)stack + stack_bytes) & ~0xFULL;
int tid = (int)syscall3(montauk::abi::SYS_THREAD_SPAWN,
@@ -94,6 +113,11 @@ namespace montauk {
montauk::mfree(stack);
return -1;
}
detail::records_lock();
ctx->next = detail::g_thread_records;
detail::g_thread_records = ctx;
__atomic_store_n(&ctx->tid, tid, __ATOMIC_RELEASE);
detail::records_unlock();
return tid;
}
@@ -101,8 +125,22 @@ namespace montauk {
// success (with the thread's exit code in *out_code if non-null) or
// -1 if `tid` is not a joinable sibling.
inline int thread_join(int tid, int* out_code = nullptr) {
return (int)syscall2(montauk::abi::SYS_THREAD_JOIN,
(uint64_t)tid, (uint64_t)out_code);
int result = (int)syscall2(montauk::abi::SYS_THREAD_JOIN,
(uint64_t)tid, (uint64_t)out_code);
if (result == 0) {
detail::records_lock();
detail::ThreadCtx** link = &detail::g_thread_records;
while (*link != nullptr && (*link)->tid != tid)
link = &(*link)->next;
detail::ThreadCtx* ctx = *link;
if (ctx != nullptr) *link = ctx->next;
detail::records_unlock();
if (ctx != nullptr) {
montauk::mfree(ctx->stack_base);
montauk::mfree(ctx);
}
}
return result;
}
// Return the calling thread's TID (== getpid() for the main thread).