feat(alloc): support wait-free allocator

[ytakano]

This PR introduces wf_alloc as a feature-selectable heap allocator backend for awkernel_lib, while keeping TLSF as the fallback path on architectures where wf_alloc is not enabled/supported. It also updates heap allocator documentation to match current behavior.

Description

• Make heap backend selection explicit by compile-time feature:
  • heap-wf-alloc selects wf_alloc.
  • otherwise, TLSF is used.
• Keep allocator behavior under Talloc policy (primary/backup) unchanged:
  • user-space path still uses primary-first behavior.
  • kernel path uses primary-then-backup behavior.
  • deallocation routing remains range-based.
• Add wf_alloc plumbing for supported targets (x86_64, aarch64):
  • add optional dependency and feature gating in awkernel_lib.
  • add compile-time guard that rejects heap-wf-alloc on unsupported architectures.
  • implement backend switch via type Allocator = ... alias.
• Update architecture heap initialization to pass active CPU count into allocator setup:
  • x86_64 and aarch64 now initialize primary/backup heaps via init_*_with_num_cpu(...).
• Default feature wiring for kernel is aligned with this backend selection (heap-wf-alloc on x86 by default, TLSF remains for non-CAS2 paths).
• Update awkernel_lib/src/heap.rs doc comments:
  • remove stale claim that TLSF is default.
  • document wf_alloc initialization flow, CPU-token mapping, interrupt-guarded alloc/dealloc, and fallback to existing OOM path when allocator init fails.
• Initialize AcpiTable before heap memory initialization for x86_64.

Related links

Benchmarking:

• https://ytakano.github.io/wf_alloc/memalloc_report_20260613_135953.html
• wf_alloc is faster than mimalloc and TLSF
• wf_alloc's alloc/free worst case latency is better mimalloc and TLSF

How was this PR tested?

Tested on Qemu and physical machines

Notes for reviewers

[Copilot]

Pull request overview

This PR adds a feature-selectable heap allocator backend to awkernel_lib, enabling wf_alloc on supported architectures while preserving the existing TLSF-based path as the fallback.

Changes:

• Introduces heap-wf-alloc / heap-tlsf feature wiring to explicitly select the heap backend (kernel defaults updated per-arch).
• Plumbs CPU-count-aware heap initialization (init_*_with_num_cpu) into x86_64 and aarch64 kernel bring-up.
• Refactors awkernel_lib::heap to select backend via a type Allocator = ... alias, adds wf_alloc initialization logic, and updates heap module documentation.

Reviewed changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
kernel/src/arch/x86_64/kernel_main.rs	Detects CPU count earlier and passes num_cpu into backup/primary heap initialization.
kernel/src/arch/aarch64/kernel_main.rs	Passes the detected CPU count into primary/backup heap initialization.
kernel/Cargo.toml	Adds kernel-level feature flags to select allocator backend and sets per-target defaults.
awkernel_lib/src/heap.rs	Implements backend selection, adds wf_alloc backend, adds *_with_num_cpu init APIs, and updates docs.
awkernel_lib/Cargo.toml	Adds optional wf_alloc dependency (target-gated) and exposes heap-wf-alloc / heap-tlsf features.

Comments suppressed due to low confidence (1)

awkernel_lib/src/heap.rs:42

• The documentation states that only 32/64 CPUs are supported, but the implementation uses NUM_MAX_CPU (currently 512) for the CPU flag bitmap and for wf_alloc's active_threads upper bound. This limitation note looks stale and can mislead callers configuring num_cpu.

//! # Limitation
//!
//! Only 32 or 64 CPUs are supported for 32 or 64 bits CPU architectures.
//!

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[atsushi421]

The mdbook chapter at mdbook/src/internal/memory_allocator.md still describes the old TLSF-only design (Allocator/BackUpAllocator structs, init_primary/init_backup only). It should be updated to cover the new HeapBackend trait, backend selection by feature, and the *_with_num_cpu APIs (and explain why an explicit num_cpu is needed for wf_alloc).

---

The new ~200 lines of layout math in WfAllocBackend::init (alignment, overflow guards, metadata sizing, active_threads validation) have no unit tests. align_up alone is a pure helper that is trivially testable. Consider adding #[cfg(test)] mod tests that at least covers align_up overflow/normal cases and the rejection paths of init (too small heap, active_threads == 0 or > NUM_MAX_CPU, misaligned heap_start).

---

is_primary_mem (around line 384 in awkernel_lib/src/heap.rs) reads primary_start and primary_size with Relaxed, and the matching stores in init_*_with_num_cpu are also Relaxed. Another CPU can observe "new primary_start, old primary_size = 0" between the two stores; in that window end == start, the range check is always false, and primary-owned deallocations are routed to the backup, which would treat them as foreign and corrupt its free list. Use Release for the size store (last) and Acquire for the load, or document that init must happen-before any allocation.

[ytakano]

The new ~200 lines of layout math in WfAllocBackend::init (...) have no unit tests. align_up alone is a pure helper that is trivially testable. Consider adding #[cfg(test)] mod tests that at least covers align_up overflow/normal cases and the rejection paths of init ...

Added in da52eab. While doing this I hit a structural snag worth noting: the heap module is #[cfg(not(feature = "std"))] (it installs the #[global_allocator]), and the host test build runs with --features std, so a #[cfg(test)] mod tests placed inside heap / wf_alloc_backend is never compiled by make test.

So I extracted align_up into a new always-declared heap_util module (its body gated to the wf backend or cfg(test)), which the std test harness does compile and run — mirroring how local_heap's tests run on the host. The tests cover align_up's already-aligned, round-up, alignment-of-one, large power-of-two, and overflow (None) cases (make test: 4 new tests, all passing).

I did not unit-test the init rejection paths: they depend on wf_alloc types (WfAlloc::metadata_region_align/size, from_metadata_region), so testing them on the host would require enabling heap-wf-alloc for the whole test binary, which also swaps in an uninitialized WfAllocBackend as the global allocator and breaks every other test. If you'd like that coverage, I think the right move is a dedicated wf-enabled test target rather than changing the default one — happy to do that as a follow-up if you want.

[ytakano]

is_primary_mem (around line 384 in awkernel_lib/src/heap.rs) reads primary_start and primary_size with Relaxed (...) Another CPU can observe "new primary_start, old primary_size = 0" between the two stores (...) Use Release for the size store (last) and Acquire for the load, or document that init must happen-before any allocation.

Took the "document that init must happen-before any allocation" option, in 471c5cd — and it turns out that's the technically accurate fix here rather than just the convenient one.

The window you describe ("new primary_start, old primary_size == 0") only exists while the BSP runs heap init single-threaded; the APs never read primary_start / primary_size until after init has fully completed, because the boot sequence puts a happens-before edge between init and any AP allocation:

• x86_64: the BSP does all heap init while the APs are parked, then publishes BSP_READY with a release store; each AP observes it and runs an fence(Acquire) before its first allocation (use_primary_then_backup etc.).
• aarch64: the BSP publishes PRIMARY_INITIALIZED with a SeqCst store after heap init, and each AP spins on a SeqCst load of it before proceeding.

So no other CPU ever observes the intermediate state, and a primary-owned deallocation is never misrouted to the backup. I also realized that bumping the two stores/loads to Release/Acquire wouldn't actually close the hole on its own: primary_start and primary_size are two separate atomics and can't be read as one atomic pair regardless of ordering — it's the init-before-sharing contract, not the per-atomic ordering, that provides safety. I documented this on init_primary_with_num_cpu / init_backup_with_num_cpu and is_primary_mem, tying it to the single-init # Safety contract on HeapBackend::init.

Happy to additionally bump these to Release/Acquire as defense-in-depth if you'd prefer the local atomics to be self-describing, but it isn't required for correctness given the boot handshake.

[ytakano]

The mdbook chapter at mdbook/src/internal/memory_allocator.md still describes the old TLSF-only design (Allocator/BackUpAllocator structs, init_primary/init_backup only). It should be updated to cover the new HeapBackend trait, backend selection by feature, and the *_with_num_cpu APIs (and explain why an explicit num_cpu is needed for wf_alloc).

Updated in 94bc5e2. The chapter now covers:

• compile-time backend selection (heap-wf-alloc → wf_alloc on x86_64/aarch64, otherwise TLSF; mutually exclusive via compile_error!);
• the HeapBackend trait and the type Allocator = ... alias, plus the real Talloc { primary: Allocator, backup: Allocator, ... } (replacing the old Allocator / BackUpAllocator structs);
• both the bare init_primary / init_backup and the init_*_with_num_cpu APIs, with a paragraph on why wf_alloc needs the explicit CPU count (it maps each CPU id to a per-CPU token and sizes its metadata from active_threads, bailing out on 0), why the bare variants read cpu::num_cpu() and so the boot code uses the explicit-count variants on x86_64, and that TLSF ignores the count;
• the x86_64 / aarch64 boot snippets now passing num_cpu.

Also fixed the "For x86_64" copy-paste in the AArch64 section and the Tallock typo. mdbook build renders cleanly.

This was the last open item from the review — thanks for the thorough pass.

[Copilot]

Pull request overview

Copilot reviewed 11 out of 12 changed files in this pull request and generated 2 comments.