Properly implement rehydration

Signed-off-by: Klim Tsoutsman <klim@tsoutsman.com>

kernel/scheduler_epoch/src/lib.rs

@@ -5,6 +5,9 @@
 //!
 //! The scheduler is comprised of two run queues: an .
 //!
+//! Note that our implementation is not constant-time since we store
+//! non-runnable tasks on the run queue.
+//!
 //! [linux-scheduler]: https://litux.nl/mirror/kerneldevelopment/0672327201/ch04lev1sec2.html
 
 #![no_std]
@@ -103,13 +106,19 @@ impl task::scheduler::Scheduler for Scheduler {
             }
         }
         self.active
-            .next(&mut self.expired)
+            .next(&mut self.expired, self.total_weight)
             .unwrap_or(self.idle_task.clone())
     }
 
     #[inline]
     fn add(&mut self, task: TaskRef) {
-        let (task, weight) = EpochTaskRef::new(task, self.total_weight);
+        let (task, weight) = EpochTaskRef::new(
+            task,
+            TaskConfiguration {
+                priority: DEFAULT_PRIORITY as usize,
+                total_weight: self.total_weight,
+            },
+        );
         self.total_weight += weight;
         self.expired.push(task, DEFAULT_PRIORITY);
     }
@@ -170,21 +179,26 @@ struct EpochTaskRef {
 }
 
 impl EpochTaskRef {
+    /// Creates a new task.
+    ///
+    /// Returns the task and the weight of the task.
     #[must_use]
     pub(crate) fn new(task: TaskRef, config: TaskConfiguration) -> (Self, usize) {
-        const NUM_TOKENS: usize =
-            TARGET_LATENCY / kernel_config::time::CONFIG_TIMESLICE_PERIOD_MICROSECONDS;
+        let mut task = Self { task, tokens: 0 };
+        let weight = task.recalculate_tokens(config);
+        (task, weight)
+    }
+
+    #[inline]
+    pub(crate) fn recalculate_tokens(&mut self, config: TaskConfiguration) -> usize {
+        const TOTAL_TOKENS: usize = TARGET_LATENCY.as_micros() as usize
+            / kernel_config::time::CONFIG_TIMESLICE_PERIOD_MICROSECONDS as usize;
 
         // TODO
         let weight = config.priority + 1;
+        self.tokens = core::cmp::max(TOTAL_TOKENS * weight / config.total_weight, 1);
 
-        (
-            Self {
-                task,
-                tokens: core::cmp::max(NUM_TOKENS * weight / config.total_weight, 1),
-            },
-            weight,
-        )
+        weight
     }
 }
 
@@ -209,16 +223,6 @@ impl DerefMut for EpochTaskRef {
     }
 }
 
-impl EpochTaskRef {
-    #[inline]
-    fn new(task: TaskRef) -> EpochTaskRef {
-        EpochTaskRef {
-            task,
-            tokens: INITIAL_TOKENS,
-        }
-    }
-}
-
 impl From<EpochTaskRef> for TaskRef {
     #[inline]
     fn from(value: EpochTaskRef) -> Self {

kernel/scheduler_epoch/src/queue.rs

@@ -3,7 +3,7 @@ use alloc::collections::VecDeque;
 use bit_set::BitSet;
 use task::TaskRef;
 
-use crate::{EpochTaskRef, MAX_PRIORITY};
+use crate::{EpochTaskRef, TaskConfiguration, MAX_PRIORITY};
 
 /// A singular run queue.
 ///
@@ -29,13 +29,16 @@ impl RunQueue {
     }
 
     #[inline]
-    pub(crate) const fn len(&self) -> usize {
-        debug_assert_eq!(self.inner.iter().map(|queue| queue.len()).sum(), self.len);
+    pub(crate) fn len(&self) -> usize {
+        debug_assert_eq!(
+            self.inner.iter().map(|queue| queue.len()).sum::<usize>(),
+            self.len
+        );
         self.len
     }
 
     #[inline]
-    pub(crate) const fn is_empty(&self) -> bool {
+    pub(crate) fn is_empty(&self) -> bool {
         self.len() == 0
     }
 
@@ -47,37 +50,93 @@ impl RunQueue {
     }
 
     #[inline]
-    pub(crate) fn next(&mut self, expired: &mut Self) -> Option<TaskRef> {
-        loop {
-            let top_index = self.top_index()?;
-            let top_queue = &mut self.inner[top_index];
+    pub(crate) fn next(&mut self, expired: &mut Self, total_weight: usize) -> Option<TaskRef> {
+        let mut priorities = self.priorities.clone();
+
+        let mut top_index = priorities.max()?;
+        // TODO: top_queue.len() == 1 optimisation
+        let mut top_queue = &mut self.inner[top_index as usize];
+        let mut next_task = top_queue.front().unwrap();
+
+        if !next_task.is_runnable() {
+            // TODO: This incredibly convoluted code is necessary because we store
+            // non-runnable tasks on the run queue.
+
+            // Iterate through the queue to find the next runnable task and bring it to the
+            // front of its respective run queue.
+
+            let mut vec_index = 0;
+
+            while !next_task.is_runnable() {
+                vec_index += 1;
+
+                if vec_index + 1 == top_queue.len() {
+                    priorities.remove(top_index);
+                    top_index = match priorities.max() {
+                        Some(top) => top,
+                        None => {
+                            // There are no runnable tasks on the run queue. We
+                            // must transfer all the tasks to the expired run
+                            // queue and return None.
+
+                            let mut priorities = self.priorities.clone();
+
+                            while let Some(top_index) = priorities.max() {
+                                let top_queue = &mut self.inner[top_index as usize];
+
+                                while let Some(mut task) = top_queue.pop_front() {
+                                    task.recalculate_tokens(TaskConfiguration {
+                                        priority: top_index as usize,
+                                        total_weight,
+                                    });
+                                    expired.push(task, top_index);
+                                }
+
+                                priorities.remove(top_index);
+                            }
+
+                            return None;
+                        }
+                    };
+                    vec_index = 0;
+                }
 
-            // TODO: top_queue.len() == 1 optimisation
+                top_queue = &mut self.inner[top_index as usize];
+                next_task = &top_queue[vec_index];
+            }
 
-            let mut next_task = top_queue.pop_front().unwrap();
+            for _ in 0..vec_index {
+                let task = top_queue.pop_front().unwrap();
+                top_queue.push_back(task);
+            }
+        }
 
-            if !next_task.is_runnable() {
-                self.len -= 1;
-                expired.push(next_task.clone(), top_index as u8);
+        let queue = &mut self.inner[top_index as usize];
+        let next_task = queue.front().unwrap();
 
-                if top_queue.is_empty() {
-                    self.priorities.remove(top_index as u8);
-                }
-            } else if next_task.tokens <= 1 {
-                self.len -= 1;
-                expired.push(next_task.clone(), top_index as u8);
+        Some(if next_task.tokens <= 1 {
+            let mut next_task = queue.pop_front().unwrap();
+            self.len -= 1;
 
-                if top_queue.is_empty() {
-                    self.priorities.remove(top_index as u8);
-                }
+            next_task.recalculate_tokens(TaskConfiguration {
+                priority: top_index as usize,
+                total_weight,
+            });
+            expired.push(next_task.clone(), top_index);
 
-                return Some(next_task.task);
-            } else {
-                next_task.tokens -= 1;
-                top_queue.push_back(next_task.clone());
-                return Some(next_task.task);
+            if queue.is_empty() {
+                self.priorities.remove(top_index);
             }
-        }
+
+            next_task.clone().task
+        } else {
+            let mut next_task = queue.pop_front().unwrap();
+
+            next_task.tokens -= 1;
+            queue.push_back(next_task.clone());
+
+            next_task.task
+        })
     }
 
     #[inline]
@@ -175,18 +234,12 @@ impl Iterator for Drain {
 
     fn next(&mut self) -> Option<Self::Item> {
         let top_index = self.inner.top_index()?;
+        let top_queue = &mut self.inner.inner[top_index];
 
-        if top_index == 64 {
-            None
-        } else {
-            let top_queue = &mut self.inner.inner[top_index];
-
-            if top_queue.len() == 1 {
-                let priority = 64 - top_index;
-                self.inner.priorities.remove(priority as u8);
-            }
-
-            Some(top_queue.pop_front().unwrap().into())
+        if top_queue.len() == 1 {
+            self.inner.priorities.remove(top_index as u8);
         }
+
+        Some(top_queue.pop_front().unwrap().into())
     }
 }