From 5e2e0b629e6615922020c4b4c6cc9777068191b7 Mon Sep 17 00:00:00 2001 From: Philipp Oppermann Date: Thu, 23 Jan 2020 10:42:37 +0100 Subject: [PATCH] Refactor threading code --- src/threads.rs | 202 ++++++++++++++++++++++++++++++++++++------------- 1 file changed, 149 insertions(+), 53 deletions(-) diff --git a/src/threads.rs b/src/threads.rs index 924a6530..7796bd2e 100644 --- a/src/threads.rs +++ b/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> = 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, + frame_allocator: &mut impl FrameAllocator, + ) -> Result { + crate::memory::alloc_stack(stack_size, mapper, frame_allocator).map(Self) + } + + unsafe fn push_to_stack(&mut self, value: T) { + self.0 -= core::mem::size_of::(); + 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(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, + frame_allocator: &mut impl FrameAllocator, + ) -> Result { + let stack_top = StackPointer::allocate(stack_size, mapper, frame_allocator)?; + Ok(unsafe { Self::new(entry_point, stack_top) }) + } + + pub fn create_from_closure( + entry_point: F, + stack_size: u64, + mapper: &mut impl Mapper, + frame_allocator: &mut impl FrameAllocator, + ) -> Result + 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, +} + +impl Scheduler { + fn new() -> Self { + Scheduler { + current_thread_id: ThreadId(0), + paused_threads: VecDeque::new(), + } + } + + fn next_thread(&mut self) -> Option { + 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>> = 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, frame_allocator: &mut impl FrameAllocator, -) { - let mut stack = crate::memory::alloc_stack(stack_size, mapper, frame_allocator).unwrap(); - stack -= core::mem::size_of::(); - let ptr: *mut u64 = stack.as_mut_ptr(); - unsafe { ptr.write(f as u64) }; - stack -= core::mem::size_of::(); - 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, + closure: F, stack_size: u64, mapper: &mut impl Mapper, frame_allocator: &mut impl FrameAllocator, -) 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::(); - let ptr: *mut u64 = stack.as_mut_ptr(); - unsafe { ptr.write(call_closure_entry as u64) }; - stack -= core::mem::size_of::(); - 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 !>;