diff --git a/crates/domain/src/scheduler/mod.rs b/crates/domain/src/scheduler/mod.rs index b89dbb6..8d5934f 100644 --- a/crates/domain/src/scheduler/mod.rs +++ b/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; diff --git a/crates/domain/src/scheduler/multi_core.rs b/crates/domain/src/scheduler/multi_core.rs new file mode 100644 index 0000000..3174ac4 --- /dev/null +++ b/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; 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 { + 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, + }); + } + } + } +} diff --git a/crates/domain/src/scheduler/per_core.rs b/crates/domain/src/scheduler/per_core.rs index 2ad3d06..b7d504c 100644 --- a/crates/domain/src/scheduler/per_core.rs +++ b/crates/domain/src/scheduler/per_core.rs @@ -43,10 +43,15 @@ impl PerCoreScheduler { self.current = None; } + pub fn highest_waiting_priority(&self) -> Option { + 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 } } diff --git a/crates/domain/src/scheduler/priority_queue.rs b/crates/domain/src/scheduler/priority_queue.rs index d80b18e..9d5dc13 100644 --- a/crates/domain/src/scheduler/priority_queue.rs +++ b/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 { + 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 } diff --git a/crates/domain/src/scheduler/tests.rs b/crates/domain/src/scheduler/tests.rs index 7606590..ea3bb03 100644 --- a/crates/domain/src/scheduler/tests.rs +++ b/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))); +}