From a2c5ba1494489514c676e22850b06e41865f1b27 Mon Sep 17 00:00:00 2001
From: Paul Pan <i@0x7f.app>
Date: Wed, 3 Jul 2024 22:51:21 +0800
Subject: [PATCH] feat: add yield

---
 kernel/src/entry.rs       |  1 +
 kernel/src/objects/tcb.rs | 13 ++++++++++---
 kernel/src/scheduler.rs   |  8 +++-----
 kernel/src/syscall.rs     |  5 +++++
 root/root.S               | 10 +++++++---
 5 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/kernel/src/entry.rs b/kernel/src/entry.rs
index 9e172fb..a528b10 100644
--- a/kernel/src/entry.rs
+++ b/kernel/src/entry.rs
@@ -40,6 +40,7 @@ pub fn rust_main() -> ! {
 
     setup_root_server(&fdt);
 
+    info!("Kernel Initialization Completed");
     SCHEDULER.schedule();
 
     error!("[rust_main] Should not reach here! Maybe scheduler is not working?");
diff --git a/kernel/src/objects/tcb.rs b/kernel/src/objects/tcb.rs
index d8adb4e..04eda0c 100644
--- a/kernel/src/objects/tcb.rs
+++ b/kernel/src/objects/tcb.rs
@@ -12,6 +12,9 @@ use utils::{
 
 pub static TID_ALLOCATOR: AtomicUsize = AtomicUsize::new(0);
 
+// one process can run for 10 ticks
+pub const TIME_TICK: usize = 10;
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub enum ThreadState {
     Inactive,
@@ -103,8 +106,12 @@ impl TcbObject {
         self.time_tick
     }
 
-    pub fn set_timetick(&mut self, timeslice: usize) {
-        self.time_tick = timeslice;
+    pub fn refill_timetick(&mut self) {
+        self.time_tick = if self.tid != 0 { TIME_TICK } else { 1 };
+    }
+
+    pub fn clear_timetick(&mut self) {
+        self.time_tick = 0;
     }
 
     pub fn do_tick(&mut self) {
@@ -124,7 +131,7 @@ impl TcbObject {
     }
 
     pub fn handle_syscall(&mut self) {
-        // TODO: don't panic here, return error to user instead
+        // TODO: don't panic here, return error to user instead, create SysRespInfo
         handle_syscall(self).unwrap();
     }
 }
diff --git a/kernel/src/scheduler.rs b/kernel/src/scheduler.rs
index 94586ec..a0fb0d5 100644
--- a/kernel/src/scheduler.rs
+++ b/kernel/src/scheduler.rs
@@ -1,6 +1,6 @@
 use crate::objects::*;
 use core::ptr::NonNull;
-use log::error;
+use log::{error, trace};
 use spin::lazy::Lazy;
 use utils::{container_of, linked_list::Link};
 
@@ -16,9 +16,6 @@ pub static SCHEDULER: Scheduler = Scheduler::new();
 
 // TODO: add a shared buffer to transfer TCB between cores
 
-// one process can run for 10 ticks
-pub const TIME_SLICE: usize = 10;
-
 pub struct Scheduler {
     head: Link<TcbObject>,
 }
@@ -41,9 +38,10 @@ impl Scheduler {
     pub fn schedule(&self) {
         while let Some(next) = self.head.next_mut() {
             if next.timetick() > 0 && next.schedulable() {
+                trace!("Scheduling thread {}", next.tid());
                 next.activate();
             } else if next.timetick() == 0 {
-                next.set_timetick(TIME_SLICE);
+                next.refill_timetick();
             }
 
             // put to the end of the queue
diff --git a/kernel/src/syscall.rs b/kernel/src/syscall.rs
index 8c6c43f..a8196c6 100644
--- a/kernel/src/syscall.rs
+++ b/kernel/src/syscall.rs
@@ -14,6 +14,7 @@ pub fn handle_syscall(tcb: &mut TcbObject) -> SysResult {
     trace!("handle_syscall: info={:?} cap={:?}", info, cap);
 
     match info.label() {
+        Syscall::Invalid => Ok(()),
         Syscall::DebugPutChar => {
             if info.length() != 3 {
                 return Err(SysError::InvalidArgument);
@@ -24,6 +25,10 @@ pub fn handle_syscall(tcb: &mut TcbObject) -> SysResult {
 
             Ok(())
         },
+        Syscall::Yield => {
+            tcb.clear_timetick();
+            Ok(())
+        },
         _ => unimplemented!("handle_syscall: {:?}, {:?}", info.label(), cap),
     }
 }
diff --git a/root/root.S b/root/root.S
index 452adef..9f61c8a 100644
--- a/root/root.S
+++ b/root/root.S
@@ -10,13 +10,17 @@ _start:
 	la	t0, msg
 .loop:
 	lb	a2, 0(t0)
-	beqz	a2, _start
-	# see uapi/syscall.rs
-	# msg: label = DebugPutChar, length = 3
+	beqz	a2, .yield
+	# msg: label = DebugPutChar(1), length = 3
 	# ptr: self-referential to cspace, just to create a valid cptr
 	syscall (1<<12 | 3), 0x0202020202020202
 	addi	t0, t0, 1
 	j	.loop
+.yield:
+	# msg: label = Yield(2), length = 2
+	# ptr: self-referential to cspace, just to create a valid cptr
+	syscall (2<<12 | 2), 0x0202020202020202
+	j	_start
 
 .data
 	msg : .string "Hello, world!\n\0"