lancer / kernel / src / sched / switch.rs

use core::arch::global_asm;

use lancer_core::object_layout::SchedContextObject;
use x86_64::structures::paging::PhysFrame;

use crate::arch::syscall as arch_syscall;
use crate::proc::context::CpuContext;
use crate::proc::{ProcessManager, ProcessState};
use crate::sync::SyncUnsafe;
use crate::types::{MAX_CPUS, Pid};
use crate::wcet::region::WcetTracker;

const IA32_FS_BASE: u32 = 0xC000_0100;

#[inline]
pub(crate) fn read_fs_base() -> u64 {
    unsafe {
        let lo: u32;
        let hi: u32;
        core::arch::asm!(
            "rdmsr",
            in("ecx") IA32_FS_BASE,
            out("eax") lo,
            out("edx") hi,
            options(nomem, nostack, preserves_flags),
        );
        ((hi as u64) << 32) | (lo as u64)
    }
}

#[inline]
pub(crate) fn write_fs_base(val: u64) {
    unsafe {
        core::arch::asm!(
            "wrmsr",
            in("ecx") IA32_FS_BASE,
            in("eax") val as u32,
            in("edx") (val >> 32) as u32,
            options(nomem, nostack, preserves_flags),
        );
    }
}

unsafe extern "C" {
    fn restore_context_and_iretq(ctx: *const CpuContext);
}

static SWITCH_WCET: [SyncUnsafe<WcetTracker<16>>; MAX_CPUS] =
    [const { SyncUnsafe::new(WcetTracker::new("context_switch")) }; MAX_CPUS];

#[allow(dead_code)]
pub fn log_switch_wcet() {
    let idx = arch_syscall::cpu_index();
    unsafe { SWITCH_WCET[idx].as_ref_unchecked() }.log_summary();
}

pub fn enter_context(ctx: &CpuContext) -> ! {
    unsafe { restore_context_and_iretq(ctx as *const CpuContext) };
    unreachable!()
}

pub fn switch_to(
    ptable: &mut ProcessManager,
    from_pid: Pid,
    to_pid: Pid,
    frame: &mut CpuContext,
) -> bool {
    if from_pid == to_pid {
        return false;
    }

    let idx = arch_syscall::cpu_index();
    let tracker = unsafe { SWITCH_WCET[idx].as_mut_unchecked() };
    let _guard = crate::wcet::region::WcetGuard::new(tracker);

    let (from_was_running, old_pml4, new_pml4, sc_ctx, fs_base) = {
        let Some(((from_sched, from_exec), (to_sched, to_exec))) =
            ptable.pair_mut(from_pid, to_pid)
        else {
            return false;
        };

        from_exec.fs_base = read_fs_base();
        from_exec.saved_context = *frame;
        from_exec.seal_context();

        from_sched.run_cpu = None;

        let from_was_running = from_sched.state() == ProcessState::Running;
        if from_was_running && from_sched.transition_to(ProcessState::Ready).is_err() {
            crate::kprintln!(
                "[sched] BUG: pid {} failed Running -> Ready",
                from_pid.raw()
            );
        }

        let old_pml4 = from_exec.pml4_phys.raw();

        if to_sched.transition_to(ProcessState::Running).is_err() {
            crate::kprintln!(
                "[sched] BUG: pid {} failed -> Running (state={:?}), skipping switch",
                to_pid.raw(),
                to_sched.state()
            );
            if from_was_running {
                let r = from_sched.transition_to(ProcessState::Running);
                debug_assert!(r.is_ok());
            }
            return false;
        }

        let new_pml4 = to_exec.pml4_phys.raw();

        to_exec.verify_context(to_pid);
        *frame = to_exec.saved_context;
        frame.validate_user(to_pid.raw());

        arch_syscall::set_fpu_state_ptr(to_exec.fpu_state.data.as_mut_ptr());

        to_sched.run_cpu = Some(idx as u16);

        let sc_ctx = to_sched.sched_context();
        let fs_base = to_exec.fs_base;

        (from_was_running, old_pml4, new_pml4, sc_ctx, fs_base)
    };

    if from_was_running {
        ptable.enqueue_ready(from_pid);
    }

    if let Some((sc_id, sc_gen)) = sc_ctx {
        let now_us = crate::wcet::tsc::cycles_to_ns(crate::wcet::tsc::read_tsc()) / 1000;
        let mut pool = crate::cap::pool::POOL.lock();
        if let Ok(sc) = pool.write_as::<SchedContextObject>(sc_id, sc_gen) {
            super::context::mark_dispatched(sc, now_us);
        }
    }

    write_fs_base(fs_base);

    arch_syscall::set_current_process(to_pid);
    super::set_current(to_pid);

    if old_pml4 != new_pml4 {
        unsafe {
            x86_64::registers::control::Cr3::write(
                PhysFrame::containing_address(new_pml4),
                x86_64::registers::control::Cr3Flags::empty(),
            );
        }
    }

    true
}

global_asm!(
    r#"
.global restore_context_and_iretq
.type restore_context_and_iretq, @function
restore_context_and_iretq:
    cli

    push qword ptr [rdi + 152]
    push qword ptr [rdi + 128]
    push qword ptr [rdi + 136]
    push qword ptr [rdi + 144]
    push qword ptr [rdi + 120]

    mov rax, [rdi + 0]
    mov rbx, [rdi + 8]
    mov rcx, [rdi + 16]
    mov rdx, [rdi + 24]
    mov rsi, [rdi + 32]
    mov rbp, [rdi + 48]
    mov r8,  [rdi + 56]
    mov r9,  [rdi + 64]
    mov r10, [rdi + 72]
    mov r11, [rdi + 80]
    mov r12, [rdi + 88]
    mov r13, [rdi + 96]
    mov r14, [rdi + 104]
    mov r15, [rdi + 112]
    mov rdi, [rdi + 40]

    push rax
    mov rax, [rsp + 16]
    test rax, 3
    pop rax
    jz .restore_no_swap
    swapgs
.restore_no_swap:
    iretq

.size restore_context_and_iretq, . - restore_context_and_iretq
"#
);