kernel/sched/
scheduler.rs

1//! Scheduler module
2//! 
3//! The scheduler module is responsible for scheduling tasks on the CPU.
4//! Currently, the scheduler is a simple round-robin scheduler.
5
6extern crate alloc;
7
8use core::panic;
9
10use alloc::{collections::vec_deque::VecDeque, string::ToString};
11
12use crate::{arch::{enable_interrupt, get_cpu, get_user_trap_handler, instruction::idle, set_trapframe, set_trapvector, Arch}, environment::NUM_OF_CPUS, task::{new_kernel_task, TaskState}, timer::get_kernel_timer, vm::{get_kernel_vm_manager, get_trampoline_trap_vector, get_trampoline_trapframe}};
13use crate::println;
14use crate::print;
15
16
17use super::dispatcher::Dispatcher;
18use crate::task::Task;
19
20static mut SCHEDULER: Option<Scheduler> = None;
21
22pub fn get_scheduler() -> &'static mut Scheduler {
23    unsafe {
24        match SCHEDULER {
25            Some(ref mut s) => s,
26            None => {
27                SCHEDULER = Some(Scheduler::new());
28                get_scheduler()
29            }
30        }
31    }
32}
33
34pub struct Scheduler {
35    task_queue: [VecDeque<Task>; NUM_OF_CPUS],
36    add_req_queue: [VecDeque<Task>; NUM_OF_CPUS],
37    dispatcher: [Dispatcher; NUM_OF_CPUS],
38    interval: u64, /* in microseconds */
39    current_task_id: [Option<usize>; NUM_OF_CPUS],
40}
41
42impl Scheduler {
43    pub const fn new() -> Self {
44        Scheduler {
45            task_queue: [const { VecDeque::new() }; NUM_OF_CPUS],
46            add_req_queue: [const { VecDeque::new() }; NUM_OF_CPUS],
47            dispatcher: [const { Dispatcher::new() }; NUM_OF_CPUS],
48            interval: 10000, /* 1ms */
49            current_task_id: [const { None }; NUM_OF_CPUS],
50        }
51    }
52
53    pub fn add_task(&mut self, task: Task, cpu_id: usize) {
54        self.task_queue[cpu_id].push_back(task);
55    }
56
57    fn run(&mut self, cpu: &mut Arch) {
58        let cpu_id = cpu.get_cpuid();
59
60        if !self.add_req_queue[cpu_id].is_empty() {
61            for task in self.add_req_queue[cpu_id].drain(..) {
62                self.task_queue[cpu_id].push_back(task);
63            }
64        }
65
66        loop {
67            let task = self.task_queue[cpu_id].pop_front();
68
69            /* If there are no subsequent tasks */
70            if self.task_queue[cpu_id].is_empty() {
71                match task {
72                    Some(mut t) => {
73                        match t.state {
74                            TaskState::Zombie => {
75                                panic!("At least one task must be scheduled");
76                            },
77                            TaskState::Terminated => {
78                                panic!("At least one task must be scheduled");
79                            },
80                            _ => {
81                                if self.current_task_id[cpu_id].is_none() {
82                                    self.dispatcher[cpu_id].dispatch(cpu, &mut t, None);
83                                }
84                                self.current_task_id[cpu_id] = Some(t.get_id());
85                                self.task_queue[cpu_id].push_back(t);
86                                break;
87                            }
88                        }
89                    }
90                    /* If the task queue is empty, run the same task again */
91                    None => break,
92                }
93            } else {
94                match task {
95                    Some(mut t) => {
96                        match t.state {
97                            TaskState::Zombie => {
98                                self.task_queue[cpu_id].push_back(t);
99                                continue;
100                            },
101                            TaskState::Terminated => {
102                                continue;
103                            },
104                            _ => {
105                                let prev_task = match self.current_task_id[cpu_id] {
106                                    Some(task_id) => self.task_queue[cpu_id].iter_mut().find(|t| t.get_id() == task_id),
107                                    None => None
108                                };
109                                if prev_task.is_some() {
110                                    self.dispatcher[cpu_id].dispatch(cpu, &mut t, prev_task);
111                                }
112                                self.current_task_id[cpu_id] = Some(t.get_id());
113                                self.task_queue[cpu_id].push_back(t);
114                                break;
115                            }
116                        }
117                    }
118                    None => break,
119                }
120            }
121        }
122    }
123
124    pub fn schedule(&mut self, cpu: &mut Arch) {
125        let cpu_id = cpu.get_cpuid();
126
127        let timer = get_kernel_timer();
128        timer.stop(cpu_id);
129        timer.set_interval_us(cpu_id, self.interval);
130
131        if !self.task_queue[cpu_id].is_empty() {
132            self.run(cpu);
133            timer.start(cpu_id);
134        }
135    }
136
137    /* MUST NOT raise any exception in this function before the idle loop */
138    pub fn start_scheduler(&mut self) {
139        let cpu = get_cpu();
140        let cpu_id = cpu.get_cpuid();
141        let timer = get_kernel_timer();
142        timer.stop(cpu_id);
143
144        let trap_vector = get_trampoline_trap_vector();
145        let trapframe = get_trampoline_trapframe(cpu_id);
146        set_trapvector(trap_vector);
147        set_trapframe(trapframe);
148        // cpu.get_trapframe().set_trap_handler(get_user_trap_handler());
149        let trapframe = cpu.get_trapframe();
150        trapframe.set_trap_handler(get_user_trap_handler());
151        trapframe.set_next_address_space(get_kernel_vm_manager().get_asid());
152
153        /* Jump to trap handler immediately */
154        timer.set_interval_us(cpu_id, 0);
155        enable_interrupt();
156        timer.start(cpu_id);
157        idle();
158    }
159
160    pub fn get_current_task(&mut self, cpu_id: usize) -> Option<&mut Task> {
161        match self.current_task_id[cpu_id] {
162            Some(task_id) => self.task_queue[cpu_id].iter_mut().find(|t| t.get_id() == task_id),
163            None => None
164        }
165    }
166
167    /// Returns a mutable reference to the task with the specified ID, if found.
168    /// 
169    /// # Arguments
170    /// * `task_id` - The ID of the task to search for.
171    /// 
172    /// # Returns
173    /// A mutable reference to the task if found, or None otherwise.
174    pub fn get_task_by_id(&mut self, task_id: usize) -> Option<&mut Task> {
175        for task_queue in self.task_queue.iter_mut() {
176            if let Some(task) = task_queue.iter_mut().find(|t| t.get_id() == task_id) {
177                return Some(task);
178            }
179        }
180        None
181    }
182}
183
184pub fn make_test_tasks() {
185    println!("Making test tasks...");
186    let sched = get_scheduler();
187    let mut task0 = new_kernel_task("Task0".to_string(), 0, || {
188        println!("Task0");
189        let mut counter: usize = 0;
190        loop {
191            if counter % 500000 == 0 {
192                print!("\nTask0: ");
193            }
194            if counter % 10000 == 0 {
195                print!(".");
196            }
197            counter += 1;
198            if counter >= 100000000 {
199                break;
200            }
201        }
202        println!("");
203        println!("Task0: Done");
204        idle();
205    });
206    task0.init();
207    sched.add_task(task0, 0);
208
209    let mut task1 = new_kernel_task("Task1".to_string(), 0, || {
210        println!("Task1");
211        let mut counter: usize = 0;
212        loop {
213            if counter % 500000 == 0 {
214                print!("\nTask1: {} %", counter / 1000000);
215            }
216            counter += 1;
217            if counter >= 100000000 {
218                break;
219            }
220        }
221        println!("\nTask1: 100 %");
222        println!("Task1: Completed");
223        idle();
224    });
225    task1.init();
226    sched.add_task(task1, 0);
227
228    let mut task2 = new_kernel_task("Task2".to_string(), 0, || {
229        println!("Task2");
230        /* Fizz Buzz */
231        for i in 1..=1000000 {
232            if i % 1000 > 0 {
233                continue;
234            }
235            let c = i / 1000;
236            if c % 15 == 0 {
237                println!("FizzBuzz");
238            } else if c % 3 == 0 {
239                println!("Fizz");
240            } else if c % 5 == 0 {
241                println!("Buzz");
242            } else {
243                println!("{}", c);
244            }
245        }
246        println!("Task2: Done");
247        idle();
248    });
249    task2.init();
250    sched.add_task(task2, 0);
251}
252
253// late_initcall!(make_test_tasks);
254
255#[cfg(test)]
256mod tests {
257    use crate::task::TaskType;
258
259    use super::*;
260
261    #[test_case]
262    fn test_add_task() {
263        let mut scheduler = Scheduler::new();
264        let task = Task::new("TestTask".to_string(), 1, TaskType::Kernel);
265        scheduler.add_task(task, 0);
266        assert_eq!(scheduler.task_queue[0].len(), 1);
267    }
268}
kernel/sched/scheduler.rs

kernel/sched/
scheduler.rs
