rustubs/proc/
sched.rs

1use crate::arch::x86_64::is_int_enabled;
2use crate::machine::interrupt::{irq_restore, irq_save};
3use crate::proc::sync::*;
4use crate::proc::task::*;
5use alloc::collections::VecDeque;
6use core::sync::atomic::AtomicBool;
7use core::sync::atomic::Ordering;
8pub static GLOBAL_SCHEDULER: L2Sync<Scheduler> = L2Sync::new(Scheduler::new());
9/// A global flag indicating whether reschedule is required.
10pub static NEED_RESCHEDULE: AtomicBool = AtomicBool::new(false);
11
12/// set NEED_RESCHEDULE to true regardless its value; return the previous state.
13#[inline(always)]
14#[allow(non_snake_case)]
15pub fn SET_NEED_RESCHEDULE() -> bool {
16	NEED_RESCHEDULE.swap(true, Ordering::Relaxed)
17}
18
19/// set NEED_RESCHEDULE to false regardless its value; return the previous
20/// state.
21#[inline(always)]
22#[allow(non_snake_case)]
23pub fn CLEAR_NEED_RESCHEDULE() -> bool {
24	NEED_RESCHEDULE.swap(false, Ordering::Relaxed)
25}
26
27pub struct Scheduler {
28	pub run_queue: VecDeque<TaskId>,
29	pub need_schedule: bool,
30}
31
32impl Scheduler {
33	pub const MIN_TASK_CAP: usize = 16;
34	pub const fn new() -> Self {
35		return Self {
36			run_queue: VecDeque::new(),
37			need_schedule: false,
38		};
39	}
40
41	// maybe we reject inserting existing tasks?
42	pub fn insert_task(&mut self, tid: TaskId) {
43		self.run_queue.push_back(tid);
44	}
45
46	pub fn try_remove(&mut self, _tid: TaskId) {
47		todo!("not implemented");
48	}
49
50	/// unsafe because this must be called on a linearization point on Epilogue
51	/// level (l2); It will check the NEED_RESCHEDULE flag.
52	pub unsafe fn try_reschedule() {
53		// this assert doesn't check if you own the L2, but at least a sanity
54		// check.
55		debug_assert!(is_int_enabled());
56		// TODO maybe refine memory ordering here
57		let r = NEED_RESCHEDULE.compare_exchange(
58			true,
59			false,
60			Ordering::Relaxed,
61			Ordering::Relaxed,
62		);
63		if r != Ok(true) {
64			return;
65		}
66		Self::do_schedule();
67	}
68
69	/// do_schedule is only called from epilogue level, so we don't need to lock
70	/// here. For cooperative scheduling call [Self::yield_cpu] instead.
71	pub unsafe fn do_schedule() {
72		let me = Task::current().unwrap();
73		let next_task;
74		let next_tid;
75		{
76			let r = irq_save();
77			// begin L3 critical section
78			// make sure we drop the mutable borrow before doing context swap
79			let sched = GLOBAL_SCHEDULER.get_ref_mut_unguarded();
80			if sched.run_queue.is_empty() && me.state == TaskState::Run {
81				// I'm the only one, just return;
82				irq_restore(r);
83				return;
84			}
85			next_tid = sched.run_queue.pop_front().expect("no runnable task");
86			next_task = next_tid.get_task_ref_mut();
87			debug_assert_eq!(next_task.state, TaskState::Run);
88			match me.state {
89				TaskState::Run => {
90					sched.run_queue.push_back(me.taskid());
91				}
92				TaskState::Dead => {
93					// TODO push the task to cleanup queue
94				}
95				_ => {}
96			}
97			// end L3 critical section
98			irq_restore(r);
99		}
100		if me.taskid() == next_task.taskid() {
101			return;
102		}
103		unsafe {
104			context_swap(
105				&(me.context) as *const _ as u64,
106				&(next_task.context) as *const _ as u64,
107			);
108		}
109	}
110
111	/// guards do_schedule and makes sure it's also sequentialized at L2. Must
112	/// not call this in interrupt context
113	pub fn yield_cpu() {
114		debug_assert!(is_int_enabled());
115		ENTER_L2();
116		unsafe {
117			Self::do_schedule();
118		}
119		LEAVE_L2();
120	}
121
122	/// like do_schedule but we there is no running context to save
123	pub unsafe fn kickoff() {
124		let irq = irq_save();
125		// must not lock the GLOBAL_SCHEDULER here because we never return.
126		// well, the "LEAVE_L2" call in the task entries logically release
127		// the GLOBAL_SCHEDULER but semantically that's too weird
128		let sched = GLOBAL_SCHEDULER.get_ref_mut_unguarded();
129		let tid = sched
130			.run_queue
131			.pop_front()
132			.expect("run queue empty, can't start");
133		let first_task = tid.get_task_ref_mut();
134		irq_restore(irq);
135		// kickoff simulates a do_schedule, so we need to enter l2 here.
136		// new tasks must leave l2 explicitly on their first run
137		ENTER_L2();
138		unsafe {
139			context_swap_to(&(first_task.context) as *const _ as u64);
140		}
141	}
142}
rustubs/proc/sched.rs

rustubs/proc/
sched.rs
