Complete implementation

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

kernel/scheduler_epoch/src/lib.rs

@@ -58,9 +58,9 @@ impl Scheduler {
         }
     }
 
-    fn apply<F, R>(&mut self, f: F) -> R
+    fn apply<F, R>(&mut self, mut f: F) -> R
     where
-        F: Fn(&mut RunQueue) -> R,
+        F: FnMut(&mut RunQueue) -> R,
         R: Returnable,
     {
         let (first, second) = if self.active.len() >= self.expired.len() {
@@ -112,14 +112,16 @@ impl task::scheduler::Scheduler for Scheduler {
 
     #[inline]
     fn add(&mut self, task: TaskRef) {
-        let (task, weight) = EpochTaskRef::new(
+        let weight = weight(DEFAULT_PRIORITY);
+        self.total_weight += weight;
+
+        let task = EpochTaskRef::new(
             task,
             TaskConfiguration {
-                priority: DEFAULT_PRIORITY as usize,
+                weight,
                 total_weight: self.total_weight,
             },
         );
-        self.total_weight += weight;
         self.expired.push(task, DEFAULT_PRIORITY);
     }
 
@@ -130,7 +132,13 @@ impl task::scheduler::Scheduler for Scheduler {
 
     #[inline]
     fn remove(&mut self, task: &TaskRef) -> bool {
-        self.apply(|run_queue| run_queue.remove(task))
+        match self.apply(|run_queue| run_queue.remove(task)) {
+            Some(weight) => {
+                self.total_weight -= weight;
+                true
+            }
+            None => false,
+        }
     }
 
     #[inline]
@@ -145,6 +153,7 @@ impl task::scheduler::Scheduler for Scheduler {
 
         mem::swap(&mut self.active, &mut active);
         mem::swap(&mut self.expired, &mut expired);
+        self.total_weight = 0;
 
         Box::new(active.drain().chain(expired.drain()))
     }
@@ -172,6 +181,11 @@ impl task::scheduler::PriorityScheduler for Scheduler {
     }
 }
 
+#[inline]
+fn weight(priority: u8) -> usize {
+    priority as usize + 1
+}
+
 #[derive(Debug, Clone)]
 struct EpochTaskRef {
     task: TaskRef,
@@ -183,27 +197,26 @@ impl EpochTaskRef {
     ///
     /// Returns the task and the weight of the task.
     #[must_use]
-    pub(crate) fn new(task: TaskRef, config: TaskConfiguration) -> (Self, usize) {
+    pub(crate) fn new(task: TaskRef, config: TaskConfiguration) -> Self {
         let mut task = Self { task, tokens: 0 };
-        let weight = task.recalculate_tokens(config);
-        (task, weight)
+        task.recalculate_tokens(config);
+        task
     }
 
     #[inline]
-    pub(crate) fn recalculate_tokens(&mut self, config: TaskConfiguration) -> usize {
+    pub(crate) fn recalculate_tokens(&mut self, config: TaskConfiguration) {
         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);
-
-        weight
+        self.tokens = core::cmp::max(TOTAL_TOKENS * config.weight / config.total_weight, 1);
     }
 }
 
 pub(crate) struct TaskConfiguration {
-    pub(crate) priority: usize,
+    /// The weight of the task.
+    pub(crate) weight: usize,
+    /// The sum of the weights of all tasks on the run queue.
     pub(crate) total_weight: usize,
 }
 

kernel/scheduler_epoch/src/queue.rs

@@ -3,7 +3,7 @@ use alloc::collections::VecDeque;
 use bit_set::BitSet;
 use task::TaskRef;
 
-use crate::{EpochTaskRef, TaskConfiguration, MAX_PRIORITY};
+use crate::{weight, EpochTaskRef, TaskConfiguration, MAX_PRIORITY};
 
 /// A singular run queue.
 ///
@@ -68,7 +68,7 @@ impl RunQueue {
             let mut vec_index = 0;
 
             while !next_task.is_runnable() {
-                vec_index += 1;
+                assert!(!top_queue.is_empty());
 
                 if vec_index + 1 == top_queue.len() {
                     priorities.remove(top_index);
@@ -86,7 +86,7 @@ impl RunQueue {
 
                                 while let Some(mut task) = top_queue.pop_front() {
                                     task.recalculate_tokens(TaskConfiguration {
-                                        priority: top_index as usize,
+                                        weight: weight(top_index),
                                         total_weight,
                                     });
                                     expired.push(task, top_index);
@@ -95,13 +95,17 @@ impl RunQueue {
                                 priorities.remove(top_index);
                             }
 
+                            self.priorities = BitSet::new();
+                            self.len = 0;
+
                             return None;
                         }
                     };
+                    top_queue = &mut self.inner[top_index as usize];
                     vec_index = 0;
                 }
 
-                top_queue = &mut self.inner[top_index as usize];
+                vec_index += 1;
                 next_task = &top_queue[vec_index];
             }
 
@@ -119,7 +123,7 @@ impl RunQueue {
             self.len -= 1;
 
             next_task.recalculate_tokens(TaskConfiguration {
-                priority: top_index as usize,
+                weight: weight(top_index),
                 total_weight,
             });
             expired.push(next_task.clone(), top_index);
@@ -145,9 +149,9 @@ impl RunQueue {
     }
 
     #[inline]
-    pub(crate) fn remove(&mut self, task: &TaskRef) -> bool {
-        for i in self.priorities.iter() {
-            let queue = &mut self.inner[i];
+    pub(crate) fn remove(&mut self, task: &TaskRef) -> Option<usize> {
+        for priority in self.priorities.iter() {
+            let queue = &mut self.inner[priority];
 
             for j in 0..queue.len() {
                 let element = &queue[j];
@@ -157,14 +161,14 @@ impl RunQueue {
                     self.len -= 1;
 
                     if queue.is_empty() {
-                        self.priorities.remove(i as u8);
+                        self.priorities.remove(priority as u8);
                     }
 
-                    return true;
+                    return Some(weight(priority as u8));
                 }
             }
         }
-        false
+        None
     }
 
     /// Returns the priority of the given task.
@@ -187,6 +191,7 @@ impl RunQueue {
     /// run queue.
     #[inline]
     pub(crate) fn set_priority(&mut self, task: &TaskRef, priority: u8) -> bool {
+        // FIXME: Recalculate weights.
         for i in self.priorities.iter() {
             let queue = &mut self.inner[i];
 
@@ -206,7 +211,8 @@ impl RunQueue {
                 }
             }
         }
-        false
+        // false
+        todo!();
     }
 
     #[inline]