sched: Refactor scheduler

Author: marv7000

kernel/src/arch/executor.rs

@@ -0,0 +1,33 @@
+use super::internal;
+use crate::{memory::stack::KernelStack, posix::errno::EResult, process::task::Task};
+
+pub use internal::executor::{Executor, ForkedImage, FrameImage};
+
+pub fn new_kernel(
+    entry: extern "C" fn(usize, usize),
+    arg1: usize,
+    arg2: usize,
+    stack: &KernelStack,
+) -> EResult<Executor> {
+    internal::executor::Executor::new_kernel(entry, arg1, arg2, stack)
+}
+
+/// Switches the current CPU context from one task executor to another.
+///
+/// # Safety
+/// The caller must ensure that `from` and `to` are valid, distinct live tasks.
+pub unsafe fn switch(from: *const Task, to: *const Task) -> *mut Task {
+    unsafe { internal::executor::switch(from, to) }
+}
+
+pub fn fork_current(entry: extern "C" fn(*const ForkedImage, usize) -> !, arg: usize) {
+    internal::executor::fork_current(entry, arg)
+}
+
+pub fn run_on_stack_raw(
+    stack_top: usize,
+    entry: extern "C" fn(usize, usize) -> !,
+    arg: usize,
+) -> ! {
+    internal::executor::run_on_stack_raw(stack_top, entry, arg)
+}

kernel/src/arch/mod.rs

@@ -22,6 +22,7 @@ mod internal {
 }
 
 pub mod cpu;
+pub mod executor;
 pub mod irq;
 pub mod sched;
 pub mod virt;

kernel/src/arch/riscv64/executor.rs

@@ -0,0 +1,38 @@
+use super::sched;
+use crate::{memory::stack::KernelStack, posix::errno::EResult, process::task::Task};
+
+pub type Executor = sched::TaskContext;
+
+#[repr(C)]
+pub struct ForkedImage;
+
+pub struct FrameImage;
+
+impl Executor {
+    pub fn new_kernel(
+        entry: extern "C" fn(usize, usize),
+        arg1: usize,
+        arg2: usize,
+        stack: &KernelStack,
+    ) -> EResult<Self> {
+        let mut this = Self::default();
+        sched::init_task(&mut this, entry, arg1, arg2, stack, false)?;
+        Ok(this)
+    }
+}
+
+pub unsafe fn switch(from: *const Task, to: *const Task) -> *mut Task {
+    unsafe { sched::switch(from, to) }
+}
+
+pub fn fork_current(_entry: extern "C" fn(*const ForkedImage, usize) -> !, _arg: usize) {
+    todo!("riscv64 executor forking is not implemented yet")
+}
+
+pub fn run_on_stack_raw(
+    stack_top: usize,
+    entry: extern "C" fn(usize, usize) -> !,
+    arg: usize,
+) -> ! {
+    sched::run_on_stack_raw(stack_top, entry, arg)
+}

kernel/src/arch/riscv64/mod.rs

@@ -1,4 +1,5 @@
 pub mod cpu;
+pub mod executor;
 pub mod irq;
 pub mod sched;
 pub mod system;

kernel/src/arch/sched.rs

@@ -7,15 +7,15 @@ use crate::{
 use core::{fmt::Debug, mem::MaybeUninit};
 
 pub use internal::sched::{Context, SyscallRestart};
-assert_trait_impl!(TaskContext, Debug);
+assert_trait_impl!(Context, Debug);
 assert_trait_impl!(Context, Default);
 assert_trait_impl!(Context, Clone);
 assert_trait_impl!(Context, Copy);
 
-pub use internal::sched::TaskContext;
-assert_trait_impl!(TaskContext, Debug);
-assert_trait_impl!(TaskContext, Default);
-assert_trait_impl!(TaskContext, Clone);
+pub use crate::arch::executor::Executor;
+assert_trait_impl!(Executor, Debug);
+assert_trait_impl!(Executor, Default);
+assert_trait_impl!(Executor, Clone);
 
 /// Disables preemption.
 /// # Safety
@@ -55,7 +55,7 @@ pub fn broadcast_shootdown() {
 
 /// Initializes a new task.
 pub fn init_task(
-    task: &mut TaskContext,
+    task: &mut Executor,
     entry: extern "C" fn(usize, usize),
     arg1: usize,
     arg2: usize,

kernel/src/arch/x86_64/executor.rs

@@ -0,0 +1,156 @@
+use super::{sched, sched::Context};
+use crate::{memory::stack::KernelStack, posix::errno::EResult, process::task::Task};
+use alloc::boxed::Box;
+use core::{arch::asm, mem::offset_of};
+
+#[repr(C)]
+#[derive(Default, Debug, Clone)]
+pub struct Executor {
+    pub rsp: u64,
+    pub fpu_region: Box<[u8]>,
+    pub ds: u16,
+    pub es: u16,
+    pub fs: u16,
+    pub gs: u16,
+    pub fsbase: u64,
+    pub gsbase: u64,
+    pub restarted: bool,
+}
+
+impl Executor {
+    pub fn new_kernel(
+        entry: extern "C" fn(usize, usize),
+        arg1: usize,
+        arg2: usize,
+        stack: &KernelStack,
+    ) -> EResult<Self> {
+        let mut this = Self::default();
+        sched::init_task(&mut this, entry, arg1, arg2, stack, false)?;
+        Ok(this)
+    }
+}
+
+#[repr(C)]
+pub struct ForkedImage {
+    rbx: usize,
+    rbp: usize,
+    r12: usize,
+    r13: usize,
+    r14: usize,
+    r15: usize,
+    rsp: usize,
+    rip: usize,
+    rflags: usize,
+}
+
+impl ForkedImage {
+    pub const fn ip(&self) -> usize {
+        self.rip
+    }
+
+    pub const fn sp(&self) -> usize {
+        self.rsp
+    }
+
+    pub const fn flags(&self) -> usize {
+        self.rflags
+    }
+}
+
+pub struct FrameImage<'a> {
+    frame: &'a mut Context,
+}
+
+impl<'a> FrameImage<'a> {
+    pub fn new(frame: &'a mut Context) -> Self {
+        Self { frame }
+    }
+
+    pub const fn frame(&self) -> &Context {
+        self.frame
+    }
+
+    pub fn frame_mut(&mut self) -> &mut Context {
+        self.frame
+    }
+
+    pub const fn ip(&self) -> usize {
+        self.frame.rip as usize
+    }
+
+    pub const fn sp(&self) -> usize {
+        self.frame.rsp as usize
+    }
+
+    pub const fn flags(&self) -> usize {
+        self.frame.rflags as usize
+    }
+}
+
+pub unsafe fn switch(from: *const Task, to: *const Task) -> *mut Task {
+    unsafe { sched::switch(from, to) }
+}
+
+pub fn run_on_stack_raw(
+    stack_top: usize,
+    entry: extern "C" fn(usize, usize) -> !,
+    arg: usize,
+) -> ! {
+    sched::run_on_stack_raw(stack_top, entry, arg)
+}
+
+pub fn fork_current(entry: extern "C" fn(*const ForkedImage, usize) -> !, arg: usize) {
+    let mut image = ForkedImage {
+        rbx: 0,
+        rbp: 0,
+        r12: 0,
+        r13: 0,
+        r14: 0,
+        r15: 0,
+        rsp: 0,
+        rip: 0,
+        rflags: 0,
+    };
+
+    unsafe {
+        asm!(
+            "mov [{image} + {rbx}], rbx",
+            "mov [{image} + {rbp}], rbp",
+            "mov [{image} + {r12}], r12",
+            "mov [{image} + {r13}], r13",
+            "mov [{image} + {r14}], r14",
+            "mov [{image} + {r15}], r15",
+            "mov [{image} + {rsp}], rsp",
+            "lea rax, [rip + 2f]",
+            "mov [{image} + {rip}], rax",
+            "pushfq",
+            "pop [{image} + {rflags}]",
+            "mov rdi, {image}",
+            "mov rsi, {arg}",
+            "call {entry}",
+            "ud2",
+            "2:",
+            image = in(reg) &mut image,
+            arg = in(reg) arg,
+            entry = in(reg) entry,
+            rbx = const offset_of!(ForkedImage, rbx),
+            rbp = const offset_of!(ForkedImage, rbp),
+            r12 = const offset_of!(ForkedImage, r12),
+            r13 = const offset_of!(ForkedImage, r13),
+            r14 = const offset_of!(ForkedImage, r14),
+            r15 = const offset_of!(ForkedImage, r15),
+            rsp = const offset_of!(ForkedImage, rsp),
+            rip = const offset_of!(ForkedImage, rip),
+            rflags = const offset_of!(ForkedImage, rflags),
+            out("rax") _,
+            out("rdi") _,
+            out("rsi") _,
+            out("rdx") _,
+            out("rcx") _,
+            out("r8") _,
+            out("r9") _,
+            out("r10") _,
+            out("r11") _,
+        );
+    }
+}

kernel/src/arch/x86_64/mod.rs

@@ -3,6 +3,7 @@ use crate::util::once::Once;
 mod asm;
 mod consts;
 pub mod cpu;
+pub mod executor;
 pub mod irq;
 pub mod sched;
 pub mod system;

kernel/src/arch/x86_64/sched.rs

@@ -7,6 +7,7 @@ use super::{
     system::{apic::LAPIC, gdt::TSS},
 };
 use crate::{
+    arch::sched::Executor,
     irq::lock::IrqLock,
     memory::{UserPtr, VirtAddr, stack::KernelStack},
     percpu::CpuData,
@@ -19,27 +20,12 @@ use crate::{
     sched::Scheduler,
     uapi::signal::{SA_ONSTACK, siginfo_t},
 };
-use alloc::boxed::Box;
 use core::{
     arch::{asm, naked_asm},
     fmt::Write,
     mem::offset_of,
 };
 
-#[repr(C)]
-#[derive(Default, Debug, Clone)]
-pub struct TaskContext {
-    pub rsp: u64,
-    pub fpu_region: Box<[u8]>,
-    pub ds: u16,
-    pub es: u16,
-    pub fs: u16,
-    pub gs: u16,
-    pub fsbase: u64,
-    pub gsbase: u64,
-    pub restarted: bool,
-}
-
 #[repr(C)]
 #[derive(Default, Clone, Copy)]
 pub struct Context {
@@ -187,8 +173,8 @@ pub(in crate::arch) unsafe fn switch(from: *const Task, to: *const Task) -> *mut
         let from = from.as_ref().unwrap();
         let to = to.as_ref().unwrap();
 
-        let from_context = &mut *from.task_context.get();
-        let to_context = &mut *to.task_context.get();
+        let from_context = &mut *from.executor.get();
+        let to_context = &mut *to.executor.get();
 
         let cpu = ARCH_DATA.get();
         TSS.get().lock().rsp0 = to.kernel_stack.top().into();
@@ -279,7 +265,7 @@ pub(in crate::arch) extern "C" fn run_on_stack_raw(
 }
 
 pub(in crate::arch) fn init_task(
-    context: &mut TaskContext,
+    context: &mut Executor,
     entry: extern "C" fn(usize, usize),
     arg1: usize,
     arg2: usize,

kernel/src/arch/x86_64/system/smp.rs

@@ -219,6 +219,7 @@ extern "C" fn ap_entry(info: PhysAddr) -> ! {
 
     let reaper_task =
         Arc::new(Task::new(sched::reaper_fn, 0, 0, Process::get_kernel(), false).unwrap());
+    reaper_task.bound.store(true, Ordering::Release);
     cpu_ctx.scheduler.set_reaper_task(reaper_task);
 
     // Create a dummy task to drop right after the first reschedule.