Discuss: Blocked resched

Signed-off-by: guoweikang <guoweikang.kernel@gmail.com>

modules/axtask/src/run_queue.rs

@@ -10,7 +10,6 @@ use scheduler::BaseScheduler;
 use axhal::cpu::this_cpu_id;
 
 use crate::task::{CurrentTask, TaskState};
-use crate::wait_queue::{WaitQueueGuard, WaitTaskNode};
 use crate::{AxCpuMask, AxTaskRef, Scheduler, TaskInner, WaitQueue};
 
 macro_rules! percpu_static {
@@ -260,18 +259,27 @@ impl<'a, G: BaseGuard> CurrentRunQueueRef<'a, G> {
         }
     }
 
+    /// Reschedule current task.
+    ///
+    /// Once current is running, need to put RQ again.
+    pub fn reschedule(&mut self) {
+        let curr = &self.current_task;
+        if curr.is_running() {
+            self.inner
+                .put_task_with_state(curr.clone(), TaskState::Running, false);
+        }
+
+        self.inner.resched();
+    }
+
     /// Yield the current task and reschedule.
     /// This function will put the current task into this run queue with `Ready` state,
     /// and reschedule to the next task on this run queue.
     pub fn yield_current(&mut self) {
         let curr = &self.current_task;
         trace!("task yield: {}", curr.id_name());
         assert!(curr.is_running());
-
-        self.inner
-            .put_task_with_state(curr.clone(), TaskState::Running, false);
-
-        self.inner.resched();
+        self.reschedule();
     }
 
     /// Migrate the current task to a new run queue matching its CPU affinity and reschedule.
@@ -360,42 +368,6 @@ impl<'a, G: BaseGuard> CurrentRunQueueRef<'a, G> {
         unreachable!("task exited!");
     }
 
-    /// Block the current task, put current task into the wait queue and reschedule.
-    /// Mark the state of current task as `Blocked`, set the `in_wait_queue` flag as true.
-    /// Note:
-    ///     1. The caller must hold the lock of the wait queue.
-    ///     2. The caller must ensure that the current task is in the running state.
-    ///     3. The caller must ensure that the current task is not the idle task.
-    ///     4. The lock of the wait queue will be released explicitly after current task is pushed into it.
-    pub fn blocked_resched(
-        &mut self,
-        mut wq_guard: WaitQueueGuard,
-        curr_waiter: Arc<WaitTaskNode>,
-    ) {
-        let curr = &self.current_task;
-        assert!(curr.is_running());
-        assert!(!curr.is_idle());
-        // we must not block current task with preemption disabled.
-        // Current expected preempt count is 2.
-        // 1 for `NoPreemptIrqSave`, 1 for wait queue's `SpinNoIrq`.
-        #[cfg(feature = "preempt")]
-        assert!(curr.can_preempt(2));
-
-        // Mark the task as blocked, this has to be done before adding it to the wait queue
-        // while holding the lock of the wait queue.
-        curr.set_state(TaskState::Blocked);
-        wq_guard.push_back(curr_waiter);
-        // Drop the lock of wait queue explictly.
-        drop(wq_guard);
-
-        // Current task's state has been changed to `Blocked` and added to the wait queue.
-        // Note that the state may have been set as `Ready` in `unblock_task()`,
-        // see `unblock_task()` for details.
-
-        debug!("task block: {}", curr.id_name());
-        self.inner.resched();
-    }
-
     #[cfg(feature = "irq")]
     pub fn sleep_until(&mut self, deadline: axhal::time::TimeValue) {
         let curr = &self.current_task;

modules/axtask/src/wait_queue.rs

@@ -5,9 +5,7 @@ use kspin::{SpinNoIrq, SpinNoIrqGuard};
 use linked_list::{def_node, List};
 
 use crate::{current_run_queue, select_run_queue, AxTaskRef};
-
-#[cfg(feature = "irq")]
-use crate::CurrentTask;
+use crate::task::TaskState;
 
 def_node!(WaitTaskNode, AxTaskRef);
 
@@ -55,6 +53,18 @@ impl WaitQueue {
         }
     }
 
+    fn waiter_blocked(&self, waiter: &Node) {
+        let curr = waiter.inner();
+        assert!(curr.is_running());
+        assert!(!curr.is_idle());
+        let mut wq_locked = self.list.lock();
+        debug!("task block: {}", curr.id_name());
+        // Mark the task as blocked, this has to be done before adding it to the wait queue
+        // while holding the lock of the wait queue.
+        curr.set_state(TaskState::Blocked);
+        wq_locked.push_back(waiter.clone());
+    }
+
     /// Cancel events by removing the task from the wait list.
     /// If `from_timer_list` is true, try to remove the task from the timer list.
     fn cancel_events(&self, waiter: &Node, _from_timer_list: bool) {
@@ -82,7 +92,8 @@ impl WaitQueue {
     /// notifies it.
     pub fn wait(&self) {
         declare_current_waiter!(waiter);
-        current_run_queue::<NoPreemptIrqSave>().blocked_resched(self.list.lock(), waiter.clone());
+        self.waiter_blocked(&waiter);
+        current_run_queue::<NoPreemptIrqSave>().reschedule();
         self.cancel_events(&waiter, false);
     }
 
@@ -97,13 +108,12 @@ impl WaitQueue {
     {
         declare_current_waiter!(waiter);
         loop {
-            let mut rq = current_run_queue::<NoPreemptIrqSave>();
-            let wq = self.list.lock();
             if condition() {
                 break;
             }
-            rq.blocked_resched(wq, waiter.clone());
+            self.waiter_blocked(&waiter);
             // Preemption may occur here.
+            current_run_queue::<NoPreemptIrqSave>().reschedule();
         }
 
         self.cancel_events(&waiter, false);
@@ -114,20 +124,16 @@ impl WaitQueue {
     #[cfg(feature = "irq")]
     pub fn wait_timeout(&self, dur: core::time::Duration) -> bool {
         declare_current_waiter!(waiter);
-        let mut rq = current_run_queue::<NoPreemptIrqSave>();
-        let curr = crate::current();
         let deadline = axhal::time::wall_time() + dur;
         debug!(
             "task wait_timeout: {} deadline={:?}",
-            curr.id_name(),
+            waiter.inner().id_name(),
             deadline
         );
-        crate::timers::set_alarm_wakeup(deadline, curr.clone());
-
-        rq.blocked_resched(self.list.lock(), waiter.clone());
-
+        self.waiter_blocked(&waiter);
+        crate::timers::set_alarm_wakeup(deadline, waiter.inner().clone());
+        current_run_queue::<NoPreemptIrqSave>().reschedule();
         let timeout = axhal::time::wall_time() >= deadline;
-
         // Always try to remove the task from the timer list.
         self.cancel_events(&waiter, true);
         timeout
@@ -144,28 +150,27 @@ impl WaitQueue {
         F: Fn() -> bool,
     {
         declare_current_waiter!(waiter);
-        let curr = crate::current();
         let deadline = axhal::time::wall_time() + dur;
         debug!(
             "task wait_timeout: {}, deadline={:?}",
-            curr.id_name(),
+            waiter.inner().id_name(),
             deadline
         );
-        crate::timers::set_alarm_wakeup(deadline, curr.clone());
+        crate::timers::set_alarm_wakeup(deadline, waiter.inner().clone());
 
         let mut timeout = true;
         loop {
-            let mut rq = current_run_queue::<NoPreemptIrqSave>();
             if axhal::time::wall_time() >= deadline {
                 break;
             }
-            let mut wq = self.list.lock();
             if condition() {
                 timeout = false;
                 break;
             }
-            rq.blocked_resched(wq, waiter.clone());
+
+            self.waiter_blocked(&waiter);
             // Preemption may occur here.
+            current_run_queue::<NoPreemptIrqSave>().reschedule();
         }
 
         // Always try to remove the task from the timer list.
@@ -217,7 +222,7 @@ impl WaitQueue {
             }
             cursor.move_next();
         }
-    }
+    } 
 }
 
 fn unblock_one_task(task: AxTaskRef, resched: bool) {