opencoreC/kernel/src/sched/mod.rs

//! Preemptive & Cooperative Microkernel Scheduler

pub mod thread;

use thread::{Thread, ThreadState, MAX_THREADS};
use crate::arch::gdt::set_kernel_stack;
use crate::mm::vmm::VMM;
use crate::kprintln;

extern "C" {
    fn switch_to(prev_rsp: *mut u64, next_rsp: u64);
    fn thread_entry_trampoline();
}

const DEFAULT_QUANTUM_TICKS: u32 = 5; // 5 ticks @ 100Hz = 50ms time slice

pub struct Scheduler {
    pub threads: [Thread; MAX_THREADS],
    pub current_tid: usize,
    pub next_tid: u64,
    pub quantum_remaining: u32,
    pub ticks: u64,
}

pub static mut SCHEDULER: Scheduler = Scheduler {
    threads: [const { Thread::empty() }; MAX_THREADS],
    current_tid: 0,
    next_tid: 1,
    quantum_remaining: DEFAULT_QUANTUM_TICKS,
    ticks: 0,
};

impl Scheduler {
    pub fn init(&mut self) {
        self.threads[0].id = 0;
        self.threads[0].name = "kernel_idle";
        self.threads[0].state = ThreadState::Running;
        self.current_tid = 0;
        self.next_tid = 1;
        self.quantum_remaining = DEFAULT_QUANTUM_TICKS;
        self.ticks = 0;
        kprintln!("[SCHED] Preemptive Multi-Tasking Scheduler initialized (Quantum: {} ticks).", DEFAULT_QUANTUM_TICKS);
    }

    pub fn create_user_thread(
        &mut self,
        name: &'static str,
        entry_rip: u64,
        user_rsp: u64,
        pml4_paddr: u64,
        kernel_stack_top: u64,
    ) -> Option<u64> {
        for i in 1..MAX_THREADS {
            if self.threads[i].state == ThreadState::Unused {
                let tid = self.next_tid;
                self.next_tid += 1;

                self.threads[i].id = tid;
                self.threads[i].name = name;
                self.threads[i].state = ThreadState::Ready;
                self.threads[i].pml4_paddr = pml4_paddr;
                self.threads[i].kernel_stack_top = kernel_stack_top;
                self.threads[i].user_rsp = user_rsp;
                self.threads[i].user_rip = entry_rip;

        // Build initial kernel stack frame for new thread:
        // switch_to will pop r15, r14, r13, r12, rbx, rbp, then ret.
        // We configure r13 = entry_rip, r12 = user_rsp, ret = thread_entry_trampoline.
        unsafe {
            let k = kernel_stack_top as *mut u64;
            *k.sub(1) = thread_entry_trampoline as u64; // ret address
            *k.sub(2) = 0;                              // rbp
            *k.sub(3) = 0;                              // rbx
            *k.sub(4) = user_rsp;                       // r12 -> user_rsp
            *k.sub(5) = entry_rip;                      // r13 -> entry_rip
            *k.sub(6) = 0;                              // r14
            *k.sub(7) = 0;                              // r15

            self.threads[i].context.rsp = k.sub(7) as u64;
        }

                kprintln!("[SCHED] Created Preemptible Thread [TID {}] '{}' (RIP: {:#x}, RSP: {:#x})",
                    tid, name, entry_rip, user_rsp);
                return Some(tid);
            }
        }
        None
    }

    pub fn get_current_mut(&mut self) -> &mut Thread {
        &mut self.threads[self.current_tid]
    }

    pub fn get_thread_mut(&mut self, tid: u64) -> Option<&mut Thread> {
        for t in self.threads.iter_mut() {
            if t.id == tid && t.state != ThreadState::Unused {
                return Some(t);
            }
        }
        None
    }

    pub fn block_current(&mut self, new_state: ThreadState) {
        self.threads[self.current_tid].state = new_state;
        self.schedule();
    }

    pub fn unblock(&mut self, tid: u64) {
        if let Some(t) = self.get_thread_mut(tid) {
            if t.state != ThreadState::Unused && t.state != ThreadState::Dead {
                t.state = ThreadState::Ready;
            }
        }
    }

    pub fn kill_thread(&mut self, tid: u64) -> bool {
        if tid <= 3 {
            return false; // Protect kernel_idle, init, uart, and sh
        }
        for t in self.threads.iter_mut() {
            if t.id == tid && t.state != ThreadState::Unused && t.state != ThreadState::Dead {
                t.state = ThreadState::Dead;
                return true;
            }
        }
        false
    }

    pub fn yield_current(&mut self) {
        if self.threads[self.current_tid].state == ThreadState::Running {
            self.threads[self.current_tid].state = ThreadState::Ready;
        }
        self.quantum_remaining = 0;
    }

    /// Preemptive tick called directly from timer ISR handler
    pub fn preempt_tick(&mut self) {
        self.ticks += 1;

        if self.quantum_remaining > 0 {
            self.quantum_remaining -= 1;
        }

        if self.quantum_remaining == 0 {
            self.quantum_remaining = DEFAULT_QUANTUM_TICKS;
            self.schedule();
        }
    }

    pub fn schedule(&mut self) {
        let prev_idx = self.current_tid;
        let mut next_idx = (prev_idx + 1) % MAX_THREADS;

        let mut found = false;
        for _ in 0..MAX_THREADS {
            if self.threads[next_idx].state == ThreadState::Ready {
                found = true;
                break;
            }
            next_idx = (next_idx + 1) % MAX_THREADS;
        }

        if !found {
            if self.threads[prev_idx].state == ThreadState::Running {
                return;
            }
            next_idx = 0;
        }

        if prev_idx == next_idx && self.threads[prev_idx].state == ThreadState::Running {
            return;
        }

        if self.threads[prev_idx].state == ThreadState::Running {
            self.threads[prev_idx].state = ThreadState::Ready;
        }

        self.current_tid = next_idx;
        self.threads[next_idx].state = ThreadState::Running;

        let next_pml4 = self.threads[next_idx].pml4_paddr;
        let next_kstack = self.threads[next_idx].kernel_stack_top;
        let next_rsp = self.threads[next_idx].context.rsp;
        let prev_rsp_ptr = core::ptr::addr_of_mut!(self.threads[prev_idx].context.rsp);

        unsafe {
            if next_pml4 != 0 {
                (*core::ptr::addr_of_mut!(VMM)).load_cr3(next_pml4);
            }
            if next_kstack != 0 {
                set_kernel_stack(next_kstack);
            }
            switch_to(prev_rsp_ptr, next_rsp);
        }
    }
}

pub fn init() {
    unsafe {
        (*core::ptr::addr_of_mut!(SCHEDULER)).init();
    }
}