Try to keep tasks on the same e-core

2 files changed, 6 insertions, 5 deletions

diff --git a/src/e_core_selector.rs b/src/e_core_selector.rs
index 6952e80..251cbe8 100644
--- a/src/e_core_selector.rs
+++ b/src/e_core_selector.rs
@@ -1,10 +1,9 @@
 use std::ops::Range;
 
 pub trait ECoreSelector {
-    fn next_core(&mut self) -> i32;
+    fn next_core(&mut self, previous_cpu: i32) -> i32;
 }
 
-//TODO: better core selection e.g. try to keep task on the same core
 pub struct RoundRobinSelector {
     offset: u32,
     num_cores: u32,
@@ -22,8 +21,10 @@ impl RoundRobinSelector {
 }
 
 impl ECoreSelector for RoundRobinSelector {
-    // TODO: use task cpu if suitable
-    fn next_core(&mut self) -> i32 {
+    fn next_core(&mut self, previous_cpu: i32) -> i32 {
+        if (self.offset..(self.offset + self.num_cores)).contains(&(previous_cpu as u32)) {
+            return previous_cpu;
+        }
         self.last_used += 1;
         (self.offset + (self.last_used % self.num_cores)) as i32
     }
diff --git a/src/scheduler.rs b/src/scheduler.rs
index 7645498..8618b78 100644
--- a/src/scheduler.rs
+++ b/src/scheduler.rs
@@ -149,7 +149,7 @@ impl<'a> Scheduler<'a> {
             let mut dispatched_task = DispatchedTask::new(&task);
 
             // Low budget tasks go to e-cores
-            let cpu = self.e_core_selector.next_core();
+            let cpu = self.e_core_selector.next_core(task.cpu);
             if cpu >= 0 {
                 dispatched_task.cpu = cpu;
             } else {