atmos/kernel/

allocator.rs

//! # TLSF (Two-Level Segregated Fit) ヒープアロケータ
//!
//! `allocator.rs` - 512MBのカーネルヒープ領域を管理するアロケータです。
//!
//! ## 採用理由と仕様
//! 以前の first-fit フリーリスト方式では、ブラウザのパース処理などでフリーリストが数万ノードに肥大化した際、
//! 確保・解放が $O(N)$ になり数秒間のフリーズを引き起こしていました。
//! 本実装に用いられている TLSF 方式（`rlsf` クレートを利用）は、確保・解放ともに **最悪計算量 $O(1)$** が保証されており、
//! メモリ断片化も有界に抑えられるため、ベアメタル環境において非常に高い安定性とパフォーマンスを発揮します。
//!
//! ## パニック回復とスレッド安全性
//! パニック時のデッドロックを防止するための `Locked` ラッパー（コア別のロック保持追跡と強制解放機能）と、
//! 割り込みセーフな `spin::Mutex` を統合しています。これにより、ロックを保持したままコアが例外で落ちた場合でも、
//! 例外ハンドラのログ出力時にアロケータロックの再デッドロックを防止できます。

#![allow(dead_code)]

use core::alloc::{GlobalAlloc, Layout};
use core::ptr::{null_mut, NonNull};
use core::sync::atomic::{AtomicBool, Ordering};
use rlsf::Tlsf;

// 512MB ヒープに適合する TLSF 構成:
//   GRANULARITY = size_of::<usize>()*4 = 32B、最大ブロック = (32 << FLLEN) - 32。
//   FLLEN=25 → 最大ブロック ≈ 1GB、MAX_POOL_SIZE = 1<<(5+25) = 1GB ≥ 512MB。
//   FLBitmap=u32（≥25bit）、SLLEN=16（2 のべき・≤16bit）→ SLBitmap=u16。
type TheTlsf = Tlsf<'static, u32, u16, 25, 16>;

/// TLSFアロケータの実体を保持し、確保・解放および使用量計測を行う内部構造体。
struct TlsfHeap {
    tlsf: TheTlsf,
    total: usize,
    used: usize,
}

impl TlsfHeap {
    /// 空のアロケータ実体を初期生成します。
    const fn empty() -> Self {
        Self {
            tlsf: Tlsf::new(),
            total: 0,
            used: 0,
        }
    }

    /// メモリプールを登録してヒープ領域を初期化します。
    ///
    /// # Safety
    /// 与えられた `start` アドレスから `size` バイトのメモリ領域が、排他的に利用可能な物理アドレス領域である必要があります。
    unsafe fn init(&mut self, start: usize, size: usize) {
        self.total = size;
        self.used = 0;
        let slice = core::ptr::slice_from_raw_parts_mut(start as *mut u8, size);
        if let Some(nn) = NonNull::new(slice) {
            self.tlsf.insert_free_block_ptr(nn);
        }
    }

    unsafe fn alloc(&mut self, layout: Layout) -> *mut u8 {
        match self.tlsf.allocate(layout) {
            Some(p) => {
                self.used = self.used.saturating_add(layout.size());
                p.as_ptr()
            }
            None => null_mut(),
        }
    }

    unsafe fn dealloc(&mut self, ptr: NonNull<u8>, layout: Layout) {
        // rlsf::deallocate は確保時と同じ align を要求する（layout.align() で一致）。
        self.tlsf.deallocate(ptr, layout.align());
        self.used = self.used.saturating_sub(layout.size());
    }

    fn free_space(&self) -> usize {
        self.total.saturating_sub(self.used)
    }
}

/// カーネル全体のアロケータのグローバルなインスタンス。
#[global_allocator]
static ALLOCATOR: Locked<TlsfHeap> = Locked::new(TlsfHeap::empty());

/// カーネルヒープの大きさ。
///
/// 【2026-08-25】この値は **2 箇所で使われる**。
///
/// 1. `main.rs` … ヒープをこの大きさで初期化する
/// 2. `fs/mod.rs` … ヒープの**直後**から RAMDisk を確保する
///
/// 両方が別々に `512 * 1024 * 1024` と書いていた。
/// 片方だけ変えると RAMDisk がヒープに**重なって**、
/// 記憶装置の書き込みがヒープを破壊する（発見が極めて難しい）。
/// 定数を 1 つにして、ずれようがないようにする。
pub const KERNEL_HEAP_SIZE: usize = 512 * 1024 * 1024;

/// カーネルヒープ領域全体の総バイト数。
static mut TOTAL_HEAP_SIZE: usize = 0;

/// 指定されたメモリ開始アドレスとサイズで、カーネルアロケータを初期化します。
pub fn init(start: usize, size: usize) {
    unsafe {
        TOTAL_HEAP_SIZE = size;
        ALLOCATOR.lock().init(start, size);
    }
}

/// ヒープの (使用中の概算バイト数, 総バイト数) を取得します。
///
/// ※ 使用中サイズはアロケータのヘッダやパディング分を考慮しない概算値です。
pub fn get_heap_stats() -> (usize, usize) {
    let total = unsafe { TOTAL_HEAP_SIZE };
    let free = ALLOCATOR.lock().free_space();
    (total.saturating_sub(free), total)
}

/// 旧 first-fit の canary 整合性検査との互換スタブです。
///
/// TLSF 方式では境界タグ方式によりアロケータが独自に健全性を管理しているため、
/// この検査は常に `false`（破損未検出）を返します。
pub fn check_heap_integrity() -> bool {
    false
}

/// パニック回復処理において、指定された CPU コア ID が現在ヒープロックを保持している場合、
/// そのロックを強制解放します。
pub fn force_release_lock(cid: i16) {
    unsafe {
        ALLOCATOR.force_release_if_core(cid);
    }
}

unsafe impl GlobalAlloc for Locked<TlsfHeap> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.lock().alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        if let Some(nn) = NonNull::new(ptr) {
            self.lock().dealloc(nn, layout);
        }
    }
}

// ---------------------------------------------------------------------------
// パニック回復対応ロックラッパー
// ---------------------------------------------------------------------------

/// 割り込みセーフな `spin::Mutex` をラップし、パニック発生時のデッドロックを防止するために
/// ロックを現在保持している CPU コア ID の追跡記録機能を追加したラッパー構造体。
pub struct Locked<A> {
    inner: spin::Mutex<A>,
    /// ロックを現在保持しているコア ID（-1 の場合は未保持）。
    holder: core::sync::atomic::AtomicI16,
}

impl<A> Locked<A> {
    /// 与えられたアロケータをラップする `Locked` インスタンスを新規生成します。
    pub const fn new(inner: A) -> Self {
        Self {
            inner: spin::Mutex::new(inner),
            holder: core::sync::atomic::AtomicI16::new(-1),
        }
    }

    /// ロックを獲得し、保持するコア ID を自動記録してガードオブジェクトを返します。
    ///
    /// # プリエンプション禁止は不要（2026-08-05 検討済み）
    ///
    /// 一度 `preempt_disable()` を足したが、**誤りだったので取り消した**。
    /// 下の `spin::Mutex::lock` は取得時に**割り込みを禁止**する
    /// （`is_interrupt_enabled` / `disable_interrupts`）。
    /// 割り込みが来ないのだからタイマ割り込み経由のプリエンプションも起きず、
    /// 「ロックを握ったまま切り替わって、他スレッドが空回りする」経路は存在しない。
    ///
    /// 効果が無いばかりか、`preempt_disable`/`enable` は SeqCst のアトミック
    /// 読み書きなので、**全確保・全解放**に余計なコストを足すだけだった。
    /// ここは最も呼ばれる経路なので、意味の無い命令を置いてはいけない。
    pub fn lock(&self) -> LockedGuard<'_, A> {
        let guard = self.inner.lock();
        self.holder.store(
            crate::kernel::scheduler::core_id() as i16,
            Ordering::Relaxed,
        );
        LockedGuard {
            guard: Some(guard),
            holder: &self.holder,
        }
    }

    /// パニック回復時に、指定された CPU コア `cid` が現在このロックを保持している場合、
    /// そのロックを強制解放してリカバリ可能にします。
    ///
    /// # Safety
    /// コア `cid` の現在実行中のタスク（スレッド）がアボートされ、その文脈が破棄される状況でのみ
    /// 安全に呼び出すことができます。
    pub unsafe fn force_release_if_core(&self, cid: i16) {
        if self.holder.load(Ordering::Relaxed) == cid {
            self.holder.store(-1, Ordering::Relaxed);
            self.inner.force_unlock();
        }
    }
}

pub struct LockedGuard<'a, A> {
    guard: Option<spin::MutexGuard<'a, A>>,
    holder: &'a core::sync::atomic::AtomicI16,
}

impl<'a, A> Drop for LockedGuard<'a, A> {
    fn drop(&mut self) {
        self.holder.store(-1, Ordering::Relaxed);
        self.guard.take(); // spin ロック解放
    }
}

impl<'a, A> core::ops::Deref for LockedGuard<'a, A> {
    type Target = A;
    fn deref(&self) -> &A {
        match self.guard.as_ref() {
            Some(g) => g,
            None => loop {
                core::hint::spin_loop();
            },
        }
    }
}

impl<'a, A> core::ops::DerefMut for LockedGuard<'a, A> {
    fn deref_mut(&mut self) -> &mut A {
        match self.guard.as_mut() {
            Some(g) => g,
            None => loop {
                core::hint::spin_loop();
            },
        }
    }
}

// ---------------------------------------------------------------------------
// IRQ セーフな自前スピンロック（旧実装から温存）
// ---------------------------------------------------------------------------

mod spin {
    use super::*;

    pub struct Mutex<T> {
        lock: AtomicBool,
        data: core::cell::UnsafeCell<T>,
    }

    unsafe impl<T: Send> Sync for Mutex<T> {}
    unsafe impl<T: Send> Send for Mutex<T> {}

    pub struct MutexGuard<'a, T> {
        mutex: &'a Mutex<T>,
        interrupt_state: bool,
    }

    impl<T> Mutex<T> {
        pub const fn new(value: T) -> Self {
            Self {
                lock: AtomicBool::new(false),
                data: core::cell::UnsafeCell::new(value),
            }
        }

        pub fn lock(&self) -> MutexGuard<'_, T> {
            let irq_enabled = is_interrupt_enabled();
            if irq_enabled {
                disable_interrupts();
            }

            while self
                .lock
                .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
                .is_err()
            {
                core::hint::spin_loop();
            }

            MutexGuard {
                mutex: self,
                interrupt_state: irq_enabled,
            }
        }

        /// パニック回復専用: ガードを経由せずロックを強制解放する。
        /// 割込み状態は復元しない（呼び出し直後にコンテキスト切替で SPSR から復元される前提）。
        ///
        /// # Safety
        /// 当該クリティカルセクションが破棄される（パニックしたプロセスが exit する）場合のみ。
        pub unsafe fn force_unlock(&self) {
            self.lock.store(false, Ordering::Release);
        }
    }

    impl<'a, T> core::ops::Deref for MutexGuard<'a, T> {
        type Target = T;
        fn deref(&self) -> &Self::Target {
            unsafe { &*self.mutex.data.get() }
        }
    }

    impl<'a, T> core::ops::DerefMut for MutexGuard<'a, T> {
        fn deref_mut(&mut self) -> &mut Self::Target {
            unsafe { &mut *self.mutex.data.get() }
        }
    }

    impl<'a, T> Drop for MutexGuard<'a, T> {
        fn drop(&mut self) {
            self.mutex.lock.store(false, Ordering::Release);
            if self.interrupt_state {
                enable_interrupts();
            }
        }
    }

    pub(super) fn is_interrupt_enabled() -> bool {
        let daif: u64;
        unsafe {
            core::arch::asm!("mrs {}, daif", out(reg) daif);
        }
        (daif & (1 << 7)) == 0
    }

    pub(super) fn disable_interrupts() {
        unsafe {
            core::arch::asm!("msr daifset, #2", options(nomem, nostack));
        }
    }

    pub(super) fn enable_interrupts() {
        unsafe {
            core::arch::asm!("msr daifclr, #2", options(nomem, nostack));
        }
    }
}