diff --git a/src/proc/scheduler/mod.rs b/src/proc/scheduler/mod.rs index 17e4084..1c7ec52 100644 --- a/src/proc/scheduler/mod.rs +++ b/src/proc/scheduler/mod.rs @@ -31,6 +31,13 @@ impl ThreadStream { waker: AtomicWaker::new(), } } + + pub fn register(&mut self, id: ThreadId) { + self.ids + .push(id) + .expect("Thread queue full"); + self.waker.wake(); + } } impl Stream for ThreadStream { @@ -69,6 +76,7 @@ impl Scheduler { entry_point: 0, started: true, rsp: 0, + base_stack: 0, }; res.register(Arc::new(RefCell::new(k_thread))); res @@ -77,9 +85,11 @@ impl Scheduler { pub async fn run(&mut self) { while let Some(id) = self.thread_queue.next().await { - let thread = self.get_thread(id).unwrap(); - unsafe { - (&mut*thread.as_ptr()).run(); + if let Some(thread) = self.get_thread(id) { // Thread still exists + unsafe { + (&mut*thread.as_ptr()).run(); + } + self.thread_queue.register(id); } } } @@ -89,11 +99,13 @@ impl Scheduler { if self.threads.insert(thread_id, thread).is_some() { panic!("Duplicate thread ID") } - self.thread_queue - .ids - .push(thread_id) - .expect("Thread queue full"); - self.thread_queue.waker.wake(); + if thread_id != ThreadId(0) { + self.thread_queue.register(thread_id); + } + } + + pub fn exit(&mut self, id: ThreadId) { + self.threads.remove(&id).unwrap().borrow().exit(); } pub fn get_thread(&mut self, id: ThreadId) -> Option { diff --git a/src/proc/thread/mod.rs b/src/proc/thread/mod.rs index 60f513b..29de36a 100644 --- a/src/proc/thread/mod.rs +++ b/src/proc/thread/mod.rs @@ -34,6 +34,7 @@ pub fn exit() { .get_thread(ThreadId(0)) .unwrap() .as_ptr(); + scheduler.exit(*RUNNING_THREAD.try_lock().unwrap()); } // Drop scheduler mutex guard unsafe { @@ -51,36 +52,62 @@ pub struct Thread { pub entry_point: u64, pub started: bool, pub rsp: u64, + pub base_stack: u64 } impl Thread { pub fn new(entry_point: u64) -> Self { unsafe { + let stack_bottom = alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64; Thread { id: ThreadId::new(), entry_point: entry_point, started: false, - rsp: alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64 + STACK_SIZE as u64, + rsp: stack_bottom + STACK_SIZE as u64, + base_stack: stack_bottom, } } } + pub fn exit(&self) { + unsafe { + dealloc(self.base_stack as *mut u8, Layout::new::<[u8; STACK_SIZE]>()); + } + } + pub fn run(&mut self) { println!("Running thread {:?}", self.id); unsafe { let mut current_thread_guard = RUNNING_THREAD.try_lock().unwrap(); - let current_rsp: u64; - asm!( - "push rsp", // Recover current rsp - "pop {out}", - "sub {out}, 56", // Offset to saved registers - out = out(reg) current_rsp, // Save thread rsp - ); let mut scheduler = SCHEDULER.try_lock().unwrap(); - // TODO: check if the thread still exists - let current_thread = scheduler.get_thread(*current_thread_guard).unwrap(); - current_thread.borrow_mut().rsp = current_rsp; + if let Some(current_thread) = scheduler.get_thread(*current_thread_guard) { + let current_rsp: u64; + asm!( + "push rsp", // Recover current rsp + "pop {out}", + "sub {out}, 56", // Offset to saved registers + out = out(reg) current_rsp, // Save thread rsp + ); + current_thread.borrow_mut().rsp = current_rsp; + } + else { // Thread does not exists anymore + *current_thread_guard = self.id; // change running thread + asm!( // Just switch to new thead without saving registers + "push {rsp}", // Set stack pointer to the new thread + "pop rsp", + + "pop rdi", // Restore new thread regs + "pop rsi", + "pop rbp", + "pop rdx", + "pop rcx", + "pop rbx", + "pop rax", + rsp = in(reg) self.rsp, + ); + return; + } *current_thread_guard = self.id; // change running thread } // The scheduler and running thread guards are dropped here