Num tasks

Cargo.toml

@@ -11,4 +11,4 @@ documentation = "https://arceos-org.github.io/scheduler"
 
 [dependencies]
 linked_list = { git = "https://github.com/arceos-org/linked_list.git", tag = "v0.1.0" }
-
+kspin = { version = "0.1" }

src/cfs.rs

@@ -2,6 +2,8 @@ use alloc::{collections::BTreeMap, sync::Arc};
 use core::ops::Deref;
 use core::sync::atomic::{AtomicIsize, Ordering};
 
+use kspin::SpinNoIrq;
+
 use crate::BaseScheduler;
 
 /// task for CFS
@@ -104,6 +106,7 @@ pub struct CFScheduler<T> {
     ready_queue: BTreeMap<(isize, isize), Arc<CFSTask<T>>>, // (vruntime, taskid)
     min_vruntime: Option<AtomicIsize>,
     id_pool: AtomicIsize,
+    lock: SpinNoIrq<()>,
 }
 
 impl<T> CFScheduler<T> {
@@ -113,6 +116,7 @@ impl<T> CFScheduler<T> {
             ready_queue: BTreeMap::new(),
             min_vruntime: None,
             id_pool: AtomicIsize::new(0_isize),
+            lock: SpinNoIrq::new(()),
         }
     }
     /// get the name of scheduler
@@ -127,6 +131,7 @@ impl<T> BaseScheduler for CFScheduler<T> {
     fn init(&mut self) {}
 
     fn add_task(&mut self, task: Self::SchedItem) {
+        let _lock = self.lock.lock();
         if self.min_vruntime.is_none() {
             self.min_vruntime = Some(AtomicIsize::new(0_isize));
         }
@@ -143,6 +148,7 @@ impl<T> BaseScheduler for CFScheduler<T> {
     }
 
     fn remove_task(&mut self, task: &Self::SchedItem) -> Option<Self::SchedItem> {
+        let _lock = self.lock.lock();
         if let Some((_, tmp)) = self
             .ready_queue
             .remove_entry(&(task.clone().get_vruntime(), task.clone().get_id()))
@@ -159,6 +165,7 @@ impl<T> BaseScheduler for CFScheduler<T> {
     }
 
     fn pick_next_task(&mut self) -> Option<Self::SchedItem> {
+        let _lock = self.lock.lock();
         if let Some((_, v)) = self.ready_queue.pop_first() {
             Some(v)
         } else {
@@ -167,6 +174,7 @@ impl<T> BaseScheduler for CFScheduler<T> {
     }
 
     fn put_prev_task(&mut self, prev: Self::SchedItem, _preempt: bool) {
+        let _lock = self.lock.lock();
         let taskid = self.id_pool.fetch_add(1, Ordering::Release);
         prev.set_id(taskid);
         self.ready_queue
@@ -187,4 +195,9 @@ impl<T> BaseScheduler for CFScheduler<T> {
             false
         }
     }
+
+    fn num_tasks(&self) -> usize {
+        // Do not lock the ready queue for efficiency considerations.
+        self.ready_queue.len()
+    }
 }

src/fifo.rs

@@ -1,6 +1,8 @@
 use alloc::sync::Arc;
 use core::ops::Deref;
+use core::sync::atomic::{AtomicUsize, Ordering};
 
+use kspin::SpinNoIrq;
 use linked_list::{Adapter, Links, List};
 
 use crate::BaseScheduler;
@@ -55,14 +57,16 @@ impl<T> Deref for FifoTask<T> {
 ///
 /// It internally uses a linked list as the ready queue.
 pub struct FifoScheduler<T> {
-    ready_queue: List<Arc<FifoTask<T>>>,
+    ready_queue: SpinNoIrq<List<Arc<FifoTask<T>>>>,
+    num_tasks: AtomicUsize,
 }
 
 impl<T> FifoScheduler<T> {
     /// Creates a new empty [`FifoScheduler`].
     pub const fn new() -> Self {
         Self {
-            ready_queue: List::new(),
+            ready_queue: SpinNoIrq::new(List::new()),
+            num_tasks: AtomicUsize::new(0),
         }
     }
     /// get the name of scheduler
@@ -77,19 +81,31 @@ impl<T> BaseScheduler for FifoScheduler<T> {
     fn init(&mut self) {}
 
     fn add_task(&mut self, task: Self::SchedItem) {
-        self.ready_queue.push_back(task);
+        self.num_tasks.fetch_add(1, Ordering::AcqRel);
+        self.ready_queue.lock().push_back(task);
     }
 
     fn remove_task(&mut self, task: &Self::SchedItem) -> Option<Self::SchedItem> {
-        unsafe { self.ready_queue.remove(task) }
+        let res = unsafe { self.ready_queue.lock().remove(task) };
+        if res.is_some() {
+            // Only decrement the number of tasks if the task is removed.
+            self.num_tasks.fetch_sub(1, Ordering::AcqRel);
+        }
+        res
     }
 
     fn pick_next_task(&mut self) -> Option<Self::SchedItem> {
-        self.ready_queue.pop_front()
+        let res = self.ready_queue.lock().pop_front();
+        if res.is_some() {
+            // Only decrement the number of tasks if the task is picked.
+            self.num_tasks.fetch_sub(1, Ordering::AcqRel);
+        }
+        res
     }
 
     fn put_prev_task(&mut self, prev: Self::SchedItem, _preempt: bool) {
-        self.ready_queue.push_back(prev);
+        self.num_tasks.fetch_add(1, Ordering::AcqRel);
+        self.ready_queue.lock().push_back(prev);
     }
 
     fn task_tick(&mut self, _current: &Self::SchedItem) -> bool {
@@ -99,4 +115,9 @@ impl<T> BaseScheduler for FifoScheduler<T> {
     fn set_priority(&mut self, _task: &Self::SchedItem, _prio: isize) -> bool {
         false
     }
+
+    fn num_tasks(&self) -> usize {
+        // Don't need to lock the ready queue to get the number of tasks.
+        self.num_tasks.load(Ordering::Acquire)
+    }
 }

src/lib.rs

@@ -63,6 +63,9 @@ pub trait BaseScheduler {
     /// `current` is the current running task.
     fn task_tick(&mut self, current: &Self::SchedItem) -> bool;
 
-    /// set priority for a task
+    /// Set priority for a task.
     fn set_priority(&mut self, task: &Self::SchedItem, prio: isize) -> bool;
+
+    /// Returns the number of tasks in the scheduler.
+    fn num_tasks(&self) -> usize;
 }

src/round_robin.rs

@@ -2,6 +2,8 @@ use alloc::{collections::VecDeque, sync::Arc};
 use core::ops::Deref;
 use core::sync::atomic::{AtomicIsize, Ordering};
 
+use kspin::SpinNoIrq;
+
 use crate::BaseScheduler;
 
 /// A task wrapper for the [`RRScheduler`].
@@ -57,13 +59,15 @@ impl<T, const S: usize> Deref for RRTask<T, S> {
 /// [`FifoScheduler`]: crate::FifoScheduler
 pub struct RRScheduler<T, const MAX_TIME_SLICE: usize> {
     ready_queue: VecDeque<Arc<RRTask<T, MAX_TIME_SLICE>>>,
+    lock: SpinNoIrq<()>,
 }
 
 impl<T, const S: usize> RRScheduler<T, S> {
     /// Creates a new empty [`RRScheduler`].
     pub const fn new() -> Self {
         Self {
             ready_queue: VecDeque::new(),
+            lock: SpinNoIrq::new(()),
         }
     }
     /// get the name of scheduler
@@ -78,10 +82,12 @@ impl<T, const S: usize> BaseScheduler for RRScheduler<T, S> {
     fn init(&mut self) {}
 
     fn add_task(&mut self, task: Self::SchedItem) {
+        let _lock = self.lock.lock();
         self.ready_queue.push_back(task);
     }
 
     fn remove_task(&mut self, task: &Self::SchedItem) -> Option<Self::SchedItem> {
+        let _lock = self.lock.lock();
         // TODO: more efficient
         self.ready_queue
             .iter()
@@ -90,10 +96,12 @@ impl<T, const S: usize> BaseScheduler for RRScheduler<T, S> {
     }
 
     fn pick_next_task(&mut self) -> Option<Self::SchedItem> {
+        let _lock = self.lock.lock();
         self.ready_queue.pop_front()
     }
 
     fn put_prev_task(&mut self, prev: Self::SchedItem, preempt: bool) {
+        let _lock = self.lock.lock();
         if prev.time_slice() > 0 && preempt {
             self.ready_queue.push_front(prev)
         } else {
@@ -110,4 +118,9 @@ impl<T, const S: usize> BaseScheduler for RRScheduler<T, S> {
     fn set_priority(&mut self, _task: &Self::SchedItem, _prio: isize) -> bool {
         false
     }
+
+    fn num_tasks(&self) -> usize {
+        // Do not lock the ready queue for efficiency considerations.
+        self.ready_queue.len()
+    }
 }