diff --git a/Cargo.toml b/Cargo.toml index 9133e80..d3ab55a 100644 --- a/Cargo.toml +++ b/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" } diff --git a/src/cfs.rs b/src/cfs.rs index 30043aa..e551266 100644 --- a/src/cfs.rs +++ b/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 { ready_queue: BTreeMap<(isize, isize), Arc>>, // (vruntime, taskid) min_vruntime: Option, id_pool: AtomicIsize, + lock: SpinNoIrq<()>, } impl CFScheduler { @@ -113,6 +116,7 @@ impl CFScheduler { 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 BaseScheduler for CFScheduler { 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 BaseScheduler for CFScheduler { } fn remove_task(&mut self, task: &Self::SchedItem) -> Option { + 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 BaseScheduler for CFScheduler { } fn pick_next_task(&mut self) -> Option { + let _lock = self.lock.lock(); if let Some((_, v)) = self.ready_queue.pop_first() { Some(v) } else { @@ -167,6 +174,7 @@ impl BaseScheduler for CFScheduler { } 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 BaseScheduler for CFScheduler { false } } + + fn num_tasks(&self) -> usize { + // Do not lock the ready queue for efficiency considerations. + self.ready_queue.len() + } } diff --git a/src/fifo.rs b/src/fifo.rs index f2469e5..1ce9ea5 100644 --- a/src/fifo.rs +++ b/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 Deref for FifoTask { /// /// It internally uses a linked list as the ready queue. pub struct FifoScheduler { - ready_queue: List>>, + ready_queue: SpinNoIrq>>>, + num_tasks: AtomicUsize, } impl FifoScheduler { /// 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 BaseScheduler for FifoScheduler { 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 { - 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.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 BaseScheduler for FifoScheduler { 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) + } } diff --git a/src/lib.rs b/src/lib.rs index de15661..9a16d3e 100644 --- a/src/lib.rs +++ b/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; } diff --git a/src/round_robin.rs b/src/round_robin.rs index bf7e2c4..4732363 100644 --- a/src/round_robin.rs +++ b/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,6 +59,7 @@ impl Deref for RRTask { /// [`FifoScheduler`]: crate::FifoScheduler pub struct RRScheduler { ready_queue: VecDeque>>, + lock: SpinNoIrq<()>, } impl RRScheduler { @@ -64,6 +67,7 @@ impl RRScheduler { pub const fn new() -> Self { Self { ready_queue: VecDeque::new(), + lock: SpinNoIrq::new(()), } } /// get the name of scheduler @@ -78,10 +82,12 @@ impl BaseScheduler for RRScheduler { 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 { + let _lock = self.lock.lock(); // TODO: more efficient self.ready_queue .iter() @@ -90,10 +96,12 @@ impl BaseScheduler for RRScheduler { } fn pick_next_task(&mut self) -> Option { + 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 BaseScheduler for RRScheduler { 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() + } }