rustubs/proc/
sched.rs

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

rustubs/proc/
sched.rs
