From fba4a9405ecedd32d06446678b26520f8a37906c Mon Sep 17 00:00:00 2001
From: Chuandong Li <lichuandong.lcd@antgroup.com>
Date: Mon, 3 Jul 2023 00:29:59 +0800
Subject: [PATCH] Support the wait queue that allows sleeping

---
 framework/jinux-frame/src/sync/wait.rs      | 10 +++++++---
 framework/jinux-frame/src/task/mod.rs       |  4 ++--
 framework/jinux-frame/src/task/processor.rs |  6 ++++--
 framework/jinux-frame/src/task/scheduler.rs | 12 ++++++------
 framework/jinux-frame/src/task/task.rs      |  2 --
 5 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/framework/jinux-frame/src/sync/wait.rs b/framework/jinux-frame/src/sync/wait.rs
index c34a45c3b..cf5ebfa58 100644
--- a/framework/jinux-frame/src/sync/wait.rs
+++ b/framework/jinux-frame/src/sync/wait.rs
@@ -4,7 +4,7 @@ use super::SpinLock;
 use alloc::{collections::VecDeque, sync::Arc};
 use bitflags::bitflags;
 
-use crate::task::schedule;
+use crate::task::{add_task, current_task, schedule, Task, TaskStatus};
 
 /// A wait queue.
 ///
@@ -85,7 +85,6 @@ impl WaitQueue {
     }
 }
 
-#[derive(Debug)]
 struct Waiter {
     /// Whether the
     is_woken_up: AtomicBool,
@@ -93,6 +92,8 @@ struct Waiter {
     flag: WaiterFlag,
     /// if the wait condition is ture, then the waiter is finished and can be removed from waitqueue
     wait_finished: AtomicBool,
+    /// The `Task` held by the waiter.
+    task: Arc<Task>,
 }
 
 impl Waiter {
@@ -101,16 +102,18 @@ impl Waiter {
             is_woken_up: AtomicBool::new(false),
             flag: WaiterFlag::empty(),
             wait_finished: AtomicBool::new(false),
+            task: current_task().unwrap(),
         }
     }
 
     /// make self into wait status until be called wake up
     pub fn wait(&self) {
         self.is_woken_up.store(false, Ordering::SeqCst);
+        self.task.inner_exclusive_access().task_status = TaskStatus::Sleeping;
         while !self.is_woken_up.load(Ordering::SeqCst) {
-            // yield the execution, to allow other task to continue
             schedule();
         }
+        self.task.inner_exclusive_access().task_status = TaskStatus::Runnable;
     }
 
     pub fn is_woken_up(&self) -> bool {
@@ -119,6 +122,7 @@ impl Waiter {
 
     pub fn wake_up(&self) {
         self.is_woken_up.store(true, Ordering::SeqCst);
+        add_task(self.task.clone());
     }
 
     pub fn set_finished(&self) {
diff --git a/framework/jinux-frame/src/task/mod.rs b/framework/jinux-frame/src/task/mod.rs
index 7cf2fde57..5548defac 100644
--- a/framework/jinux-frame/src/task/mod.rs
+++ b/framework/jinux-frame/src/task/mod.rs
@@ -5,6 +5,6 @@ mod scheduler;
 #[allow(clippy::module_inception)]
 mod task;
 
-pub use self::processor::schedule;
-pub use self::scheduler::{set_scheduler, Scheduler};
+pub use self::processor::{current_task, schedule};
+pub use self::scheduler::{add_task, set_scheduler, Scheduler};
 pub use self::task::{Task, TaskAdapter, TaskStatus};
diff --git a/framework/jinux-frame/src/task/processor.rs b/framework/jinux-frame/src/task/processor.rs
index 2e7b93d3d..e2892d9bc 100644
--- a/framework/jinux-frame/src/task/processor.rs
+++ b/framework/jinux-frame/src/task/processor.rs
@@ -71,8 +71,10 @@ pub fn switch_to_task(next_task: Arc<Task>) {
         PROCESSOR.lock().get_idle_task_cx_ptr()
     } else {
         current_task = current_task_option.unwrap();
-        if current_task.status() != TaskStatus::Exited {
-            GLOBAL_SCHEDULER.lock().enqueue(current_task.clone());
+        if current_task.status() == TaskStatus::Runnable {
+            GLOBAL_SCHEDULER
+                .lock_irq_disabled()
+                .enqueue(current_task.clone());
         }
         &mut current_task.inner_exclusive_access().ctx as *mut TaskContext
     };
diff --git a/framework/jinux-frame/src/task/scheduler.rs b/framework/jinux-frame/src/task/scheduler.rs
index 52b949798..78c07a988 100644
--- a/framework/jinux-frame/src/task/scheduler.rs
+++ b/framework/jinux-frame/src/task/scheduler.rs
@@ -1,12 +1,12 @@
 use crate::prelude::*;
+use crate::sync::SpinLock;
 use crate::task::Task;
 
 use lazy_static::lazy_static;
-use spin::Mutex;
 
 lazy_static! {
-    pub(crate) static ref GLOBAL_SCHEDULER: Mutex<GlobalScheduler> =
-        Mutex::new(GlobalScheduler { scheduler: None });
+    pub(crate) static ref GLOBAL_SCHEDULER: SpinLock<GlobalScheduler> =
+        SpinLock::new(GlobalScheduler { scheduler: None });
 }
 
 /// A scheduler for tasks.
@@ -43,13 +43,13 @@ impl GlobalScheduler {
 ///
 /// This must be called before invoking `Task::spawn`.
 pub fn set_scheduler(scheduler: &'static dyn Scheduler) {
-    GLOBAL_SCHEDULER.lock().scheduler = Some(scheduler);
+    GLOBAL_SCHEDULER.lock_irq_disabled().scheduler = Some(scheduler);
 }
 
 pub fn fetch_task() -> Option<Arc<Task>> {
-    GLOBAL_SCHEDULER.lock().dequeue()
+    GLOBAL_SCHEDULER.lock_irq_disabled().dequeue()
 }
 
 pub fn add_task(task: Arc<Task>) {
-    GLOBAL_SCHEDULER.lock().enqueue(task);
+    GLOBAL_SCHEDULER.lock_irq_disabled().enqueue(task);
 }
diff --git a/framework/jinux-frame/src/task/task.rs b/framework/jinux-frame/src/task/task.rs
index 0e5ec39e6..1f79989f7 100644
--- a/framework/jinux-frame/src/task/task.rs
+++ b/framework/jinux-frame/src/task/task.rs
@@ -216,8 +216,6 @@ impl Task {
 pub enum TaskStatus {
     /// The task is runnable.
     Runnable,
-    /// The task is running.
-    Running,
     /// The task is sleeping.
     Sleeping,
     /// The task has exited.