120 lines
3.9 KiB
Rust
120 lines
3.9 KiB
Rust
use types::{
|
|
bounded::Priority,
|
|
id::{CoreId, ThreadId},
|
|
};
|
|
|
|
use super::{Scheduler, SchedulerDecision, per_core::PerCoreScheduler};
|
|
|
|
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],
|
|
}
|
|
}
|
|
|
|
pub fn current(&self, core: CoreId) -> Option<ThreadId> {
|
|
self.cores[core.raw() as usize].current()
|
|
}
|
|
|
|
pub fn set_current(&mut self, core: CoreId, thread: ThreadId) {
|
|
self.cores[core.raw() as usize].set_current(thread);
|
|
}
|
|
|
|
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) -> Option<CoreId> {
|
|
(0..self.num_cores)
|
|
.map(|i| CoreId::from_raw(i))
|
|
.min_by_key(|&c| self.thread_count_on(c))
|
|
}
|
|
|
|
fn most_loaded_core(&self) -> Option<CoreId> {
|
|
(0..self.num_cores)
|
|
.map(|i| CoreId::from_raw(i))
|
|
.max_by_key(|&c| self.thread_count_on(c))
|
|
}
|
|
}
|
|
|
|
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() {
|
|
let core_idx = placement.core.raw() as usize;
|
|
self.cores[core_idx].remove(thread, placement.priority);
|
|
if self.cores[core_idx].current() == Some(thread) {
|
|
self.cores[core_idx].clear_current();
|
|
}
|
|
}
|
|
}
|
|
|
|
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 {
|
|
if let Some(next) = self.cores[core_idx].pick_next() {
|
|
self.cores[core_idx].add(current, current_priority);
|
|
return SchedulerDecision::Preempt(next);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
SchedulerDecision::Continue
|
|
}
|
|
|
|
fn balance(&mut self) {
|
|
let (busiest, lightest) = match (self.most_loaded_core(), self.least_loaded_core()) {
|
|
(Some(b), Some(l)) => (b, l),
|
|
_ => return,
|
|
};
|
|
|
|
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,
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|