feat(domain): add MultiCoreScheduler with placement tracking, priority preemption, and load balancing

crates/domain/src/scheduler/mod.rs

@@ -3,6 +3,7 @@ use types::{
     id::{CoreId, ThreadId},
 };
 
+pub mod multi_core;
 pub mod per_core;
 pub mod priority_queue;
 

crates/domain/src/scheduler/multi_core.rs

@@ -0,0 +1,106 @@
+use types::{
+    bounded::Priority,
+    id::{CoreId, ThreadId},
+};
+
+use super::{per_core::PerCoreScheduler, Scheduler, SchedulerDecision};
+
+const MAX_CORES: usize = 8;
+const MAX_SCHEDULED_THREADS: usize = 256;
+const BALANCE_THRESHOLD: usize = 2;
+
+struct ThreadPlacement {
+    core: CoreId,
+    priority: Priority,
+}
+
+pub struct MultiCoreScheduler {
+    cores: [PerCoreScheduler; MAX_CORES],
+    num_cores: u8,
+    placements: [Option<ThreadPlacement>; MAX_SCHEDULED_THREADS],
+}
+
+impl MultiCoreScheduler {
+    pub fn new(num_cores: u8) -> Self {
+        let cores = core::array::from_fn(|i| PerCoreScheduler::new(CoreId::from_raw(i as u8)));
+        Self {
+            cores,
+            num_cores,
+            placements: [const { None }; MAX_SCHEDULED_THREADS],
+        }
+    }
+
+    fn thread_count_on(&self, core: CoreId) -> usize {
+        self.placements
+            .iter()
+            .filter(|p| p.as_ref().is_some_and(|tp| tp.core == core))
+            .count()
+    }
+
+    fn least_loaded_core(&self) -> CoreId {
+        (0..self.num_cores)
+            .map(|i| CoreId::from_raw(i))
+            .min_by_key(|&c| self.thread_count_on(c))
+            .unwrap()
+    }
+
+    fn most_loaded_core(&self) -> CoreId {
+        (0..self.num_cores)
+            .map(|i| CoreId::from_raw(i))
+            .max_by_key(|&c| self.thread_count_on(c))
+            .unwrap()
+    }
+}
+
+impl Scheduler for MultiCoreScheduler {
+    fn add(&mut self, thread: ThreadId, priority: Priority, core: CoreId) {
+        let idx = thread.raw() as usize;
+        self.placements[idx] = Some(ThreadPlacement { core, priority });
+        self.cores[core.raw() as usize].add(thread, priority);
+    }
+
+    fn remove(&mut self, thread: ThreadId) {
+        let idx = thread.raw() as usize;
+        if let Some(placement) = self.placements[idx].take() {
+            self.cores[placement.core.raw() as usize].remove(thread, placement.priority);
+        }
+    }
+
+    fn next(&mut self, core: CoreId) -> Option<ThreadId> {
+        self.cores[core.raw() as usize].pick_next()
+    }
+
+    fn tick(&mut self, core: CoreId, current: ThreadId) -> SchedulerDecision {
+        let idx = current.raw() as usize;
+        if let Some(placement) = &self.placements[idx] {
+            let current_priority = placement.priority;
+            let core_idx = core.raw() as usize;
+            if let Some(highest) = self.cores[core_idx].highest_waiting_priority() {
+                if highest > current_priority {
+                    let next = self.cores[core_idx].pick_next().unwrap();
+                    self.cores[core_idx].add(current, current_priority);
+                    return SchedulerDecision::Preempt(next);
+                }
+            }
+        }
+        SchedulerDecision::Continue
+    }
+
+    fn balance(&mut self) {
+        let busiest = self.most_loaded_core();
+        let lightest = self.least_loaded_core();
+
+        if self.thread_count_on(busiest) > self.thread_count_on(lightest) + BALANCE_THRESHOLD {
+            let busiest_idx = busiest.raw() as usize;
+            if let Some((thread, priority)) = self.cores[busiest_idx].steal_lowest() {
+                let lightest_idx = lightest.raw() as usize;
+                self.cores[lightest_idx].add(thread, priority);
+                let tid = thread.raw() as usize;
+                self.placements[tid] = Some(ThreadPlacement {
+                    core: lightest,
+                    priority,
+                });
+            }
+        }
+    }
+}

crates/domain/src/scheduler/per_core.rs

@@ -43,10 +43,15 @@ impl PerCoreScheduler {
         self.current = None;
     }
 
+    pub fn highest_waiting_priority(&self) -> Option<Priority> {
+        self.queue.highest_priority()
+    }
+
+    pub fn steal_lowest(&mut self) -> Option<(ThreadId, Priority)> {
+        self.queue.steal_lowest()
+    }
+
     pub fn len(&self) -> usize {
-        // for load balancing — count of ready threads
-        // could track this incrementally instead of computing
-        // but for now, bitmap popcount works
         self.queue.bitmap().count_ones() as usize
     }
 }

crates/domain/src/scheduler/priority_queue.rs

@@ -102,6 +102,26 @@ impl PriorityQueue {
         Some((thread, Priority::new(idx as u8).unwrap()))
     }
 
+    pub fn highest_priority(&self) -> Option<Priority> {
+        if self.bitmap == 0 {
+            return None;
+        }
+        let idx = 31 - self.bitmap.leading_zeros() as usize;
+        Some(Priority::new(idx as u8).unwrap())
+    }
+
+    pub fn steal_lowest(&mut self) -> Option<(ThreadId, Priority)> {
+        if self.bitmap == 0 {
+            return None;
+        }
+        let idx = self.bitmap.trailing_zeros() as usize;
+        let thread = self.queues[idx].pop()?;
+        if self.queues[idx].is_empty() {
+            self.bitmap &= !(1 << idx);
+        }
+        Some((thread, Priority::new(idx as u8).unwrap()))
+    }
+
     pub fn is_empty(&self) -> bool {
         self.bitmap == 0
     }

crates/domain/src/scheduler/tests.rs

@@ -1,8 +1,14 @@
+use super::multi_core::MultiCoreScheduler;
 use super::per_core::PerCoreScheduler;
 use super::priority_queue::PriorityQueue;
+use super::{Scheduler, SchedulerDecision};
 use types::bounded::Priority;
 use types::id::{CoreId, ThreadId};
 
+fn core(id: u8) -> CoreId {
+    CoreId::from_raw(id)
+}
+
 fn thread(id: u16) -> ThreadId {
     ThreadId::from_raw(id)
 }
@@ -122,3 +128,115 @@ fn clear_current() {
     core.clear_current();
     assert!(core.current().is_none());
 }
+
+// --- PriorityQueue new methods ---
+
+#[test]
+fn highest_priority_peeks_without_popping() {
+    let mut q = PriorityQueue::new();
+    q.add(thread(1), prio(5));
+    q.add(thread(2), prio(20));
+    assert_eq!(q.highest_priority(), Some(prio(20)));
+    assert_eq!(q.next().unwrap().0, thread(2));
+}
+
+#[test]
+fn highest_priority_empty_returns_none() {
+    let q = PriorityQueue::new();
+    assert_eq!(q.highest_priority(), None);
+}
+
+#[test]
+fn steal_lowest_takes_lowest_priority_thread() {
+    let mut q = PriorityQueue::new();
+    q.add(thread(1), prio(5));
+    q.add(thread(2), prio(20));
+    let (t, p) = q.steal_lowest().unwrap();
+    assert_eq!(t, thread(1));
+    assert_eq!(p, prio(5));
+    assert_eq!(q.next().unwrap().0, thread(2));
+}
+
+#[test]
+fn steal_lowest_empty_returns_none() {
+    let mut q = PriorityQueue::new();
+    assert!(q.steal_lowest().is_none());
+}
+
+// --- MultiCoreScheduler tests ---
+
+#[test]
+fn multi_add_and_next() {
+    let mut s = MultiCoreScheduler::new(2);
+    s.add(thread(1), prio(10), core(0));
+    assert_eq!(s.next(core(0)), Some(thread(1)));
+}
+
+#[test]
+fn multi_remove_then_next_returns_none() {
+    let mut s = MultiCoreScheduler::new(1);
+    s.add(thread(1), prio(10), core(0));
+    s.remove(thread(1));
+    assert_eq!(s.next(core(0)), None);
+}
+
+#[test]
+fn multi_tick_continue_when_no_higher_priority() {
+    let mut s = MultiCoreScheduler::new(1);
+    s.add(thread(1), prio(10), core(0));
+    let current = s.next(core(0)).unwrap();
+    s.add(thread(2), prio(5), core(0));
+    assert!(matches!(
+        s.tick(core(0), current),
+        SchedulerDecision::Continue
+    ));
+}
+
+#[test]
+fn multi_tick_preempts_for_higher_priority() {
+    let mut s = MultiCoreScheduler::new(1);
+    s.add(thread(1), prio(5), core(0));
+    let current = s.next(core(0)).unwrap();
+    s.add(thread(2), prio(20), core(0));
+    match s.tick(core(0), current) {
+        SchedulerDecision::Preempt(next) => assert_eq!(next, thread(2)),
+        SchedulerDecision::Continue => panic!("should have preempted"),
+    }
+}
+
+#[test]
+fn multi_tick_preempt_returns_current_to_queue() {
+    let mut s = MultiCoreScheduler::new(1);
+    s.add(thread(1), prio(5), core(0));
+    let current = s.next(core(0)).unwrap();
+    s.add(thread(2), prio(20), core(0));
+    s.tick(core(0), current); // preempts, thread(1) goes back to queue
+    // thread(2) is now running (was picked by tick), thread(1) should be next
+    assert_eq!(s.next(core(0)), Some(thread(1)));
+}
+
+#[test]
+fn multi_balance_moves_thread_from_overloaded_core() {
+    let mut s = MultiCoreScheduler::new(2);
+    for i in 0..4 {
+        s.add(thread(i), prio(10), core(0));
+    }
+    s.balance();
+    assert!(s.next(core(1)).is_some());
+}
+
+#[test]
+fn multi_next_on_empty_core_returns_none() {
+    let mut s = MultiCoreScheduler::new(2);
+    s.add(thread(1), prio(10), core(0));
+    assert_eq!(s.next(core(1)), None);
+}
+
+#[test]
+fn multi_threads_on_different_cores() {
+    let mut s = MultiCoreScheduler::new(2);
+    s.add(thread(1), prio(10), core(0));
+    s.add(thread(2), prio(10), core(1));
+    assert_eq!(s.next(core(0)), Some(thread(1)));
+    assert_eq!(s.next(core(1)), Some(thread(2)));
+}