Change scheduler to only store runnable tasks on run queue

applications/test_sync_block/src/lib.rs

@@ -71,7 +71,7 @@ fn sync_block_task(lock: Arc<Mutex<usize>>) -> Result<(), &'static str> {
         .map_err(|_| "couldn't get current task")?;
     warn!("ENTERED TASK {}", curr_task);
 
-    for _i in 0..1000 {
+    for _i in 0..1 {
         scheduler::schedule(); // give other tasks a chance to acquire the lock
         warn!("{} trying to acquire lock...", curr_task);
         let mut locked = lock.lock();
@@ -131,17 +131,20 @@ fn lockstep_task((lock, remainder): (Arc<Mutex<usize>>, usize)) -> Result<(), &'
 
     for _i in 0..20 {
         loop { 
-            warn!("{} top of loop, remainder {}", curr_task, remainder);
+            if remainder == 1 {
+                warn!("{} top of loop, remainder {}", curr_task, remainder);
+            }
+            // warn!("schedulers: {:#?}", task::scheduler::tasks());
             scheduler::schedule(); // give other tasks a chance to acquire the lock
             let mut locked = lock.lock();
             scheduler::schedule();
             if *locked % 3 == remainder {
-                warn!("Task {} Time to shine, value is {}!", curr_task, *locked);
+                // warn!("Task {} Time to shine, value is {}!", curr_task, *locked);
                 *locked += 1;
                 break;
             } else {
                 scheduler::schedule();
-                warn!("Task {} going back to sleep, value {}, remainder {}!", curr_task, *locked, remainder);
+                // warn!("Task {} going back to sleep, value {}, remainder {}!", curr_task, *locked, remainder);
             }
             scheduler::schedule();
         }

kernel/scheduler_epoch/src/lib.rs

@@ -8,17 +8,25 @@
 //! ```
 //! * Each time a task is picked, its token count is decremented by 1.
 //! * A task can only be selected for next execution if it has tokens remaining.
-//! * When all tokens of all runnable task are exhausted, a new scheduling epoch begins.
+//! * When all tokens of all runnable task are exhausted, a new scheduling epoch
+//!   begins.
 //!
 //! This epoch scheduler is also a priority-based scheduler, so it allows
 //! getting and setting the priorities of each task.
 
 #![no_std]
+#![feature(core_intrinsics)]
 
 extern crate alloc;
 
 use alloc::{boxed::Box, collections::VecDeque, vec::Vec};
-use core::ops::{Deref, DerefMut};
+use core::{
+    cmp::max,
+    intrinsics::unlikely,
+    ops::{Deref, DerefMut},
+    sync::atomic::Ordering,
+};
+
 use task::TaskRef;
 
 const MAX_PRIORITY: u8 = 40;
@@ -28,142 +36,141 @@ const INITIAL_TOKENS: usize = 10;
 /// An instance of an epoch scheduler, typically one per CPU.
 pub struct Scheduler {
     idle_task: TaskRef,
-    queue: VecDeque<EpochTaskRef>,
+    have_tokens: VecDeque<EpochTaskRef>,
+    out_of_tokens: Vec<EpochTaskRef>,
 }
 
 impl Scheduler {
     /// Creates a new epoch scheduler instance with the given idle task.
     pub const fn new(idle_task: TaskRef) -> Self {
         Self {
             idle_task,
-            queue: VecDeque::new(),
+            have_tokens: VecDeque::new(),
+            out_of_tokens: Vec::new(),
         }
     }
 
-    /// Moves the `TaskRef` at the given `index` in this scheduler's runqueue
-    /// to the end (back) of the runqueue.
-    ///
-    /// Sets the number of tokens for that task to the given `tokens`
-    /// and increments that task's number of context switches.
-    ///
-    /// Returns a cloned reference to the `TaskRef` at the given `index`.
-    fn update_and_move_to_end(&mut self, index: usize, tokens: usize) -> Option<TaskRef> {
-        if let Some(mut priority_task_ref) = self.queue.remove(index) {
-            priority_task_ref.tokens_remaining = tokens;
-            let task_ref = priority_task_ref.task.clone();
-            self.queue.push_back(priority_task_ref);
-            Some(task_ref)
-        } else {
-            None
+    fn try_next(&mut self) -> Option<TaskRef> {
+        while let Some(task) = self.have_tokens.pop_front() {
+            if task.task.is_runnable() {
+                if let Some(task) = self.add_epoch_task(task) {
+                    return Some(task);
+                }
+            } else {
+                task.task
+                    .expose_is_on_run_queue()
+                    .store(false, Ordering::Release);
+                // This check prevents an interleaving where `TaskRef::unblock` wouldn't add
+                // the task back onto the run queue. `TaskRef::unblock` sets the run state and
+                // then checks `is_on_run_queue` so we have to do the inverse.
+                if unlikely(task.task.is_runnable()) {
+                    if let Some(task) = self.add_epoch_task(task) {
+                        return Some(task);
+                    }
+                }
+            }
         }
+        None
     }
 
-    fn try_next(&mut self) -> Option<TaskRef> {
-        if let Some((task_index, _)) = self
-            .queue
-            .iter()
-            .enumerate()
-            .find(|(_, task)| task.is_runnable() && task.tokens_remaining > 0)
-        {
-            let chosen_task = self.queue.get(task_index).unwrap();
-            let modified_tokens = chosen_task.tokens_remaining.saturating_sub(1);
-            self.update_and_move_to_end(task_index, modified_tokens)
-        } else {
-            None
+    fn add_epoch_task(&mut self, mut task: EpochTaskRef) -> Option<TaskRef> {
+        task.task
+            .expose_is_on_run_queue()
+            .store(true, Ordering::Release);
+        match task.tokens_remaining.checked_sub(1) {
+            Some(new_tokens) => {
+                task.tokens_remaining = new_tokens;
+                let task_ref = task.task.clone();
+                self.have_tokens.push_back(task);
+                Some(task_ref)
+            }
+            None => {
+                self.out_of_tokens.push(task);
+                None
+            }
         }
     }
 
     fn assign_tokens(&mut self) {
-        // We begin with total priorities = 1 to avoid division by zero
-        let mut total_priorities: usize = 1;
-
-        // This loop calculates the total priorities of the runqueue
-        for (_i, t) in self.queue.iter().enumerate() {
-            // we assign tokens only to runnable tasks
-            if !t.is_runnable() {
-                continue;
-            }
-
-            total_priorities = total_priorities
-                .saturating_add(1)
-                .saturating_add(t.priority as usize);
+        while let Some(task) = self.out_of_tokens.pop() {
+            self.have_tokens.push_back(task);
         }
 
-        // Each epoch lasts for a total of 100 tokens by default.
-        // However, as this granularity could skip over low priority tasks
-        // when many concurrent tasks are running, we increase the epoch in such cases.
-        let epoch: usize = core::cmp::max(total_priorities, 100);
-
-        for (_i, t) in self.queue.iter_mut().enumerate() {
-            // we give zero tokens to the idle tasks
-            if t.is_an_idle_task {
-                continue;
-            }
+        let mut total_priorities = 1;
+        for task in self.have_tokens.iter() {
+            total_priorities += 1 + task.priority as usize;
+        }
 
-            // we give zero tokens to non-runnable tasks
-            if !t.is_runnable() {
-                continue;
-            }
+        // Each epoch lasts for a total of 100 tokens by default. However, as this
+        // granularity could skip over low priority tasks when many concurrent tasks are
+        // running, we increase the epoch in such cases.
+        let epoch = max(total_priorities, 100);
 
-            // task_tokens = epoch * (taskref + 1) / total_priorities;
-            let task_tokens = epoch
-                .saturating_mul((t.priority as usize).saturating_add(1))
+        for task in self.have_tokens.iter_mut() {
+            task.tokens_remaining = epoch
+                .saturating_mul((task.priority as usize).saturating_add(1))
                 .wrapping_div(total_priorities);
-
-            t.tokens_remaining = task_tokens;
-            // debug!("assign_tokens(): CPU {} chose Task {:?}", cpu_id, &*t);
         }
     }
+
+    fn len(&self) -> usize {
+        self.have_tokens.len() + self.out_of_tokens.len()
+    }
 }
 
 impl task::scheduler::Scheduler for Scheduler {
-    fn next(&mut self) -> TaskRef {
-        self.try_next()
+    fn next(&mut self, _current_task: TaskRef) -> Option<TaskRef> {
+        // FIXME: Return none
+        Some(self.try_next()
             .or_else(|| {
                 self.assign_tokens();
                 self.try_next()
             })
-            .unwrap_or_else(|| self.idle_task.clone())
+            .unwrap_or_else(|| self.idle_task.clone()))
     }
 
     fn add(&mut self, task: TaskRef) {
-        let priority_task_ref = EpochTaskRef::new(task);
-        self.queue.push_back(priority_task_ref);
+        let epoch_task_ref = EpochTaskRef::new(task);
+        self.add_epoch_task(epoch_task_ref);
     }
 
     fn busyness(&self) -> usize {
-        self.queue.len()
+        self.len()
     }
 
     fn remove(&mut self, task: &TaskRef) -> bool {
-        let mut task_index = None;
-        for (i, t) in self.queue.iter().enumerate() {
-            if **t == *task {
-                task_index = Some(i);
-                break;
-            }
-        }
-
-        if let Some(task_index) = task_index {
-            self.queue.remove(task_index);
-            true
-        } else {
-            false
-        }
+        let old_len = self.len();
+        self.have_tokens
+            .retain(|other_task| other_task.task != *task);
+        self.out_of_tokens
+            .retain(|other_task| other_task.task != *task);
+        let new_len = self.len();
+        debug_assert!(
+            old_len - new_len < 2,
+            "difference between run queue lengths was: {}",
+            old_len - new_len
+        );
+        new_len != old_len
     }
 
     fn as_priority_scheduler(&mut self) -> Option<&mut dyn task::scheduler::PriorityScheduler> {
         Some(self)
     }
 
     fn drain(&mut self) -> Box<dyn Iterator<Item = TaskRef> + '_> {
-        Box::new(self.queue.drain(..).map(|epoch_task| epoch_task.task))
+        Box::new(
+            self.have_tokens
+                .drain(..)
+                .chain(self.out_of_tokens.drain(..))
+                .map(|epoch_task| epoch_task.task),
+        )
     }
 
     fn tasks(&self) -> Vec<TaskRef> {
-        self.queue
+        self.have_tokens
             .clone()
             .into_iter()
+            .chain(self.out_of_tokens.clone())
             .map(|epoch_task| epoch_task.task)
             .collect()
     }
@@ -172,7 +179,12 @@ impl task::scheduler::Scheduler for Scheduler {
 impl task::scheduler::PriorityScheduler for Scheduler {
     fn set_priority(&mut self, task: &TaskRef, priority: u8) -> bool {
         let priority = core::cmp::min(priority, MAX_PRIORITY);
-        for epoch_task in self.queue.iter_mut() {
+
+        for epoch_task in self
+            .have_tokens
+            .iter_mut()
+            .chain(self.out_of_tokens.iter_mut())
+        {
             if epoch_task.task == *task {
                 epoch_task.priority = priority;
                 return true;
@@ -182,7 +194,7 @@ impl task::scheduler::PriorityScheduler for Scheduler {
     }
 
     fn priority(&mut self, task: &TaskRef) -> Option<u8> {
-        for epoch_task in self.queue.iter() {
+        for epoch_task in self.have_tokens.iter().chain(self.out_of_tokens.iter()) {
             if epoch_task.task == *task {
                 return Some(epoch_task.priority);
             }