Lazily remove exited tasks from run queue in scheduler

Avoids locking all the schedulers on task exit.

This PR also technically changes the round robin scheduler algorithm, as
blocked tasks are now moved to the end of the run queue. Previously, the
blocked tasks would be kept in place, aside from the task at the front of
the queue which would be switched with the next runnable task using
`swap_remove_front`.

Signed-off-by: Klim Tsoutsman <[email protected]>

kernel/scheduler_epoch/src/lib.rs

@@ -14,11 +14,16 @@
 //! 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::{
+    intrinsics::likely,
+    ops::{Deref, DerefMut},
+};
+
 use task::TaskRef;
 
 const MAX_PRIORITY: u8 = 40;
@@ -40,37 +45,24 @@ impl Scheduler {
         }
     }
 
-    /// 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> {
-        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
+        let len = self.queue.len();
+        let mut i = 0;
+
+        while i < len {
+            let mut task = self.queue.pop_front().unwrap();
+
+            if task.is_runnable() && task.tokens_remaining > 0 {
+                task.tokens_remaining -= 1;
+                self.queue.push_back(task.clone());
+                return Some(task.task);
+            } else if likely(!task.is_complete()) {
+                self.queue.push_back(task);
+            }
+            i += 1;
         }
+
+        None
     }
 
     fn assign_tokens(&mut self) {

kernel/scheduler_priority/src/lib.rs

@@ -1,11 +1,12 @@
 //! This scheduler implements a priority algorithm.
 
 #![no_std]
+#![feature(core_intrinsics)]
 
 extern crate alloc;
 
 use alloc::{boxed::Box, collections::BinaryHeap, vec::Vec};
-use core::cmp::Ordering;
+use core::{cmp::Ordering, intrinsics::likely};
 
 use task::TaskRef;
 use time::Instant;
@@ -39,7 +40,7 @@ impl task::scheduler::Scheduler for Scheduler {
                 task.last_ran = time::now::<time::Monotonic>();
                 self.queue.push(task.clone());
                 return task.task;
-            } else {
+            } else if likely(!task.task.is_complete()) {
                 blocked_tasks.push(task);
             }
         }

kernel/scheduler_round_robin/src/lib.rs

@@ -3,10 +3,12 @@
 //! This task is then moved to the back of the queue.
 
 #![no_std]
+#![feature(core_intrinsics)]
 
 extern crate alloc;
 
 use alloc::{boxed::Box, collections::VecDeque, vec::Vec};
+use core::intrinsics::likely;
 
 use task::TaskRef;
 
@@ -26,18 +28,22 @@ impl Scheduler {
 
 impl task::scheduler::Scheduler for Scheduler {
     fn next(&mut self) -> TaskRef {
-        if let Some((task_index, _)) = self
-            .queue
-            .iter()
-            .enumerate()
-            .find(|(_, task)| task.is_runnable())
-        {
-            let task = self.queue.swap_remove_front(task_index).unwrap();
-            self.queue.push_back(task.clone());
-            task
-        } else {
-            self.idle_task.clone()
+        let len = self.queue.len();
+        let mut i = 0;
+
+        while i < len {
+            let task = self.queue.pop_front().unwrap();
+
+            if task.is_runnable() {
+                self.queue.push_back(task.clone());
+                return task;
+            } else if likely(!task.is_complete()) {
+                self.queue.push_back(task);
+            }
+            i += 1;
         }
+
+        self.idle_task.clone()
     }
 
     fn busyness(&self) -> usize {

kernel/spawn/src/lib.rs

@@ -155,7 +155,6 @@ impl Drop for BootstrapTaskRef {
     // See the documentation for `BootstrapTaskRef::finish()` for more details.
     fn drop(&mut self) {
         // trace!("Finishing Bootstrap Task on core {}: {:?}", self.cpu_id, self.task_ref);
-        remove_current_task_from_runqueue(&self.exitable_taskref);
         self.exitable_taskref.mark_as_exited(Box::new(()))
             .expect("BUG: bootstrap task was unable to mark itself as exited");
 
@@ -998,11 +997,6 @@ where
     loop { core::hint::spin_loop() }
 }
 
-/// Helper function to remove a task from its runqueue and drop it.
-fn remove_current_task_from_runqueue(current_task: &ExitableTaskRef) {
-    task::scheduler::remove_task(current_task);
-}
-
 /// A basic idle task that does nothing but loop endlessly.
 ///
 /// Note: the current spawn API does not support spawning a task with the return type `!`,

kernel/task_struct/src/lib.rs

@@ -397,10 +397,23 @@ impl Task {
     /// [`Runnable`]: RunState::Runnable
     /// [suspended]: Task::is_suspended
     /// [running]: Task::is_running
+    #[inline]
     pub fn is_runnable(&self) -> bool {
         self.runstate() == RunState::Runnable && !self.is_suspended()
     }
 
+    /// Returns whether this `Task` is complete i.e. will never be runnable again.
+    ///
+    /// A task is complete if it is in an [`Exited`] or [`Reaped`] run state.
+    ///
+    /// [`Exited`]: RunState::Exited
+    /// [`Reaped`]: RunState::Reaped
+    #[inline]
+    pub fn is_complete(&self) -> bool {
+        let run_state = self.runstate();
+        run_state == RunState::Exited || run_state == RunState::Reaped
+    }
+
     /// Returns the namespace that this `Task` is loaded/linked into and runs within.
     pub fn get_namespace(&self) -> &Arc<CrateNamespace> {
         &self.namespace