1 files changed, 544 insertions, 0 deletions

diff --git a/rsnix/src/sched.rs b/rsnix/src/sched.rs
new file mode 100644
index 0000000..4b10503
--- /dev/null
+++ b/rsnix/src/sched.rs
@@ -0,0 +1,544 @@
+// ---------------------------------------------------------------------------
+// sched — preemptive round-robin kernel tasks.
+//
+// 16 static task slots.  The 100 Hz PIT tick (vector 32) and the
+// software yield vector (0x60) go through `isr_common`, which saves the
+// running task's full IntFrame on its own stack and asks the scheduler
+// for the next frame pointer — the asm then moves rsp there and iretq's,
+// which *is* the context switch.
+//
+//   - task 0: the shell (boot stack, never allocated by us)
+//   - task 1: idle (hlt loop, always ready — the lowest-priority
+//     fallback)
+//   - tasks 2..: spawn'd, 8 KiB stack + one unmapped guard page below
+//
+// Preemption is disabled while the scheduler lock (PREEMPT_DEPTH) is
+// held: the tick/yield handlers bail out entirely when the depth is
+// > 0, so table mutations and klog/serial line writes are atomic with
+// respect to switching.  The lock is a counter, not a mutex — re-entrant
+// by design (a panic mid-emit increments again; the panic path clears
+// it, because a dead task's lock is meaningless).
+//
+// Task page faults (vector 14) run on a dedicated IST stack (see gdt.rs)
+// so a stack overflow inside a guard page kills the task gracefully
+// instead of double-faulting the machine.
+// ---------------------------------------------------------------------------
+use core::arch::asm;
+use core::fmt;
+use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
+
+extern crate alloc;
+use alloc::string::String;
+use alloc::vec::Vec;
+
+use crate::interrupts::IntFrame;
+use crate::memory;
+use crate::paging;
+
+pub const MAX_TASKS: usize = 16;
+pub const SHELL_TID: usize = 0;
+pub const IDLE_TID: usize = 1;
+const STACK_PAGES: usize = 3; // 1 guard page + 8 KiB stack
+const NAME_MAX: usize = 24;
+const PANIC_MAX: usize = 64;
+
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub enum TaskState {
+    Ready,
+    Sleeping,
+    Blocked,
+    Done,
+}
+
+struct Task {
+    name: [u8; NAME_MAX],
+    name_len: u8,
+    state: TaskState,
+    /// Frame pointer to restore on switch-in (set at spawn; stored on
+    /// every switch-away).
+    saved_rsp: usize,
+    /// Physical base of the stack allocation (0 = none, e.g. the shell).
+    stack: usize,
+    /// Guard page address (unmapped; 0 = none).
+    guard: usize,
+    wake_at: u64,
+    ticks: u64,
+    panic: [u8; PANIC_MAX],
+    panic_len: u8,
+}
+
+const TASK_NONE: Task = Task {
+    name: [0; NAME_MAX],
+    name_len: 0,
+    state: TaskState::Done,
+    saved_rsp: 0,
+    stack: 0,
+    guard: 0,
+    wake_at: 0,
+    ticks: 0,
+    panic: [0; PANIC_MAX],
+    panic_len: 0,
+};
+
+static mut TASKS: [Task; MAX_TASKS] = [TASK_NONE; MAX_TASKS];
+static mut CURRENT: usize = SHELL_TID;
+static mut KERNEL_CS: u16 = 0;
+static mut KERNEL_SS: u16 = 0;
+
+static SCHED_READY: AtomicBool = AtomicBool::new(false);
+static IN_ISR: AtomicBool = AtomicBool::new(false);
+static PREEMPT_DEPTH: AtomicU64 = AtomicU64::new(0);
+
+// ---------------------------------------------------------------------------
+// Preemption lock
+// ---------------------------------------------------------------------------
+
+/// Disable preemption (re-entrant counter).  The tick/yield handlers
+/// skip all scheduling work while the depth is > 0.
+pub fn lock() {
+    PREEMPT_DEPTH.fetch_add(1, Ordering::Relaxed);
+}
+
+/// Re-enable preemption (must pair with `lock`).
+pub fn unlock() {
+    PREEMPT_DEPTH.fetch_sub(1, Ordering::Relaxed);
+}
+
+pub fn in_isr() -> bool {
+    IN_ISR.load(Ordering::Relaxed)
+}
+pub fn set_in_isr(v: bool) {
+    IN_ISR.store(v, Ordering::Relaxed);
+}
+pub fn is_ready() -> bool {
+    SCHED_READY.load(Ordering::Relaxed)
+}
+pub fn current() -> usize {
+    unsafe { CURRENT }
+}
+
+// ---------------------------------------------------------------------------
+// Tick / yield entry points (called from the ISR dispatch)
+// ---------------------------------------------------------------------------
+
+/// Timer tick: wake sleepers, account ticks, pick the next task.
+/// Receives the interrupted task's frame pointer; returns the frame
+/// pointer the asm should switch to (the same one when no switch).
+pub fn on_tick(cur: usize) -> usize {
+    if !is_ready() || PREEMPT_DEPTH.load(Ordering::Relaxed) > 0 {
+        return cur;
+    }
+    unsafe {
+        TASKS[CURRENT].saved_rsp = cur;
+        TASKS[CURRENT].ticks += 1;
+        let now = crate::interrupts::uptime_ms();
+        for t in 0..MAX_TASKS {
+            if TASKS[t].state == TaskState::Sleeping && TASKS[t].wake_at <= now {
+                TASKS[t].state = TaskState::Ready;
+            }
+        }
+        pick_next(cur)
+    }
+}
+
+/// Software yield (vector 0x60).
+pub fn on_yield(cur: usize) -> usize {
+    if !is_ready() || PREEMPT_DEPTH.load(Ordering::Relaxed) > 0 {
+        return cur;
+    }
+    unsafe {
+        TASKS[CURRENT].saved_rsp = cur;
+        pick_next(cur)
+    }
+}
+
+/// Round-robin over Ready tasks, skipping the idle task unless nothing
+/// else can run.  Returns the next frame pointer (or `cur` to stay).
+unsafe fn pick_next(cur: usize) -> usize {
+    unsafe {
+        for i in 1..MAX_TASKS {
+            let idx = (CURRENT + i) % MAX_TASKS;
+            if idx == IDLE_TID {
+                continue;
+            }
+            if TASKS[idx].state == TaskState::Ready {
+                CURRENT = idx;
+                return TASKS[idx].saved_rsp;
+            }
+        }
+        if CURRENT != IDLE_TID && TASKS[IDLE_TID].state == TaskState::Ready {
+            CURRENT = IDLE_TID;
+            return TASKS[IDLE_TID].saved_rsp;
+        }
+        cur
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Spawn / lifecycle
+// ---------------------------------------------------------------------------
+
+/// Spawn a kernel task running `entry` on its own 8 KiB stack with a
+/// guard page below it.  Returns the task id, or None when the table is
+/// full (Done slots are reaped first).  `entry` runs inside a
+/// trampoline, so returning from it exits the task cleanly.
+pub fn spawn(entry: fn(), name: &str) -> Option<usize> {
+    lock();
+    let id = unsafe {
+        // Reap a Done task with a stack, or take a never-used slot.
+        let mut slot = None;
+        for i in 2..MAX_TASKS {
+            if TASKS[i].state == TaskState::Done && TASKS[i].stack != 0 {
+                reap(i);
+                slot = Some(i);
+                break;
+            }
+            if TASKS[i].state == TaskState::Done && TASKS[i].stack == 0 && TASKS[i].name_len == 0 {
+                slot = Some(i);
+                break;
+            }
+        }
+        slot?
+    };
+    let ok = spawn_at(id, entry, name);
+    unlock();
+    if ok {
+        Some(id)
+    } else {
+        None
+    }
+}
+
+unsafe fn reap(id: usize) {
+    unsafe {
+        let t = &mut TASKS[id];
+        if t.stack != 0 {
+            // Restore the guard page's mapping so the pages are reusable.
+            paging::map_page(t.stack, paging::PAGE_WRITABLE);
+            memory::free_pages(t.stack, STACK_PAGES);
+        }
+        *t = TASK_NONE;
+    }
+}
+
+fn spawn_at(id: usize, entry: fn(), name: &str) -> bool {
+    let Some(base) = memory::alloc_pages(STACK_PAGES) else {
+        return false;
+    };
+    // The first page is the guard — unmap it so any stack overflow
+    // faults there (caught on the vector-14 IST stack).
+    if !paging::unmap_page(base) {
+        memory::free_pages(base, STACK_PAGES);
+        return false;
+    }
+    let top = base + STACK_PAGES * memory::PAGE_SIZE;
+    let frame = (top - core::mem::size_of::<IntFrame>()) as *mut IntFrame;
+    let mut name_buf = [0u8; NAME_MAX];
+    let n = name.len().min(NAME_MAX);
+    name_buf[..n].copy_from_slice(&name.as_bytes()[..n]);
+    unsafe {
+        core::ptr::write(
+            frame,
+            IntFrame {
+                r15: 0,
+                r14: 0,
+                r13: 0,
+                r12: 0,
+                r11: 0,
+                r10: 0,
+                r9: 0,
+                r8: 0,
+                rdi: 0,
+                rsi: 0,
+                rbp: 0,
+                rbx: 0,
+                rdx: 0,
+                rcx: entry as usize as u64, // MS x64 ABI: first arg
+                rax: 0,
+                vector: 0,
+                error: 0,
+                rip: trampoline as *const () as u64,
+                cs: KERNEL_CS as u64,
+                rflags: 0x202, // IF | reserved bit 1
+                // QEMU's iretq always pops RSP/SS from the frame tail.
+                // rsp = the stack top so the first iretq lands with a
+                // valid stack pointer (the trampoline then descends from
+                // it normally); ss mirrors the kernel data selector.
+                rsp: top as u64,
+                ss: KERNEL_SS as u64,
+            },
+        );
+        TASKS[id] = Task {
+            name: name_buf,
+            name_len: n as u8,
+            state: TaskState::Ready,
+            saved_rsp: frame as usize,
+            stack: base,
+            guard: base,
+            wake_at: 0,
+            ticks: 0,
+            panic: [0; PANIC_MAX],
+            panic_len: 0,
+        };
+    }
+    true
+}
+
+/// Entry trampoline: run the task body, then exit the task when it
+/// returns.  Runs with the synthetic frame's registers (rcx = entry).
+unsafe extern "C" fn trampoline(entry: usize) -> ! {
+    let f: fn() = unsafe { core::mem::transmute(entry as *const ()) };
+    f();
+    task_exit()
+}
+
+fn task_exit() -> ! {
+    lock();
+    unsafe {
+        TASKS[CURRENT].state = TaskState::Done;
+    }
+    unlock();
+    yield_int();
+    unreachable!()
+}
+
+/// Block the current task for `ms` milliseconds.  The shell uses this
+/// for `sleep`; background tasks keep running meanwhile.
+pub fn sleep_ms(ms: u64) {
+    lock();
+    unsafe {
+        TASKS[CURRENT].state = TaskState::Sleeping;
+        TASKS[CURRENT].wake_at = crate::interrupts::uptime_ms() + ms;
+    }
+    unlock();
+    yield_int();
+}
+
+/// Cooperatively yield the rest of the current quantum.
+pub fn yield_now() {
+    yield_int();
+}
+
+/// Software-interrupt the scheduler (vector 0x60).  Net stack usage is
+/// zero (the frame is popped by iretq), so no `nostack` — the compiler
+/// must not use the red zone around the interrupt.
+#[inline(never)]
+fn yield_int() {
+    unsafe {
+        asm!("int 0x60");
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Task panics and faults
+// ---------------------------------------------------------------------------
+
+/// Kill the current task with a panic message: logs it, marks the task
+/// Done, drops its scheduler lock (a dead task's lock is meaningless),
+/// and switches away forever.  Called by the panic hook (task context)
+/// and by the page-fault handler (IST stack, ISR context).
+pub fn abort_current(msg: &str) -> ! {
+    unsafe {
+        let id = CURRENT;
+        let n = msg.len().min(PANIC_MAX);
+        TASKS[id].panic[..n].copy_from_slice(&msg.as_bytes()[..n]);
+        TASKS[id].panic_len = n as u8;
+        TASKS[id].state = TaskState::Done;
+        crate::serial_println!(
+            "sched: task {} \"{}\" panicked: {} (rsp_saved={:#x} stack={:#x} guard={:#x})",
+            id,
+            name_str(id),
+            msg,
+            TASKS[id].saved_rsp,
+            TASKS[id].stack,
+            TASKS[id].guard
+        );
+        PREEMPT_DEPTH.store(0, Ordering::Relaxed);
+    }
+    yield_int();
+    unreachable!()
+}
+
+/// Vector 14 (page fault) handler — runs on the dedicated IST stack.
+/// A fault in a real task (id >= 2) kills the task gracefully; anything
+/// else (shell, idle, pre-scheduler) is a kernel bug: halt.
+pub fn page_fault(cr2: u64, frame_ptr: usize) -> usize {
+    if !is_ready() || unsafe { CURRENT } <= IDLE_TID {
+        crate::serial_println!("pgfault: kernel page fault at {:#x} — halting", cr2);
+        crate::interrupts::halt();
+    }
+    let g = unsafe { TASKS[CURRENT].guard };
+    let in_guard = g != 0 && cr2 >= g as u64 && cr2 < (g + memory::PAGE_SIZE) as u64;
+    let rip = unsafe { (frame_ptr as *const IntFrame).read().rip };
+    let mut buf = [0u8; 96];
+    let n = {
+        let mut f = FmtBuf { b: &mut buf, n: 0 };
+        let _ = fmt::write(
+            &mut f,
+            core::format_args!(
+                "{} at {:#x} (rip={:#x})",
+                if in_guard {
+                    "stack overflow (guard page)"
+                } else {
+                    "page fault"
+                },
+                cr2,
+                rip
+            ),
+        );
+        f.n
+    };
+    abort_current(core::str::from_utf8(&buf[..n]).unwrap_or("task fault"))
+}
+
+// ---------------------------------------------------------------------------
+// Introspection (shell: jobs / kill)
+// ---------------------------------------------------------------------------
+
+/// Snapshot the task table for `jobs` — one line per used slot.
+pub fn jobs() -> Vec<String> {
+    let mut v = Vec::new();
+    lock();
+    unsafe {
+        for i in 0..MAX_TASKS {
+            let t = &TASKS[i];
+            if t.name_len == 0 {
+                continue; // never used
+            }
+            let state = match t.state {
+                TaskState::Ready => "ready",
+                TaskState::Sleeping => "sleeping",
+                TaskState::Blocked => "blocked",
+                TaskState::Done => {
+                    if t.panic_len > 0 {
+                        "done (panicked)"
+                    } else {
+                        "done"
+                    }
+                }
+            };
+            let mut line = alloc::format!(
+                "  {:>2}  {:<12} {:<15} {:>6} ticks",
+                i,
+                name_str(i),
+                state,
+                t.ticks
+            );
+            if t.panic_len > 0 {
+                let m = core::str::from_utf8(&t.panic[..t.panic_len as usize]).unwrap_or("?");
+                line.push_str("  ");
+                line.push_str(m);
+            }
+            v.push(line);
+        }
+    }
+    unlock();
+    v
+}
+
+/// Mark a task Done so it stops running and gets reaped at the next
+/// spawn.  The shell (0) and idle (1) are protected.
+pub fn kill(id: usize) -> Result<(), &'static str> {
+    if id == SHELL_TID {
+        return Err("cannot kill the shell (task 0)");
+    }
+    if id == IDLE_TID {
+        return Err("cannot kill the idle task (task 1)");
+    }
+    if id >= MAX_TASKS {
+        return Err("no such task");
+    }
+    lock();
+    let alive = unsafe {
+        let t = &TASKS[id];
+        t.name_len != 0 && t.state != TaskState::Done
+    };
+    if !alive {
+        unlock();
+        return Err("no such task");
+    }
+    unsafe {
+        TASKS[id].state = TaskState::Done;
+    }
+    unlock();
+    Ok(())
+}
+
+pub fn name_str(id: usize) -> &'static str {
+    unsafe {
+        let t = &TASKS[id];
+        if t.name_len == 0 {
+            return "";
+        }
+        core::str::from_utf8(&t.name[..t.name_len as usize]).unwrap_or("")
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Init
+// ---------------------------------------------------------------------------
+
+/// Set up the task table: task 0 is the shell (boot stack), task 1 is
+/// the idle `hlt` fallback.  Called after `interrupts::init` (the tick
+/// must be live so `sleep`/wake scheduling works).
+pub fn init() {
+    unsafe {
+        let cs64: u64;
+        asm!("mov {0}, cs", out(reg) cs64, options(nostack, preserves_flags));
+        KERNEL_CS = cs64 as u16;
+        let ss64: u64;
+        asm!("mov {0}, ss", out(reg) ss64, options(nostack, preserves_flags));
+        KERNEL_SS = ss64 as u16;
+
+        TASKS[SHELL_TID] = Task {
+            name: {
+                let mut n = [0u8; 24];
+                n[..5].copy_from_slice(b"shell");
+                n
+            },
+            name_len: 5,
+            state: TaskState::Ready,
+            saved_rsp: 0, // captured on the first switch-away
+            stack: 0,
+            guard: 0,
+            wake_at: 0,
+            ticks: 0,
+            panic: [0; PANIC_MAX],
+            panic_len: 0,
+        };
+        spawn_at(IDLE_TID, idle_main, "idle");
+        SCHED_READY.store(true, Ordering::Relaxed);
+    }
+    crate::serial_println!(
+        "sched: ready — {} slots, shell=task 0, idle=task 1",
+        MAX_TASKS
+    );
+    crate::serial_println!(
+        "sched: trampoline={:#x} idle={:#x} emit={:#x}",
+        trampoline as *const () as usize,
+        idle_main as *const () as usize,
+        crate::klog::emit as *const () as usize
+    );
+}
+
+fn idle_main() {
+    loop {
+        unsafe {
+            asm!("hlt", options(nostack, preserves_flags));
+        }
+    }
+}
+
+struct FmtBuf<'a> {
+    b: &'a mut [u8],
+    n: usize,
+}
+
+impl fmt::Write for FmtBuf<'_> {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        let m = s.len().min(self.b.len() - self.n);
+        self.b[self.n..self.n + m].copy_from_slice(&s.as_bytes()[..m]);
+        self.n += m;
+        Ok(())
+    }
+}