rust-kernel

//! Preemptive kernel threads (beyond the phil-opp series).
//!
//! Each thread has its own stack. The timer interrupt calls `schedule()`,
//! which picks the next thread round-robin and calls `switch_context` to swap
//! stacks. Threads don't have to cooperate: even `loop {}` gets interrupted.
//!
//! How a switch works (all inside the timer interrupt, interrupts off):
//!
//!   thread A running -> timer IRQ -> CPU pushes A's RIP/RSP/RFLAGS (the
//!   interrupt frame) on A's stack -> the `x86-interrupt` handler saves the
//!   caller-saved registers (rax, rcx, rdx, rsi, rdi, r8-r11) on A's stack ->
//!   `schedule()` -> `switch_context(&A.rsp, B.rsp)`:
//!       push the callee-saved registers (rbx, rbp, r12-r15) on A's stack
//!       save rsp into A.rsp
//!       load rsp from B.rsp            <- now we're on B's stack
//!       pop B's callee-saved registers
//!       ret                            <- returns into B's earlier schedule()
//!   ... which returns into B's timer handler, which restores B's caller-saved
//!   registers and `iretq`s back to wherever B was interrupted.
//!
//! So a sleeping thread's whole CPU state is simply sitting on its own stack,
//! and all we have to remember per thread is its stack pointer.
//!
//! A *new* thread has never been interrupted, so `spawn` builds a fake stack
//! that looks like `switch_context` saved it, with the "return address"
//! pointing at `thread_trampoline`, which enables interrupts and calls the
//! thread's function.

use alloc::{boxed::Box, vec, vec::Vec};
use core::arch::global_asm;
use spin::Mutex;
use x86_64::instructions::interrupts;

/// 16 KiB per thread stack (no guard page! a deep recursion in a thread
/// silently overwrites whatever is below its stack on the heap).
const STACK_SIZE: usize = 16 * 1024;

struct Thread {
    name: &'static str,
    /// Saved stack pointer while the thread is not running.
    rsp: u64,
    /// The thread's stack. `None` for the boot thread (uses the bootloader's stack).
    _stack: Option<Vec<u8>>,
    finished: bool,
}

struct Scheduler {
    /// Boxed so each Thread has a fixed address even if the Vec grows: we
    /// hand a raw pointer to `rsp` to `switch_context`.
    threads: Vec<Box<Thread>>,
    current: usize,
}

static SCHEDULER: Mutex<Scheduler> = Mutex::new(Scheduler { threads: Vec::new(), current: 0 });

/// Start a new kernel thread running `entry`. Needs the heap.
pub fn spawn(name: &'static str, entry: fn()) {
    let mut stack = vec![0u8; STACK_SIZE];

    // Build the initial stack, top down (stacks grow towards lower addresses):
    //   [return address = thread_trampoline]
    //   rbp, rbx, r12 (= entry), r13, r14, r15     <- rsp points at r15
    // exactly what `switch_context` pops before its `ret`.
    let top = (stack.as_mut_ptr() as u64 + STACK_SIZE as u64) & !0xf; // 16-byte aligned
    let initial: [u64; 7] = [
        0,                                // r15
        0,                                // r14
        0,                                // r13
        entry as *const () as u64,        // r12: picked up by the trampoline
        0,                                // rbx
        0,                                // rbp
        thread_trampoline as *const () as u64, // return address
    ];
    let rsp = top - 8 * initial.len() as u64;
    unsafe { core::ptr::copy_nonoverlapping(initial.as_ptr(), rsp as *mut u64, initial.len()) };

    let thread = Box::new(Thread { name, rsp, _stack: Some(stack), finished: false });

    // The timer handler also takes this lock, so keep interrupts off while we
    // hold it -- otherwise the handler would spin forever on our lock.
    interrupts::without_interrupts(|| {
        let mut scheduler = SCHEDULER.lock();
        if scheduler.threads.is_empty() {
            // First spawn: register the code that's running right now (the
            // boot/main thread) as thread 0. Its rsp is filled in the first
            // time we switch away from it.
            scheduler.threads.push(Box::new(Thread {
                name: "main",
                rsp: 0,
                _stack: None,
                finished: false,
            }));
        }
        scheduler.threads.push(thread);
    });
}

/// Called from the timer interrupt (interrupts disabled, EOI already sent).
/// Switches to the next runnable thread, round-robin. Returns when this
/// thread is scheduled again.
pub fn schedule() {
    let (old_rsp, new_rsp) = {
        let mut s = SCHEDULER.lock();
        let n = s.threads.len();
        if n < 2 {
            return; // no threads spawned yet
        }
        let current = s.current;
        // Find the next thread that hasn't finished.
        let next = (1..=n).map(|i| (current + i) % n).find(|&i| !s.threads[i].finished);
        let next = match next {
            Some(next) if next != current => next,
            _ => return, // nobody else to run
        };
        s.current = next;
        let old_rsp: *mut u64 = &mut s.threads[current].rsp;
        (old_rsp, s.threads[next].rsp)
        // The lock is released here, *before* switching: the next thread
        // will never come back to release a lock we still held.
    };
    unsafe { switch_context(old_rsp, new_rsp) };
}

/// Name of the running thread (for printing).
pub fn current_name() -> &'static str {
    interrupts::without_interrupts(|| {
        let s = SCHEDULER.lock();
        s.threads.get(s.current).map_or("main", |t| t.name)
    })
}

/// Sleep for `ticks` timer ticks (~55 ms each), letting other threads run.
pub fn sleep_ticks(ticks: u64) {
    use crate::interrupts::TICKS;
    use core::sync::atomic::Ordering;
    let wake_at = TICKS.load(Ordering::Relaxed) + ticks;
    while TICKS.load(Ordering::Relaxed) < wake_at {
        x86_64::instructions::hlt(); // the timer interrupt may switch threads here
    }
}

// --- assembly ---------------------------------------------------------------

unsafe extern "C" {
    /// Save callee-saved registers + rsp into `*old_rsp`, load `new_rsp`,
    /// restore that thread's registers and return into it.
    fn switch_context(old_rsp: *mut u64, new_rsp: u64);
    /// First "return address" of every new thread.
    fn thread_trampoline();
}

// System V calling convention: 1st argument in rdi, 2nd in rsi. A normal
// function call already saved the caller-saved registers, so only the
// callee-saved ones (rbx, rbp, r12-r15) need saving here.
global_asm!(
    ".global switch_context",
    "switch_context:",
    "    push rbp",
    "    push rbx",
    "    push r12",
    "    push r13",
    "    push r14",
    "    push r15",
    "    mov [rdi], rsp", // old thread's rsp -> *old_rsp
    "    mov rsp, rsi",   // switch stacks!
    "    pop r15",
    "    pop r14",
    "    pop r13",
    "    pop r12",
    "    pop rbx",
    "    pop rbp",
    "    ret",            // pops the new thread's return address
    "",
    ".global thread_trampoline",
    "thread_trampoline:",
    "    mov rdi, r12",   // entry function pointer -> 1st argument
    "    and rsp, -16",   // ABI: stack 16-byte aligned before `call`
    "    call {start}",
    "    ud2",            // thread_start never returns
    start = sym thread_start,
);

/// Rust side of a new thread's life.
extern "C" fn thread_start(entry: usize) -> ! {
    // We arrived here from inside the timer interrupt, where interrupts are
    // off. Turn them back on, or this thread could never be preempted.
    interrupts::enable();

    let entry: fn() = unsafe { core::mem::transmute(entry) };
    entry();

    // The thread function returned: mark it finished so the scheduler skips
    // it, then idle until we're switched away for good. (Its stack is kept;
    // freeing it while we still run on it would be a bad idea.)
    interrupts::without_interrupts(|| {
        let mut s = SCHEDULER.lock();
        let current = s.current;
        s.threads[current].finished = true;
    });
    crate::hlt_loop();
}