ahrightos_process/
scheduler.rs1use crate::{ProcessControlBlock, ProcessId, ProcessState};
2use alloc::boxed::Box;
3use alloc::collections::VecDeque;
4use alloc::vec::Vec;
5
6pub struct ReadyQueue {
7 queue: VecDeque<Box<ProcessControlBlock>>,
8}
9
10impl ReadyQueue {
11 pub fn new() -> Self {
12 Self {
13 queue: VecDeque::new(),
14 }
15 }
16
17 pub fn push_back(&mut self, process: Box<ProcessControlBlock>) {
18 self.queue.push_back(process);
19 }
20
21 pub fn pop_front(&mut self) -> Option<Box<ProcessControlBlock>> {
22 self.queue.pop_front()
23 }
24
25 pub fn is_empty(&self) -> bool {
26 self.queue.is_empty()
27 }
28
29 pub fn len(&self) -> usize {
30 self.queue.len()
31 }
32
33 pub fn retain(&mut self, f: impl Fn(&ProcessControlBlock) -> bool) {
34 self.queue.retain(|p| f(p));
35 }
36
37 pub fn iter(&self) -> impl Iterator<Item = &ProcessControlBlock> {
38 self.queue.iter().map(|b| b.as_ref())
39 }
40
41 pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut ProcessControlBlock> {
42 self.queue.iter_mut().map(|b| b.as_mut())
43 }
44}
45
46impl Default for ReadyQueue {
47 fn default() -> Self {
48 Self::new()
49 }
50}
51
52#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub struct SchedulerStats {
54 pub current_pid: Option<ProcessId>,
55 pub ready: usize,
56 pub blocked: usize,
57 pub zombies: usize,
58 pub remaining_ticks: usize,
59 pub total: usize,
60}
61
62pub struct RoundRobinScheduler {
63 ready_queue: ReadyQueue,
64 blocked_queue: VecDeque<Box<ProcessControlBlock>>,
65 current: Option<Box<ProcessControlBlock>>,
66 zombies: Vec<ProcessControlBlock>,
67 time_slice: usize,
68 remaining_ticks: usize,
69}
70
71impl RoundRobinScheduler {
72 pub fn new(time_slice: usize) -> Self {
73 let time_slice = time_slice.max(1);
74 Self {
75 ready_queue: ReadyQueue::new(),
76 blocked_queue: VecDeque::new(),
77 current: None,
78 zombies: Vec::with_capacity(8),
79 time_slice,
80 remaining_ticks: time_slice,
81 }
82 }
83
84 pub fn add_process(&mut self, process: Box<ProcessControlBlock>) -> Result<(), &'static str> {
85 if process.state() != ProcessState::Ready {
86 return Err("Process must be in Ready state");
87 }
88 self.ready_queue.push_back(process);
89 Ok(())
90 }
91
92 fn park_process(&mut self, mut process: Box<ProcessControlBlock>) {
93 match process.state() {
94 ProcessState::Running => {
95 process
96 .set_state(ProcessState::Ready)
97 .expect("running process can return to ready");
98 self.ready_queue.push_back(process);
99 }
100 ProcessState::Ready => self.ready_queue.push_back(process),
101 ProcessState::Blocked => self.blocked_queue.push_back(process),
102 ProcessState::Zombie => self.zombies.push(*process),
103 }
104 }
105
106 fn select_next_ready(&mut self) -> Option<&mut ProcessControlBlock> {
107 let mut pcb = self.ready_queue.pop_front()?;
108 pcb.set_state(ProcessState::Running)
109 .expect("ready process can be scheduled");
110 self.remaining_ticks = self.time_slice;
111 self.current = Some(pcb);
112 self.current.as_deref_mut()
113 }
114
115 pub fn schedule(&mut self) -> Option<&mut ProcessControlBlock> {
116 if let Some(proc) = self.current.take() {
117 self.park_process(proc);
118 }
119 self.select_next_ready()
120 }
121
122 pub fn yield_current(&mut self) {
123 let _ = self.schedule();
124 }
125
126 pub fn tick(&mut self) {
127 if self.current.is_none() {
128 let _ = self.schedule();
129 return;
130 }
131 if self.remaining_ticks > 0 {
132 self.remaining_ticks -= 1;
133 }
134 if self.remaining_ticks == 0 {
135 let _ = self.schedule();
136 }
137 }
138
139 pub fn current(&self) -> Option<&ProcessControlBlock> {
140 self.current.as_deref()
141 }
142
143 pub fn current_mut(&mut self) -> Option<&mut ProcessControlBlock> {
144 self.current.as_deref_mut()
145 }
146
147 pub fn ready_count(&self) -> usize {
148 self.ready_queue.len()
149 }
150
151 pub fn blocked_count(&self) -> usize {
152 self.blocked_queue.len()
153 }
154
155 pub fn zombie_count(&self) -> usize {
156 self.zombies.len()
157 }
158
159 pub fn stats(&self) -> SchedulerStats {
160 let current = usize::from(self.current.is_some());
161 let ready = self.ready_queue.len();
162 let blocked = self.blocked_queue.len();
163 let zombies = self.zombies.len();
164
165 SchedulerStats {
166 current_pid: self.current.as_ref().map(|p| p.pid()),
167 ready,
168 blocked,
169 zombies,
170 remaining_ticks: self.remaining_ticks,
171 total: current + ready + blocked + zombies,
172 }
173 }
174
175 pub fn is_idle(&self) -> bool {
176 self.current.is_none() && self.ready_queue.is_empty()
177 }
178
179 pub fn find_zombie_child(&self, parent_pid: ProcessId) -> Option<(ProcessId, i32)> {
180 for pcb in &self.zombies {
181 if pcb.parent_pid() == Some(parent_pid) && pcb.state() == ProcessState::Zombie {
182 return Some((pcb.pid(), pcb.exit_code().unwrap_or(0)));
183 }
184 }
185 None
186 }
187
188 pub fn remove_zombie(&mut self, pid: ProcessId) -> bool {
189 if let Some(pos) = self.zombies.iter().position(|p| p.pid() == pid) {
190 self.zombies.remove(pos);
191 true
192 } else {
193 false
194 }
195 }
196
197 pub fn has_living_children(&self, parent_pid: ProcessId) -> bool {
198 let in_current = self
199 .current
200 .as_ref()
201 .is_some_and(|p| p.parent_pid() == Some(parent_pid));
202 let in_ready = self
203 .ready_queue
204 .iter()
205 .any(|p| p.parent_pid() == Some(parent_pid));
206 let in_blocked = self
207 .blocked_queue
208 .iter()
209 .any(|p| p.parent_pid() == Some(parent_pid));
210 in_current || in_ready || in_blocked
211 }
212
213 pub fn unblock_process(&mut self, pid: ProcessId) -> bool {
214 if let Some(pos) = self.blocked_queue.iter().position(|p| p.pid() == pid) {
215 if let Some(mut pcb) = self.blocked_queue.remove(pos) {
216 pcb.set_state(ProcessState::Ready)
217 .expect("blocked process can be readied");
218 self.ready_queue.push_back(pcb);
219 return true;
220 }
221 }
222 false
223 }
224
225 pub fn all_processes(&self) -> Vec<&ProcessControlBlock> {
226 let mut out = Vec::new();
227 if let Some(ref p) = self.current {
228 out.push(p.as_ref());
229 }
230 for p in self.ready_queue.iter() {
231 out.push(p);
232 }
233 for p in &self.blocked_queue {
234 out.push(p.as_ref());
235 }
236 for p in &self.zombies {
237 out.push(p);
238 }
239 out
240 }
241
242 pub fn block_current(&mut self) {
243 if let Some(ref mut pcb) = self.current {
244 pcb.set_state(ProcessState::Blocked)
245 .expect("running process can block");
246 }
247 }
248
249 pub fn pick_next_ready(&mut self) -> Option<&ProcessControlBlock> {
250 if let Some(old) = self.current.take() {
251 self.park_process(old);
252 }
253 self.select_next_ready().map(|p| &*p)
254 }
255}
256
257#[cfg(test)]
258mod tests {
259 use super::*;
260
261 fn process(pid: usize, state: ProcessState, parent: Option<usize>) -> Box<ProcessControlBlock> {
262 let mut pcb = ProcessControlBlock::new(
263 ProcessId::from_usize(pid),
264 "test",
265 parent.map(ProcessId::from_usize),
266 );
267 pcb.set_state_unchecked(state);
268 Box::new(pcb)
269 }
270
271 #[test]
272 fn tick_switches_after_configured_time_slice() {
273 let mut scheduler = RoundRobinScheduler::new(2);
274 scheduler
275 .add_process(process(40, ProcessState::Ready, None))
276 .unwrap();
277 scheduler
278 .add_process(process(41, ProcessState::Ready, None))
279 .unwrap();
280
281 scheduler.schedule();
282 scheduler.tick();
283 assert_eq!(scheduler.current().map(|p| p.pid().as_usize()), Some(40));
284 scheduler.tick();
285 assert_eq!(scheduler.current().map(|p| p.pid().as_usize()), Some(41));
286 }
287
288 #[test]
289 fn blocked_process_can_be_unblocked() {
290 let mut scheduler = RoundRobinScheduler::new(2);
291 scheduler
292 .add_process(process(20, ProcessState::Ready, None))
293 .unwrap();
294 scheduler
295 .add_process(process(21, ProcessState::Ready, None))
296 .unwrap();
297
298 scheduler.schedule();
299 scheduler.block_current();
300 scheduler.pick_next_ready();
301
302 assert!(scheduler.unblock_process(ProcessId::from_usize(20)));
303 assert_eq!(scheduler.blocked_count(), 0);
304 assert_eq!(scheduler.ready_count(), 1);
305 }
306
307 #[test]
308 fn zombie_current_is_retained_for_wait() {
309 let parent = ProcessId::from_usize(99);
310 let mut scheduler = RoundRobinScheduler::new(2);
311 scheduler
312 .add_process(process(30, ProcessState::Ready, Some(parent.as_usize())))
313 .unwrap();
314 scheduler
315 .add_process(process(31, ProcessState::Ready, None))
316 .unwrap();
317
318 scheduler.schedule().unwrap().set_exit_code(7);
319 scheduler.schedule();
320
321 assert_eq!(
322 scheduler.find_zombie_child(parent),
323 Some((ProcessId::from_usize(30), 7))
324 );
325 assert!(scheduler.remove_zombie(ProcessId::from_usize(30)));
326 }
327
328 #[test]
329 fn living_child_query_excludes_reapable_zombies() {
330 let parent = ProcessId::from_usize(77);
331 let mut scheduler = RoundRobinScheduler::new(1);
332 scheduler
333 .add_process(process(78, ProcessState::Ready, Some(parent.as_usize())))
334 .unwrap();
335
336 scheduler.schedule().unwrap().set_exit_code(9);
337 scheduler.schedule();
338
339 assert_eq!(
340 scheduler.find_zombie_child(parent),
341 Some((ProcessId::from_usize(78), 9))
342 );
343 assert!(!scheduler.has_living_children(parent));
344 }
345}