strat9_kernel/arch/x86_64/

smp.rs

//! SMP (Symmetric Multi-Processing) boot for x86_64.
//!
//! Boots Application Processors (APs) using the legacy INIT+SIPI sequence
//! and parks them in an idle loop. Per-CPU data is initialized but no
//! per-CPU scheduler is active yet.

use core::{
    arch::global_asm,
    sync::atomic::{AtomicBool, AtomicUsize, Ordering},
};

use alloc::{vec, vec::Vec};
use x86_64::{
    structures::paging::{Page, PageTableFlags, PhysFrame, Size4KiB},
    PhysAddr, VirtAddr,
};

use crate::{
    acpi::madt,
    arch::x86_64::{apic, idt, io::io_wait, percpu, timer},
    memory,
    process::task::KernelStack,
    sync::SpinLock,
};

/// Physical address where the SMP trampoline is copied.
pub const TRAMPOLINE_PHYS_ADDR: u64 = 0x8000;

/// Number of booted cores (starts at 1 for BSP).
static BOOTED_CORES: AtomicUsize = AtomicUsize::new(1);
/// Counter for synchronization barriers.
static SYNC_BARRIER: AtomicUsize = AtomicUsize::new(0);
/// Target count for the rendezvous barrier (set by BSP before barrier).
static BARRIER_TARGET: AtomicUsize = AtomicUsize::new(0);
/// Gate used by BSP to release APs into scheduler/timer start.
static AP_SCHED_GATE_OPEN: AtomicBool = AtomicBool::new(false);

/// Keep AP kernel stacks alive.
static AP_KERNEL_STACKS: SpinLock<Vec<KernelStack>> = SpinLock::new(Vec::new());

// Boot Application Processors using the legacy INIT+SIPI sequence.
// Returns the number of online CPUs (including BSP) or an error string on failure.
// - BSP must have already initialized the APIC and parsed MADT to get APIC IDs.
// - APs are parked in an idle loop after booting, waiting for the scheduler gate
// to open before starting the scheduler and timer on each AP.
#[cfg(target_arch = "x86_64")]
global_asm!(
    r#"
.section .text
.code16

.global smp_trampoline
.global smp_trampoline_end

.set SMP_VAR_ADDR, 0x8000 + (smp_trampoline_end - smp_trampoline)

smp_trampoline:
    cli
    cld
    ljmp 0, 0x8040

.align 16
_gdt_table:
    .long 0, 0
    .long 0x0000ffff, 0x00af9a00 # code 64
    .long 0x0000ffff, 0x00cf9200 # data
    .long 0x0000ffff, 0x00cf9a00 # code 32
_gdt:
    .word _gdt - _gdt_table - 1
    .long 0x8010
    .long 0, 0
.align 64

    xor ax, ax
    mov ds, ax
    lgdt [0x8030]
    mov eax, cr0
    or eax, 1
    mov cr0, eax
    ljmp 24, 0x8060

.align 32
.code32
    mov ax, 16
    mov ds, ax
    mov ss, ax

    # Get Local APIC ID
    mov eax, 1
    cpuid
    shr ebx, 24

    # Set PML4 physical address
    mov eax, [SMP_VAR_ADDR]
    mov cr3, eax

    # Enable PSE, PAE, OSFXSR, OSXMMEXCPT.
    # Do not force OSXSAVE here: some VMs/CPUs may fault before Rust-side
    # feature probing. init_cpu_extensions() enables OSXSAVE conditionally.
    mov eax, cr4
    or eax, 0x630
    mov cr4, eax

    # Enable LME
    mov ecx, 0xc0000080 # EFER
    xor edx, edx
    rdmsr
    or eax, 0x901
    wrmsr

    # Enable paging, write protect, and FPU/SSE
    mov eax, cr0
    and eax, 0xFFFFFFFB # Clear EM (bit 2)
    or eax, 0x80010002  # Set PG, WP, MP (bit 1)
    mov cr0, eax

    ljmp 8, 0x80c0

.align 32
.code64
    # Setup local stack
    mov rsp, [SMP_VAR_ADDR + 8]
    shl rbx, 3
    add rsp, rbx
    mov rsp, [rsp]

    push 0
    popfq

    movabs rax, offset smp_main
    jmp rax

.align 8
smp_trampoline_end:
"#
);

unsafe extern "C" {
    /// Performs the smp trampoline operation.
    fn smp_trampoline();
    /// Performs the smp trampoline end operation.
    fn smp_trampoline_end();
}

/// Busy-wait for the given number of microseconds (very rough).
fn udelay(us: u32) {
    for _ in 0..us {
        io_wait();
    }
}

/// Performs the ensure identity mapping operation.
fn ensure_identity_mapping(phys_start: u64, length: usize) {
    let start = phys_start & !0xFFFu64;
    let end = (phys_start + length as u64 + 0xFFF) & !0xFFFu64;
    let flags = PageTableFlags::PRESENT | PageTableFlags::WRITABLE;

    let mut addr = start;
    while addr < end {
        let virt = VirtAddr::new(addr);
        if let Some(mapped) = crate::memory::paging::translate(virt) {
            if mapped.as_u64() != addr {
                log::warn!(
                    "SMP: identity map collision at {:#x} -> {:#x}",
                    addr,
                    mapped.as_u64()
                );
            }
        } else {
            let page = Page::<Size4KiB>::containing_address(virt);
            let frame = PhysFrame::<Size4KiB>::containing_address(PhysAddr::new(addr));
            if let Err(e) = crate::memory::paging::map_page(page, frame, flags) {
                log::error!("SMP: failed to identity map {:#x}: {}", addr, e);
            }
        }
        addr += 0x1000;
    }
}

/// Performs the copy trampoline operation.
fn copy_trampoline(cr3_phys: u64, stacks_ptr: *const u64) {
    let tramp_len = (smp_trampoline_end as *const u8 as usize)
        .saturating_sub(smp_trampoline as *const u8 as usize);

    ensure_identity_mapping(TRAMPOLINE_PHYS_ADDR, tramp_len + 16);

    let tramp_virt = memory::phys_to_virt(TRAMPOLINE_PHYS_ADDR) as *mut u8;

    // SAFETY: trampoline destination is mapped and writable in HHDM.
    unsafe {
        core::ptr::copy_nonoverlapping(smp_trampoline as *const u8, tramp_virt, tramp_len);
        let ptrs = tramp_virt.add(tramp_len) as *mut u64;
        core::ptr::write_volatile(ptrs, cr3_phys);
        core::ptr::write_volatile(ptrs.add(1), stacks_ptr as u64);
    }
}

/// Performs the wait delivery operation.
fn wait_delivery() {
    const DELIVERY_STATUS: u32 = 1 << 12;
    for _ in 0..1_000_000 {
        // SAFETY: APIC initialized, ICR low is readable.
        let val = unsafe { apic::read_reg(apic::REG_ICR_LOW) };
        if val & DELIVERY_STATUS == 0 {
            return;
        }
        core::hint::spin_loop();
    }
    log::warn!("SMP: IPI delivery timeout");
}

/// Performs the send ipi operation.
fn send_ipi(apic_id: u32, value: u32) {
    apic::send_ipi_raw(apic_id, value);
    wait_delivery();
}

/// Performs the send init sipi operation.
fn send_init_sipi(apic_id: u32) {
    // INIT IPI (assert)
    send_ipi(apic_id, 0x0000_c500);
    udelay(10_000);

    // INIT de-assert
    send_ipi(apic_id, 0x0000_8500);
    udelay(200);

    // SIPI twice, vector = 0x8 (0x8000 >> 12)
    for _ in 0..2 {
        send_ipi(apic_id, 0x0000_0608);
        udelay(200);
    }
}

/// Broadcast a halt command to all other CPUs.
///
/// Used during panic to stop the system and prevent log corruption.
pub fn broadcast_panic_halt() {
    if !apic::is_initialized() {
        return;
    }
    // Destination Shorthand: 0b11 (All excluding self)
    // Delivery Mode: 0b100 (NMI)
    // Level: 1 (Assert)
    let icr_low = (0b11 << 18) | (0b100 << 8) | (1 << 14);
    apic::send_ipi_raw(0, icr_low);
}

/// Wait at a synchronization barrier until all expected CPUs arrive.
///
/// Every CPU (BSP + APs) calls this once.  BSP must store the target count
/// in `BARRIER_TARGET` before any CPU enters the barrier.
fn rendezvous_barrier() {
    let expected = BARRIER_TARGET.load(Ordering::Acquire);
    SYNC_BARRIER.fetch_add(1, Ordering::AcqRel);
    while SYNC_BARRIER.load(Ordering::Acquire) < expected {
        core::hint::spin_loop();
    }
}

/// Boot Application Processors.
pub fn init() -> Result<usize, &'static str> {
    if !apic::is_initialized() {
        return Err("APIC not initialized");
    }

    BOOTED_CORES.store(1, Ordering::Release);
    SYNC_BARRIER.store(0, Ordering::Release);
    BARRIER_TARGET.store(0, Ordering::Release);

    let madt_info = madt::parse_madt().ok_or("MADT not available")?;
    let bsp_apic_id = apic::lapic_id();

    if madt_info.local_apic_count <= 1 {
        log::info!("SMP: single CPU system");
        return Ok(1);
    }

    let mut max_apic_id: usize = 0;
    for i in 0..madt_info.local_apic_count {
        if let Some(ref entry) = madt_info.local_apics[i] {
            max_apic_id = max_apic_id.max(entry.apic_id as usize);
        }
    }

    let mut stacks: Vec<u64> = vec![0; max_apic_id + 1];
    let cr3_phys = crate::memory::paging::kernel_l4_phys().as_u64();
    let mut targets: Vec<u32> = Vec::new();
    let mut expected: usize = 1;

    for i in 0..madt_info.local_apic_count {
        let Some(ref entry) = madt_info.local_apics[i] else {
            continue;
        };

        let apic_id = entry.apic_id as u32;
        if apic_id == bsp_apic_id {
            continue;
        }

        // Allocate AP kernel stack from the buddy allocator.
        // boot_alloc is sealed after buddy init (to prevent double-allocation),
        // so AP stacks must come from buddy.  
        // AP stacks are permanent kernel allocations: we intentionally leak the frame so buddy never reclaims it.
        
        let stack_size = crate::process::task::Task::DEFAULT_STACK_SIZE;
        let pages = (stack_size + 4095) / 4096;
        let order = pages.next_power_of_two().trailing_zeros() as u8;
        let frame = crate::sync::with_irqs_disabled(|token| {
            crate::memory::allocate_frames(token, order)
        })
        .map_err(|_| "SMP: failed to allocate AP stack from buddy")?;
        let stack_phys = frame.start_address.as_u64();
        let stack_virt = crate::memory::phys_to_virt(stack_phys);

        // SAFETY: buddy gave us a valid, exclusive physical frame; phys_to_virt maps it.
        unsafe { core::ptr::write_bytes(stack_virt as *mut u8, 0, stack_size) };

        // Stack grows downward: top = base_virt + size.
        // The PhysFrame is Copy (no Drop): buddy keeps it marked as allocated
        // since we never call free_frames : permanent kernel allocation.
        let stack_top = stack_virt.saturating_add(stack_size as u64);

        if apic_id as usize >= stacks.len() {
            log::warn!("SMP: APIC id {} out of stack array range", apic_id);
            continue;
        }

        stacks[apic_id as usize] = stack_top;

        let cpu_index =
            percpu::register_cpu(apic_id).ok_or("SMP: exceeded MAX_CPUS for per-CPU data")?;
        percpu::set_kernel_stack_top(cpu_index, stack_top);

        // We no longer need to push to AP_KERNEL_STACKS as these are static.
        targets.push(apic_id);
        expected += 1;
    }

    copy_trampoline(cr3_phys, stacks.as_ptr());

    for apic_id in targets {
        send_init_sipi(apic_id);
    }

    // Do not spin forever if one AP fails very early (e.g. trampoline fault).
    // Keep booting with available CPUs and report partial bring-up.
    let mut spins: u64 = 0;
    const MAX_SPINS: u64 = 200_000_000;
    while BOOTED_CORES.load(Ordering::Acquire) < expected && spins < MAX_SPINS {
        core::hint::spin_loop();
        spins = spins.saturating_add(1);
    }
    if BOOTED_CORES.load(Ordering::Acquire) < expected {
        log::warn!(
            "SMP: timeout waiting APs (online={} expected={}), continuing",
            BOOTED_CORES.load(Ordering::Acquire),
            expected
        );
    }

    let online = BOOTED_CORES.load(Ordering::Acquire);
    log::info!("SMP: {} cores online (expected {})", online, expected);

    // Publish the barrier target so every CPU (BSP + APs) uses the same value.
    BARRIER_TARGET.store(online, Ordering::Release);
    // BSP reaches the rendezvous point.
    rendezvous_barrier();

    Ok(online)
}

/// First Rust function executed on APs after the trampoline.
#[unsafe(no_mangle)]
pub extern "C" fn smp_main() -> ! {
    idt::load();

    // Re-initialize Local APIC for this core (per-core registers).
    apic::init_ap();

    let apic_id = apic::lapic_id();
    let cpu_index = match percpu::cpu_index_by_apic(apic_id) {
        Some(idx) => idx,
        None => {
            // If we fall back to 0, this AP would share the BSP's
            // CPU index and double-increment TICK_COUNT 2× timer speed.
            log::error!("SMP AP: APIC id {} not registered — halting core", apic_id);
            loop {
                core::hint::spin_loop();
            }
        }
    };

    // Initialize per-CPU GS base.
    crate::arch::x86_64::percpu::init_gs_base(cpu_index);

    // Initialize per-CPU TSS/GDT (now uses O(1) current_cpu_index).
    crate::arch::x86_64::tss::init_cpu(cpu_index);
    crate::arch::x86_64::gdt::init_cpu(cpu_index);

    crate::arch::x86_64::syscall::init();
    crate::arch::x86_64::init_cpu_extensions();

    if let Some(stack_top) = percpu::kernel_stack_top(cpu_index) {
        crate::arch::x86_64::tss::set_kernel_stack_for(cpu_index, x86_64::VirtAddr::new(stack_top));
    }

    let _ = percpu::mark_online_by_apic(apic_id);
    BOOTED_CORES.fetch_add(1, Ordering::Release);

    // AP spins until BSP publishes the barrier target, then enters barrier.
    while BARRIER_TARGET.load(Ordering::Acquire) == 0 {
        core::hint::spin_loop();
    }
    rendezvous_barrier();

    crate::serial_println!(
        "[trace][ap] online cpu_index={} entering ap scheduler",
        cpu_index
    );

    // Wait until BSP has finished scheduler initialization.
    while !AP_SCHED_GATE_OPEN.load(Ordering::Acquire) {
        core::hint::spin_loop();
    }

    // Start APIC timer on this CPU (uses cached calibration from BSP).
    timer::start_apic_timer_cached();

    // Start per-CPU scheduler (never returns).
    crate::process::scheduler::schedule_on_cpu(cpu_index)
}

/// Return the number of online CPUs.
pub fn cpu_count() -> usize {
    BOOTED_CORES.load(Ordering::Acquire)
}

/// Allow APs to start their local timer and enter the scheduler.
pub fn open_ap_scheduler_gate() {
    AP_SCHED_GATE_OPEN.store(true, Ordering::Release);
}