Lazily clean up exited tasks in scheduler

Avoids locking all the schedulers on task cleanup.

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