working on thread, assembly is trolling me somehow

Signed-off-by: Julien CLEMENT <julien.clement@epita.fr>

src/lib.rs

@@ -48,6 +48,7 @@ pub fn init(boot_info: &BootInformation) {
     memory::gdt::init_gdt();
     interrupts::init_idt();
     vga::change_color(ColorCode::new(Color::LightGreen, Color::Black));
+    println!("Init kernel main thread: {:?}", proc::thread::KERNEL_THREAD.try_lock().unwrap().id);
 }
 
 #[no_mangle]
@@ -84,4 +85,7 @@ async fn get_file() {
     serial_println!("{}", alloc::str::from_utf8(&buf).unwrap());
 
     fd.borrow_mut().close().await;
+
+    let mut thread = proc::thread::Thread::new();
+    thread.start(proc::thread::exit as u64).await;
 }

src/proc/scheduler/mod.rs

@@ -1,7 +1,31 @@
-use super::thread::Thread;
+use super::thread::{Thread, ThreadId};
 
-use alloc::vec::Vec;
+use alloc::{collections::BTreeMap, sync::Arc};
+use crossbeam_queue::ArrayQueue;
 
 pub struct Scheduler {
-    threads: Vec<Thread>,
+    threads: BTreeMap<ThreadId, Arc<Thread>>,
+    thread_queue: Arc<ArrayQueue<ThreadId>>,
+}
+
+impl Scheduler {
+    pub fn new() -> Self {
+        Scheduler {
+            threads: BTreeMap::new(),
+            thread_queue: Arc::new(ArrayQueue::new(100)),
+        }
+    }
+
+    pub fn schedule(&mut self) -> Option<Arc<Thread>> {
+        if let Ok(thread_id) = self.thread_queue.pop() {
+            self.thread_queue.push(thread_id);
+            let thread = match self.threads.get_mut(&thread_id) {
+                Some(thread) => thread,
+                None => return None,
+            };
+            Some(thread.clone())
+        } else {
+            None
+        }
+    }
 }

src/proc/thread/mod.rs

@@ -1,8 +1,50 @@
+use crate::println;
+use crate::utils::mutex::AsyncMutex;
+
+use core::arch::asm;
+use core::sync::atomic::{AtomicU64, Ordering};
+
 use alloc::alloc::{alloc, dealloc, Layout};
+use lazy_static::lazy_static;
 
 const STACK_SIZE: usize = 4096 * 20;
 
+lazy_static! {
+    pub static ref RUNNING_THREAD: AsyncMutex<ThreadId> = AsyncMutex::new(ThreadId(0));
+    pub static ref KERNEL_THREAD: AsyncMutex<Thread> = {
+        let k_rsp: u64;
+        unsafe {
+            asm!(
+                "push rsp",    // Recover current rsp
+                "pop {out}",
+                out = out(reg) k_rsp, // Save current rsp
+            );
+        }
+        let thread: Thread = Thread {
+            id: ThreadId(0),
+            rsp: k_rsp,
+        };
+        AsyncMutex::new(thread)
+    };
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+pub struct ThreadId(u64);
+
+impl ThreadId {
+    fn new() -> Self {
+        static NEXT_ID: AtomicU64 = AtomicU64::new(1);
+        ThreadId(NEXT_ID.fetch_add(1, Ordering::Relaxed))
+    }
+}
+
+pub fn exit() {
+    println!("Exiting");
+    KERNEL_THREAD.try_lock().unwrap().run();
+}
+
 pub struct Thread {
+    pub id: ThreadId,
     rsp: u64
 }
 
@@ -10,8 +52,51 @@ impl Thread {
     pub fn new() -> Self {
         unsafe {
             Thread {
+                id: ThreadId::new(),
                 rsp: alloc(Layout::new::<[u8; STACK_SIZE]>()) as u64,
             }
         }
     }
+
+    pub async fn start(&mut self, rip: u64) {
+        unsafe {
+            asm!(
+                "pusha",       // Save current thread regs
+                "push rsp",    // Recover current rsp
+                "pop {out}",
+                out = out(reg) self.rsp, // Save current rsp
+            );
+        }
+
+        *RUNNING_THREAD.lock().await = self.id;
+        unsafe {
+            asm!(
+                "push {rsp}",
+                "pop rsp",
+                "jmp {rip}",
+                rsp = in(reg) self.rsp,
+                rip = in(reg) rip,
+            );
+        }
+    }
+
+    pub async fn run(&mut self) {
+        unsafe {
+            asm!(
+                "pusha",       // Save current thread regs
+                "push rsp",    // Recover current rsp
+                "pop {out}",
+                out = out(reg) self.rsp, // Save current rsp
+            );
+
+            *RUNNING_THREAD.lock().await = self.id; // change running thread
+
+            asm!(
+                "push {rsp}", // Set stack pointer to the new thread
+                "pop rsp",
+                "popa",       // Restore new thread regs
+                rsp = in(reg) self.rsp,
+            );
+        }
+    }
 }

src/utils/mutex.rs

@@ -73,7 +73,7 @@ impl<T> AsyncMutex<T> {
     }
 
     pub fn try_lock(&self) -> Option<AsyncMutexGuard<T>> {
-        if self.lock.try_lock() {
+        if !self.lock.try_lock() {
             Some(AsyncMutexGuard { mutex: self })
         } else {
             None