feat: support percpu run-queue and cpu affinity for axtask, bug in wa…

…it queue?

modules/axtask/src/api.rs

@@ -86,7 +86,10 @@ pub fn init_scheduler_secondary() {
 #[doc(cfg(feature = "irq"))]
 pub fn on_timer_tick() {
     crate::timers::check_events();
-    current_run_queue().scheduler_timer_tick();
+    current_run_queue()
+        .scheduler()
+        .lock()
+        .scheduler_timer_tick();
 }
 
 /// Spawns a new task with the given parameters.
@@ -107,6 +110,8 @@ where
         #[cfg(feature = "smp")]
         task.clone(),
     )
+    .scheduler()
+    .lock()
     .add_task(task.clone());
     task
 }
@@ -134,13 +139,16 @@ where
 ///
 /// [CFS]: https://en.wikipedia.org/wiki/Completely_Fair_Scheduler
 pub fn set_priority(prio: isize) -> bool {
-    current_run_queue().set_current_priority(prio)
+    current_run_queue()
+        .scheduler()
+        .lock()
+        .set_current_priority(prio)
 }
 
 /// Current task gives up the CPU time voluntarily, and switches to another
 /// ready task.
 pub fn yield_now() {
-    current_run_queue().yield_current();
+    current_run_queue().scheduler().lock().yield_current()
 }
 
 /// Current task is going to sleep for the given duration.
@@ -155,7 +163,7 @@ pub fn sleep(dur: core::time::Duration) {
 /// If the feature `irq` is not enabled, it uses busy-wait instead.
 pub fn sleep_until(deadline: axhal::time::TimeValue) {
     #[cfg(feature = "irq")]
-    current_run_queue().sleep_until(deadline);
+    current_run_queue().scheduler().lock().sleep_until(deadline);
     #[cfg(not(feature = "irq"))]
     axhal::time::busy_wait_until(deadline);
 }

modules/axtask/src/run_queue.rs

@@ -157,6 +157,7 @@ pub(crate) struct AxRunQueue {
 pub struct AxRunQueueInner {
     /// The ID of the CPU this run queue is associated with.
     cpu_id: usize,
+    /// The core scheduler of this run queue.
     scheduler: Scheduler,
 }
 
@@ -199,43 +200,6 @@ impl AxRunQueue {
         &self.inner
     }
 
-    pub fn add_task(&self, task: AxTaskRef) {
-        let mut scheduler_lock = self.inner.lock();
-        debug!(
-            "task spawn: {} on run_queue {}",
-            task.id_name(),
-            self.cpu_id
-        );
-        assert!(task.is_ready());
-        scheduler_lock.scheduler.add_task(task);
-        self.num_tasks.fetch_add(1, Ordering::AcqRel);
-    }
-
-    #[cfg(feature = "irq")]
-    pub fn scheduler_timer_tick(&self) {
-        let mut scheduler_lock = self.inner.lock();
-        let curr = crate::current();
-        if !curr.is_idle() && scheduler_lock.scheduler.task_tick(curr.as_task_ref()) {
-            #[cfg(feature = "preempt")]
-            curr.set_preempt_pending(true);
-        }
-    }
-
-    pub fn yield_current(&self) {
-        let mut scheduler_lock = self.scheduler().lock();
-        let curr = crate::current();
-        trace!("task yield: {}", curr.id_name());
-        assert!(curr.is_running());
-        scheduler_lock.resched(false);
-    }
-
-    pub fn set_current_priority(&self, prio: isize) -> bool {
-        self.inner
-            .lock()
-            .scheduler
-            .set_priority(crate::current().as_task_ref(), prio)
-    }
-
     pub fn exit_current(&self, exit_code: i32) -> ! {
         // We do not own an `SpinNoIrq` lock here, so we need to disable IRQ and preempt manually.
         let _kernel_guard = kernel_guard::IrqSave::new();
@@ -263,26 +227,44 @@ impl AxRunQueue {
         }
         unreachable!("task exited!");
     }
+}
+
+/// Core functions of run queue, which should be called after holding the scheduler() lock.
+impl AxRunQueueInner {
+    pub fn add_task(&mut self, task: AxTaskRef) {
+        debug!(
+            "task spawn: {} on run_queue {}",
+            task.id_name(),
+            self.cpu_id
+        );
+        assert!(task.is_ready());
+        self.scheduler.add_task(task);
+        get_run_queue(self.cpu_id)
+            .num_tasks
+            .fetch_add(1, Ordering::AcqRel);
+    }
 
     #[cfg(feature = "irq")]
-    pub fn sleep_until(&self, deadline: axhal::time::TimeValue) {
-        let mut scheduler_lock = self.inner.lock();
+    pub fn scheduler_timer_tick(&mut self) {
         let curr = crate::current();
-        debug!("task sleep: {}, deadline={:?}", curr.id_name(), deadline);
+        if !curr.is_idle() && self.scheduler.task_tick(curr.as_task_ref()) {
+            #[cfg(feature = "preempt")]
+            curr.set_preempt_pending(true);
+        }
+    }
+
+    pub fn yield_current(&mut self) {
+        let curr = crate::current();
+        trace!("task yield: {}", curr.id_name());
         assert!(curr.is_running());
-        assert!(!curr.is_idle());
+        self.resched(false);
+    }
 
-        let now = axhal::time::wall_time();
-        if now < deadline {
-            crate::timers::set_alarm_wakeup(deadline, curr.clone());
-            curr.set_state(TaskState::Blocked);
-            self.num_tasks.fetch_sub(1, Ordering::AcqRel);
-            scheduler_lock.resched(false);
-        }
+    pub fn set_current_priority(&mut self, prio: isize) -> bool {
+        self.scheduler
+            .set_priority(crate::current().as_task_ref(), prio)
     }
-}
 
-impl AxRunQueueInner {
     #[cfg(feature = "preempt")]
     pub fn preempt_resched(&mut self) {
         let curr = crate::current();
@@ -307,9 +289,28 @@ impl AxRunQueueInner {
         }
     }
 
+    pub fn block_current<F>(&mut self, wait_queue_push: F)
+    where
+        F: FnOnce(AxTaskRef),
+    {
+        let curr = crate::current();
+        debug!("task block: {}", curr.id_name());
+        assert!(curr.is_running());
+        assert!(!curr.is_idle());
+
+        // we must not block current task with preemption disabled.
+        #[cfg(feature = "preempt")]
+        assert!(curr.can_preempt(1));
+
+        curr.set_state(TaskState::Blocked);
+        current_run_queue().num_tasks.fetch_sub(1, Ordering::AcqRel);
+        wait_queue_push(curr.clone());
+        self.resched(false);
+    }
+
     pub fn unblock_task(&mut self, task: AxTaskRef, resched: bool) {
         let cpu_id = self.cpu_id;
-        debug!("task unblock: {} on run_queue {}", task.id_name(), cpu_id,);
+        debug!("task unblock: {} on run_queue {}", task.id_name(), cpu_id);
         if task.is_blocked() {
             task.set_state(TaskState::Ready);
             self.scheduler.add_task(task); // TODO: priority
@@ -327,23 +328,22 @@ impl AxRunQueueInner {
         }
     }
 
-    pub fn block_current<F>(&mut self, wait_queue_push: F)
-    where
-        F: FnOnce(AxTaskRef),
-    {
+    #[cfg(feature = "irq")]
+    pub fn sleep_until(&mut self, deadline: axhal::time::TimeValue) {
         let curr = crate::current();
-        debug!("task block: {}", curr.id_name());
+        debug!("task sleep: {}, deadline={:?}", curr.id_name(), deadline);
         assert!(curr.is_running());
         assert!(!curr.is_idle());
 
-        // we must not block current task with preemption disabled.
-        #[cfg(feature = "preempt")]
-        assert!(curr.can_preempt(1));
-
-        curr.set_state(TaskState::Blocked);
-        current_run_queue().num_tasks.fetch_sub(1, Ordering::AcqRel);
-        wait_queue_push(curr.clone());
-        self.resched(false);
+        let now = axhal::time::wall_time();
+        if now < deadline {
+            crate::timers::set_alarm_wakeup(deadline, curr.clone());
+            curr.set_state(TaskState::Blocked);
+            get_run_queue(self.cpu_id)
+                .num_tasks
+                .fetch_sub(1, Ordering::AcqRel);
+            self.resched(false);
+        }
     }
 }