fix(core_local): put scheduler into RefCell

Signed-off-by: Martin Kröning <[email protected]>

src/arch/aarch64/kernel/core_local.rs

@@ -1,7 +1,7 @@
 use alloc::boxed::Box;
 use alloc::vec::Vec;
 use core::arch::asm;
-use core::cell::{Cell, RefCell, RefMut};
+use core::cell::{RefCell, RefMut};
 use core::ptr;
 use core::sync::atomic::Ordering;
 
@@ -20,7 +20,7 @@ pub(crate) struct CoreLocal {
 	/// ID of the current Core.
 	core_id: CoreId,
 	/// Scheduler of the current Core.
-	scheduler: Cell<*mut PerCoreScheduler>,
+	scheduler: RefCell<Option<PerCoreScheduler>>,
 	/// Interface to the interrupt counters
 	irq_statistics: &'static IrqStatistics,
 	/// Queue of async tasks
@@ -44,7 +44,7 @@ impl CoreLocal {
 		let this = Self {
 			this: ptr::null_mut(),
 			core_id,
-			scheduler: Cell::new(ptr::null_mut()),
+			scheduler: RefCell::new(None),
 			irq_statistics,
 			async_tasks: RefCell::new(Vec::new()),
 			#[cfg(feature = "smp")]
@@ -91,17 +91,18 @@ pub(crate) fn core_id() -> CoreId {
 	}
 }
 
-#[inline]
-pub(crate) fn core_scheduler() -> &'static mut PerCoreScheduler {
-	unsafe { &mut *CoreLocal::get().scheduler.get() }
+pub(crate) fn core_scheduler() -> RefMut<'static, PerCoreScheduler> {
+	RefMut::map(CoreLocal::get().scheduler.borrow_mut(), |scheduler| {
+		scheduler.as_mut().unwrap()
+	})
 }
 
 pub(crate) fn async_tasks() -> RefMut<'static, Vec<AsyncTask>> {
 	CoreLocal::get().async_tasks.borrow_mut()
 }
 
-pub(crate) fn set_core_scheduler(scheduler: *mut PerCoreScheduler) {
-	CoreLocal::get().scheduler.set(scheduler);
+pub(crate) fn set_core_scheduler(scheduler: PerCoreScheduler) {
+	*CoreLocal::get().scheduler.borrow_mut() = Some(scheduler);
 }
 
 pub(crate) fn increment_irq_counter(irq_no: u8) {

src/arch/x86_64/kernel/apic.rs

@@ -229,7 +229,7 @@ extern "x86-interrupt" fn wakeup_handler(_stack_frame: interrupts::ExceptionStac
 
 	debug!("Received Wakeup Interrupt");
 	increment_irq_counter(WAKEUP_INTERRUPT_NUMBER);
-	let core_scheduler = core_scheduler();
+	let mut core_scheduler = core_scheduler();
 	core_scheduler.check_input();
 	eoi();
 	if core_scheduler.is_scheduling() {

src/arch/x86_64/kernel/core_local.rs

@@ -24,7 +24,7 @@ pub(crate) struct CoreLocal {
 	/// Sequential ID of this CPU Core.
 	core_id: CoreId,
 	/// Scheduler for this CPU Core.
-	scheduler: Cell<*mut PerCoreScheduler>,
+	scheduler: RefCell<Option<PerCoreScheduler>>,
 	/// Task State Segment (TSS) allocated for this CPU Core.
 	pub tss: Cell<*mut TaskStateSegment>,
 	/// start address of the kernel stack
@@ -54,7 +54,7 @@ impl CoreLocal {
 		let this = Self {
 			this: ptr::null_mut(),
 			core_id,
-			scheduler: Cell::new(ptr::null_mut()),
+			scheduler: RefCell::new(None),
 			tss: Cell::new(ptr::null_mut()),
 			kernel_stack: Cell::new(0),
 			irq_statistics,
@@ -101,16 +101,18 @@ pub(crate) fn core_id() -> CoreId {
 	}
 }
 
-pub(crate) fn core_scheduler() -> &'static mut PerCoreScheduler {
-	unsafe { &mut *CoreLocal::get().scheduler.get() }
+pub(crate) fn core_scheduler() -> RefMut<'static, PerCoreScheduler> {
+	RefMut::map(CoreLocal::get().scheduler.borrow_mut(), |scheduler| {
+		scheduler.as_mut().unwrap()
+	})
 }
 
 pub(crate) fn async_tasks() -> RefMut<'static, Vec<AsyncTask>> {
 	CoreLocal::get().async_tasks.borrow_mut()
 }
 
-pub(crate) fn set_core_scheduler(scheduler: *mut PerCoreScheduler) {
-	CoreLocal::get().scheduler.set(scheduler);
+pub(crate) fn set_core_scheduler(scheduler: PerCoreScheduler) {
+	*CoreLocal::get().scheduler.borrow_mut() = Some(scheduler);
 }
 
 pub(crate) fn increment_irq_counter(irq_no: u8) {

src/scheduler/mod.rs

@@ -1,9 +1,8 @@
-use alloc::boxed::Box;
 use alloc::collections::{BTreeMap, VecDeque};
 use alloc::rc::Rc;
 #[cfg(feature = "smp")]
 use alloc::vec::Vec;
-use core::cell::RefCell;
+use core::cell::{RefCell, RefMut};
 use core::sync::atomic::{AtomicU32, Ordering};
 
 use crossbeam_utils::Backoff;
@@ -88,9 +87,9 @@ pub trait PerCoreSchedulerExt {
 	fn exit(self, exit_code: i32) -> !;
 }
 
-impl PerCoreSchedulerExt for &mut PerCoreScheduler {
+impl PerCoreSchedulerExt for RefMut<'_, PerCoreScheduler> {
 	#[cfg(target_arch = "x86_64")]
-	fn reschedule(self) {
+	fn reschedule(mut self) {
 		without_interrupts(|| {
 			if let Some(last_stack_pointer) = self.scheduler() {
 				let (new_stack_pointer, is_idle) = {
@@ -102,10 +101,12 @@ impl PerCoreSchedulerExt for &mut PerCoreScheduler {
 				};
 
 				if is_idle || Rc::ptr_eq(&self.current_task, &self.fpu_owner) {
+					drop(self);
 					unsafe {
 						switch_to_fpu_owner(last_stack_pointer, new_stack_pointer.as_usize());
 					}
 				} else {
+					drop(self);
 					unsafe {
 						switch_to_task(last_stack_pointer, new_stack_pointer.as_usize());
 					}
@@ -123,6 +124,8 @@ impl PerCoreSchedulerExt for &mut PerCoreScheduler {
 
 		use crate::interrupts::SGI_RESCHED;
 
+		drop(self);
+
 		unsafe {
 			asm!("dsb nsh", "isb", options(nostack, nomem, preserves_flags));
 		}
@@ -142,13 +145,14 @@ impl PerCoreSchedulerExt for &mut PerCoreScheduler {
 	}
 
 	#[cfg(any(feature = "tcp", feature = "udp"))]
-	fn add_network_timer(self, wakeup_time: Option<u64>) {
+	fn add_network_timer(mut self, wakeup_time: Option<u64>) {
 		without_interrupts(|| {
 			self.blocked_tasks.add_network_timer(wakeup_time);
+			drop(self);
 		})
 	}
 
-	fn exit(self, exit_code: i32) -> ! {
+	fn exit(mut self, exit_code: i32) -> ! {
 		without_interrupts(|| {
 			// Get the current task.
 			let mut current_task_borrowed = self.current_task.borrow_mut();
@@ -538,7 +542,7 @@ impl PerCoreScheduler {
 		let backoff = Backoff::new();
 
 		loop {
-			let core_scheduler = core_scheduler();
+			let mut core_scheduler = core_scheduler();
 			interrupts::disable();
 
 			// run async tasks
@@ -548,6 +552,7 @@ impl PerCoreScheduler {
 			core_scheduler.cleanup_tasks();
 
 			if core_scheduler.ready_queue.is_empty() {
+				drop(core_scheduler);
 				if backoff.is_completed() {
 					interrupts::enable_and_wait();
 				} else {
@@ -698,7 +703,7 @@ pub fn add_current_core() {
 		"Initializing scheduler for core {} with idle task {}",
 		core_id, tid
 	);
-	let boxed_scheduler = Box::new(PerCoreScheduler {
+	set_core_scheduler(PerCoreScheduler {
 		#[cfg(feature = "smp")]
 		core_id,
 		current_task: idle_task.clone(),
@@ -709,9 +714,6 @@ pub fn add_current_core() {
 		finished_tasks: VecDeque::new(),
 		blocked_tasks: BlockedTaskQueue::new(),
 	});
-
-	let scheduler = Box::into_raw(boxed_scheduler);
-	set_core_scheduler(scheduler);
 	#[cfg(feature = "smp")]
 	{
 		SCHEDULER_INPUTS.lock().insert(
@@ -728,18 +730,17 @@ fn get_scheduler_input(core_id: CoreId) -> &'static InterruptTicketMutex<Schedul
 }
 
 pub fn join(id: TaskId) -> Result<(), ()> {
-	let core_scheduler = core_scheduler();
-
 	debug!(
 		"Task {} is waiting for task {}",
-		core_scheduler.get_current_task_id(),
+		core_scheduler().get_current_task_id(),
 		id
 	);
 
 	loop {
 		let mut waiting_tasks_guard = WAITING_TASKS.lock();
 
 		if let Some(queue) = waiting_tasks_guard.get_mut(&id) {
+			let mut core_scheduler = core_scheduler();
 			queue.push_back(core_scheduler.get_current_task_handle());
 			core_scheduler.block_current_task(None);
 

src/synch/futex.rs

@@ -52,14 +52,15 @@ pub(crate) fn futex_wait(
 		timeout
 	};
 
-	let scheduler = core_scheduler();
+	let mut scheduler = core_scheduler();
 	scheduler.block_current_task(wakeup_time);
 	let handle = scheduler.get_current_task_handle();
 	parking_lot.entry(addr(address)).or_default().push(handle);
 	drop(parking_lot);
+	drop(scheduler);
 
 	loop {
-		scheduler.reschedule();
+		core_scheduler().reschedule();
 
 		let mut parking_lot = PARKING_LOT.lock();
 		if matches!(wakeup_time, Some(t) if t <= get_timer_ticks()) {
@@ -88,7 +89,7 @@ pub(crate) fn futex_wait(
 			} else {
 				// A spurious wakeup occurred, sleep again.
 				// Tasks do not change core, so the handle in the parking lot is still current.
-				scheduler.block_current_task(wakeup_time);
+				core_scheduler().block_current_task(wakeup_time);
 			}
 		}
 		drop(parking_lot);
@@ -121,14 +122,15 @@ pub(crate) fn futex_wait_and_set(
 		timeout
 	};
 
-	let scheduler = core_scheduler();
+	let mut scheduler = core_scheduler();
 	scheduler.block_current_task(wakeup_time);
 	let handle = scheduler.get_current_task_handle();
 	parking_lot.entry(addr(address)).or_default().push(handle);
 	drop(parking_lot);
+	drop(scheduler);
 
 	loop {
-		scheduler.reschedule();
+		core_scheduler().reschedule();
 
 		let mut parking_lot = PARKING_LOT.lock();
 		if matches!(wakeup_time, Some(t) if t <= get_timer_ticks()) {
@@ -157,7 +159,7 @@ pub(crate) fn futex_wait_and_set(
 			} else {
 				// A spurious wakeup occurred, sleep again.
 				// Tasks do not change core, so the handle in the parking lot is still current.
-				scheduler.block_current_task(wakeup_time);
+				core_scheduler().block_current_task(wakeup_time);
 			}
 		}
 		drop(parking_lot);
@@ -178,7 +180,7 @@ pub(crate) fn futex_wake(address: &AtomicU32, count: i32) -> i32 {
 		Entry::Vacant(_) => return 0,
 	};
 
-	let scheduler = core_scheduler();
+	let mut scheduler = core_scheduler();
 	let mut woken = 0;
 	while woken != count || count == i32::MAX {
 		match queue.get_mut().pop() {
@@ -213,7 +215,7 @@ pub(crate) fn futex_wake_or_set(address: &AtomicU32, count: i32, new_value: u32)
 		}
 	};
 
-	let scheduler = core_scheduler();
+	let mut scheduler = core_scheduler();
 	let mut woken = 0;
 	while woken != count || count == i32::MAX {
 		match queue.get_mut().pop() {

src/synch/recmutex.rs

@@ -27,10 +27,10 @@ impl RecursiveMutex {
 
 	pub fn acquire(&self) {
 		// Get information about the current task.
-		let core_scheduler = core_scheduler();
-		let tid = core_scheduler.get_current_task_id();
+		let tid = core_scheduler().get_current_task_id();
 
 		loop {
+			let mut core_scheduler = core_scheduler();
 			{
 				let mut locked_state = self.state.lock();
 

src/synch/semaphore.rs

@@ -68,12 +68,12 @@ impl Semaphore {
 	pub fn acquire(&self, time: Option<u64>) -> bool {
 		#[cfg(feature = "smp")]
 		let backoff = Backoff::new();
-		let core_scheduler = core_scheduler();
 
 		let wakeup_time = time.map(|ms| crate::arch::processor::get_timer_ticks() + ms * 1000);
 
 		// Loop until we have acquired the semaphore.
 		loop {
+			let mut core_scheduler = core_scheduler();
 			let mut locked_state = self.state.lock();
 
 			if locked_state.count > 0 {

src/syscalls/tasks.rs

@@ -116,7 +116,7 @@ pub(crate) extern "C" fn __sys_usleep(usecs: u64) {
 		// Enough time to set a wakeup timer and block the current task.
 		debug!("sys_usleep blocking the task for {} microseconds", usecs);
 		let wakeup_time = arch::processor::get_timer_ticks() + usecs;
-		let core_scheduler = core_scheduler();
+		let mut core_scheduler = core_scheduler();
 		core_scheduler.block_current_task(Some(wakeup_time));
 
 		// Switch to the next task.
@@ -300,7 +300,7 @@ static BLOCKED_TASKS: InterruptTicketMutex<BTreeMap<TaskId, TaskHandle>> =
 
 extern "C" fn __sys_block_current_task(timeout: &Option<u64>) {
 	let wakeup_time = timeout.map(|t| arch::processor::get_timer_ticks() + t * 1000);
-	let core_scheduler = core_scheduler();
+	let mut core_scheduler = core_scheduler();
 	let handle = core_scheduler.get_current_task_handle();
 	let tid = core_scheduler.get_current_task_id();