k-os/crates/domain/src/scheduler/multi_core.rs

use types::{
    bounded::Priority,
    id::{CoreId, ThreadId},
};

use super::{Scheduler, SchedulerDecision, per_core::PerCoreScheduler};

const MAX_CORES: usize = if cfg!(kani) { 1 } else { 8 };
const MAX_SCHEDULED_THREADS: usize = crate::thread::MAX_THREADS;
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 const fn new(num_cores: u8) -> Self {
        let mut cores = [const { PerCoreScheduler::new(CoreId::from_raw(0)) }; MAX_CORES];
        // Each core numbered by its place (a const fn cannot index or map).
        let mut rest: &mut [PerCoreScheduler] = &mut cores;
        let mut id = 0u8;
        while let Some((core, tail)) = rest.split_first_mut() {
            *core = PerCoreScheduler::new(CoreId::from_raw(id));
            rest = tail;
            id += 1;
        }
        Self {
            cores,
            num_cores: if num_cores < MAX_CORES as u8 {
                num_cores
            } else {
                MAX_CORES as u8
            },
            placements: [const { None }; MAX_SCHEDULED_THREADS],
        }
    }

    fn core(&self, core: CoreId) -> Option<&PerCoreScheduler> {
        self.active().get(core.raw() as usize)
    }

    fn core_mut(&mut self, core: CoreId) -> Option<&mut PerCoreScheduler> {
        let n = self.num_cores as usize;
        self.cores.get_mut(..n)?.get_mut(core.raw() as usize)
    }

    fn active(&self) -> &[PerCoreScheduler] {
        self.cores.get(..self.num_cores as usize).unwrap_or(&[])
    }

    fn placement(&self, thread: ThreadId) -> Option<&ThreadPlacement> {
        self.placements.get(thread.raw() as usize)?.as_ref()
    }

    /// How many cores to schedule on, once the platform knows; before any
    /// thread is placed.
    pub fn set_num_cores(&mut self, num_cores: u8) {
        self.num_cores = num_cores.min(MAX_CORES as u8);
    }

    pub fn num_cores(&self) -> u8 {
        self.num_cores
    }

    /// The highest priority waiting in `core`'s run queue.
    pub fn highest_waiting(&self, core: CoreId) -> Option<Priority> {
        self.core(core)?.highest_waiting_priority()
    }

    /// The core whose run queue holds `thread`, if any.
    pub fn queued_on(&self, thread: ThreadId) -> Option<CoreId> {
        self.active()
            .iter()
            .find(|c| c.is_queued(thread))
            .map(|c| c.core())
    }

    /// Whether `thread` waits in some core's run queue.
    pub fn is_queued(&self, thread: ThreadId) -> bool {
        self.active().iter().any(|c| c.is_queued(thread))
    }

    pub fn current(&self, core: CoreId) -> Option<ThreadId> {
        self.core(core)?.current()
    }

    pub fn set_current(&mut self, core: CoreId, thread: ThreadId) {
        if let Some(c) = self.core_mut(core) {
            c.set_current(thread);
        }
    }

    /// Runs `thread` next on `core` regardless of its place in the queue.
    /// Fails if it is not queued there.
    pub fn run_now(&mut self, core: CoreId, thread: ThreadId) -> bool {
        self.core_mut(core).is_some_and(|c| c.run_now(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(CoreId::from_raw)
            .min_by_key(|&c| self.thread_count_on(c))
    }

    fn most_loaded_core(&self) -> Option<CoreId> {
        (0..self.num_cores)
            .map(CoreId::from_raw)
            .max_by_key(|&c| self.thread_count_on(c))
    }
}

impl Scheduler for MultiCoreScheduler {
    fn add(&mut self, thread: ThreadId, priority: Priority, core: CoreId) {
        if self.core(core).is_none() {
            return;
        }
        let Some(placement) = self.placements.get_mut(thread.raw() as usize) else {
            return;
        };
        *placement = Some(ThreadPlacement { core, priority });
        if let Some(c) = self.core_mut(core) {
            c.add(thread, priority);
        }
    }

    fn remove(&mut self, thread: ThreadId) {
        let Some(placement) = self
            .placements
            .get_mut(thread.raw() as usize)
            .and_then(Option::take)
        else {
            return;
        };
        if let Some(c) = self.core_mut(placement.core) {
            c.remove(thread);
            if c.current() == Some(thread) {
                c.clear_current();
            }
        }
    }

    fn next(&mut self, core: CoreId) -> Option<ThreadId> {
        self.core_mut(core)?.pick_next()
    }

    fn tick(&mut self, core: CoreId, current: ThreadId) -> SchedulerDecision {
        let Some(current_priority) = self.placement(current).map(|p| p.priority) else {
            return SchedulerDecision::Continue;
        };
        if let Some(c) = self.core_mut(core)
            && let Some(highest) = c.highest_waiting_priority()
            && highest > current_priority
            && let Some(next) = c.pick_next()
        {
            c.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 Some((thread, priority)) = self.core_mut(busiest).and_then(|c| c.steal_lowest())
        {
            if let Some(c) = self.core_mut(lightest) {
                c.add(thread, priority);
            }
            if let Some(p) = self.placements.get_mut(thread.raw() as usize) {
                *p = Some(ThreadPlacement {
                    core: lightest,
                    priority,
                });
            }
        }
    }
}