ableos_time/ableos/src/scheduler.rs

use alloc::collections::VecDeque;
use spin::{Lazy, Mutex};

pub static SCHEDULER: Lazy<Mutex<Scheduler>> = Lazy::new(|| Mutex::new(Scheduler::default()));

/// Task's stack size (in kiB)
const STACK_SIZE: usize = 512;

pub enum Task {
    Suspended(u64),
    Spawn { function: fn(), stack_start: usize },
}

#[derive(Default)]
pub struct Scheduler {
    task_queue: VecDeque<Task>,
}

impl Scheduler {
    pub fn enqueue_spawn(&mut self, function: fn()) {
        self.task_queue.push_back(Task::Spawn {
            function,
            stack_start: unsafe {
                let size = STACK_SIZE * 1024;
                alloc::alloc::alloc(
                    alloc::alloc::Layout::from_size_align(size, 1).expect("invalid layout"),
                )
                .add(size - 1) as _
            },
        });
    }

    pub fn enqueue_suspended(&mut self, sp: u64) {
        self.task_queue.push_back(Task::Suspended(sp));
    }

    pub fn pop(&mut self) -> Option<Task> {
        self.task_queue.pop_front()
    }
}