Refactor threading code

src/threads.rs

@@ -2,11 +2,137 @@ use alloc::collections::VecDeque;
 use x86_64::structures::paging::{FrameAllocator, Mapper, Size4KiB};
 use x86_64::VirtAddr;
 
-pub unsafe fn context_switch(stack_pointer: VirtAddr) {
+static SCHEDULER: spin::Mutex<Option<Scheduler>> = spin::Mutex::new(None);
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct ThreadId(pub u64);
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct StackPointer(VirtAddr);
+
+impl StackPointer {
+    unsafe fn new(pointer: VirtAddr) -> Self {
+        StackPointer(pointer)
+    }
+
+    pub fn allocate(
+        stack_size: u64,
+        mapper: &mut impl Mapper<Size4KiB>,
+        frame_allocator: &mut impl FrameAllocator<Size4KiB>,
+    ) -> Result<Self, ()> {
+        crate::memory::alloc_stack(stack_size, mapper, frame_allocator).map(Self)
+    }
+
+    unsafe fn push_to_stack<T>(&mut self, value: T) {
+        self.0 -= core::mem::size_of::<T>();
+        let ptr: *mut T = self.0.as_mut_ptr();
+        ptr.write(value);
+    }
+
+    fn as_u64(&self) -> u64 {
+        self.0.as_u64()
+    }
+}
+
+pub struct Thread {
+    id: ThreadId,
+    stack_pointer: StackPointer,
+}
+
+impl Thread {
+    pub unsafe fn new(entry_point: fn() -> !, stack_top: StackPointer) -> Self {
+        use core::sync::atomic::{AtomicU64, Ordering};
+        static NextThreadId: AtomicU64 = AtomicU64::new(1);
+
+        let mut stack_pointer = stack_top;
+        Self::set_up_stack(&mut stack_pointer, entry_point);
+
+        Thread {
+            id: ThreadId(NextThreadId.fetch_add(1, Ordering::SeqCst)),
+            stack_pointer,
+        }
+    }
+
+    pub unsafe fn new_from_closure<F>(closure: F, stack_top: StackPointer) -> Self
+    where
+        F: FnOnce() -> ! + Send + Sync + 'static,
+    {
+        use alloc::boxed::Box;
+        use core::{mem, raw::TraitObject};
+
+        let boxed: ThreadClosure = Box::new(closure);
+        let trait_object: TraitObject = unsafe { mem::transmute(boxed) };
+
+        // push trait object
+        let mut stack_pointer = stack_top;
+        unsafe { stack_pointer.push_to_stack(trait_object.data) };
+        unsafe { stack_pointer.push_to_stack(trait_object.vtable) };
+
+        let entry_point = call_closure_entry as unsafe fn() -> !;
+        unsafe { Self::new(mem::transmute(entry_point), stack_pointer) }
+    }
+
+    pub fn create(
+        entry_point: fn() -> !,
+        stack_size: u64,
+        mapper: &mut impl Mapper<Size4KiB>,
+        frame_allocator: &mut impl FrameAllocator<Size4KiB>,
+    ) -> Result<Self, ()> {
+        let stack_top = StackPointer::allocate(stack_size, mapper, frame_allocator)?;
+        Ok(unsafe { Self::new(entry_point, stack_top) })
+    }
+
+    pub fn create_from_closure<F>(
+        entry_point: F,
+        stack_size: u64,
+        mapper: &mut impl Mapper<Size4KiB>,
+        frame_allocator: &mut impl FrameAllocator<Size4KiB>,
+    ) -> Result<Self, ()>
+    where
+        F: FnOnce() -> ! + Send + Sync + 'static,
+    {
+        let stack_top = StackPointer::allocate(stack_size, mapper, frame_allocator)?;
+        Ok(unsafe { Self::new_from_closure(entry_point, stack_top) })
+    }
+
+    fn set_up_stack(stack_top: &mut StackPointer, entry_point: fn() -> !) {
+        unsafe { stack_top.push_to_stack(entry_point) };
+        let rflags: u64 = 0x200;
+        unsafe { stack_top.push_to_stack(rflags) };
+    }
+}
+
+struct Scheduler {
+    current_thread_id: ThreadId,
+    paused_threads: VecDeque<Thread>,
+}
+
+impl Scheduler {
+    fn new() -> Self {
+        Scheduler {
+            current_thread_id: ThreadId(0),
+            paused_threads: VecDeque::new(),
+        }
+    }
+
+    fn next_thread(&mut self) -> Option<Thread> {
+        self.paused_threads.pop_front()
+    }
+
+    fn add_paused_thread(&mut self, thread: Thread) {
+        self.paused_threads.push_back(thread);
+    }
+
+    fn add_new_thread(&mut self, thread: Thread) {
+        self.add_paused_thread(thread);
+    }
+}
+
+pub unsafe fn context_switch(thread: Thread) {
     asm!(
         "call asm_context_switch"
         :
-        : "{rdi}"(stack_pointer.as_u64())
+        : "{rdi}"(thread.stack_pointer.as_u64()), "{rsi}"(thread.id.0)
         : "rax", "rbx", "rcx", "rdx", "rsi", "rdi", "rpb", "r8", "r9", "r10",
         "r11", "r12", "r13", "r14", "r15", "rflags", "memory"
         : "intel", "volatile"
@@ -17,6 +143,7 @@ global_asm!(
     "
     .intel_syntax noprefix
 
+    // asm_context_switch(stack_pointer: u64, thread_id: u64)
     asm_context_switch:
         pushfq
 
@@ -32,11 +159,9 @@ global_asm!(
 );
 
 pub fn scheduler() {
-    let next = PAUSED_THREADS.try_lock().and_then(|mut paused_threads| {
-        paused_threads
-            .as_mut()
-            .and_then(|threads| threads.pop_front())
-    });
+    let next = SCHEDULER
+        .try_lock()
+        .and_then(|mut scheduler| scheduler.as_mut().and_then(|s| s.next_thread()));
     if let Some(next) = next {
         unsafe { context_switch(next) };
     }
@@ -45,66 +170,37 @@ pub fn scheduler() {
 static PAUSED_THREADS: spin::Mutex<Option<VecDeque<VirtAddr>>> = spin::Mutex::new(None);
 
 #[no_mangle]
-fn add_paused_thread(stack_pointer: VirtAddr) {
-    add_thread(stack_pointer)
-}
+pub extern "C" fn add_paused_thread(stack_pointer: u64, thread_id: u64) {
+    let thread = Thread {
+        stack_pointer: StackPointer(VirtAddr::new(stack_pointer)),
+        id: ThreadId(thread_id),
+    };
 
-fn add_thread(stack_pointer: VirtAddr) {
-    PAUSED_THREADS
-        .lock()
-        .get_or_insert_with(VecDeque::new)
-        .push_back(stack_pointer);
+    SCHEDULER.lock().get_or_insert_with(Scheduler::new).add_paused_thread(thread);
 }
 
 pub fn create_thread(
-    f: fn() -> !,
+    entry_point: fn() -> !,
     stack_size: u64,
     mapper: &mut impl Mapper<Size4KiB>,
     frame_allocator: &mut impl FrameAllocator<Size4KiB>,
-) {
-    let mut stack = crate::memory::alloc_stack(stack_size, mapper, frame_allocator).unwrap();
-    stack -= core::mem::size_of::<u64>();
-    let ptr: *mut u64 = stack.as_mut_ptr();
-    unsafe { ptr.write(f as u64) };
-    stack -= core::mem::size_of::<u64>();
-    let ptr: *mut u64 = stack.as_mut_ptr();
-    let rflags = 0x200;
-    unsafe { ptr.write(rflags) };
-    add_thread(stack);
+) -> Result<(), ()> {
+    let thread = Thread::create(entry_point, stack_size, mapper, frame_allocator)?;
+    SCHEDULER.lock().get_or_insert_with(Scheduler::new).add_new_thread(thread);
+    Ok(())
 }
 
 pub fn create_thread_from_closure<F>(
-    f: F,
+    closure: F,
     stack_size: u64,
     mapper: &mut impl Mapper<Size4KiB>,
     frame_allocator: &mut impl FrameAllocator<Size4KiB>,
-) where
-    F: FnOnce() -> ! + 'static,
+) -> Result<(), ()> where
+    F: FnOnce() -> ! + 'static + Send + Sync,
 {
-    use alloc::boxed::Box;
-    use core::{mem, raw::TraitObject};
-
-    let boxed: ThreadClosure = Box::new(f);
-    let trait_object: TraitObject = unsafe { mem::transmute(boxed) };
-
-    let mut stack = crate::memory::alloc_stack(stack_size, mapper, frame_allocator).unwrap();
-
-    // push trait object
-    stack -= core::mem::size_of::<*mut ()>();
-    let ptr: *mut *mut () = stack.as_mut_ptr();
-    unsafe { ptr.write(trait_object.data) };
-    stack -= core::mem::size_of::<*mut ()>();
-    let ptr: *mut *mut () = stack.as_mut_ptr();
-    unsafe { ptr.write(trait_object.vtable) };
-
-    stack -= core::mem::size_of::<u64>();
-    let ptr: *mut u64 = stack.as_mut_ptr();
-    unsafe { ptr.write(call_closure_entry as u64) };
-    stack -= core::mem::size_of::<u64>();
-    let ptr: *mut u64 = stack.as_mut_ptr();
-    let rflags = 0x200;
-    unsafe { ptr.write(rflags) };
-    add_thread(stack);
+    let thread = Thread::create_from_closure(closure, stack_size, mapper, frame_allocator)?;
+    SCHEDULER.lock().get_or_insert_with(Scheduler::new).add_new_thread(thread);
+    Ok(())
 }
 
 type ThreadClosure = alloc::boxed::Box<dyn FnOnce() -> !>;