fix(core_local): put scheduler into InterruptRefCell

Cargo.toml

@@ -73,6 +73,7 @@ dyn-clone = "1.0"
 hashbrown = { version = "0.14", default-features = false }
 hermit-entry = { version = "0.9", features = ["kernel"] }
 hermit-sync = "0.1"
+interrupt-ref-cell = { path = "../../interrupt-ref-cell" }
 lock_api = "0.4"
 log = { version = "0.4", default-features = false }
 num = { version = "0.4", default-features = false }

src/arch/aarch64/kernel/core_local.rs

@@ -1,7 +1,8 @@
 use alloc::boxed::Box;
 use alloc::vec::Vec;
+use interrupt_ref_cell::{InterruptRefCell, InterruptRefMut};
 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 +21,7 @@ pub(crate) struct CoreLocal {
 	/// ID of the current Core.
 	core_id: CoreId,
 	/// Scheduler of the current Core.
-	scheduler: Cell<*mut PerCoreScheduler>,
+	scheduler: InterruptRefCell<Option<PerCoreScheduler>>,
 	/// Interface to the interrupt counters
 	irq_statistics: &'static IrqStatistics,
 	/// Queue of async tasks
@@ -44,7 +45,7 @@ impl CoreLocal {
 		let this = Self {
 			this: ptr::null_mut(),
 			core_id,
-			scheduler: Cell::new(ptr::null_mut()),
+			scheduler: InterruptRefCell::new(None),
 			irq_statistics,
 			async_tasks: RefCell::new(Vec::new()),
 			#[cfg(feature = "smp")]
@@ -91,17 +92,23 @@ pub(crate) fn core_id() -> CoreId {
 	}
 }
 
-#[inline]
-pub(crate) fn core_scheduler() -> &'static mut PerCoreScheduler {
-	unsafe { &mut *CoreLocal::get().scheduler.get() }
+#[track_caller]
+pub(crate) fn core_scheduler() -> InterruptRefMut<'static, PerCoreScheduler> {
+	println!("{}", core::panic::Location::caller());
+	if CoreLocal::get().scheduler.try_borrow().is_err() {
+		println!("Oh no");
+	}
+	InterruptRefMut::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/aarch64/kernel/interrupts.rs

@@ -1,7 +1,8 @@
 use alloc::collections::BTreeMap;
 use alloc::vec::Vec;
 use core::arch::asm;
-use core::sync::atomic::{AtomicU64, Ordering};
+use core::ops::DerefMut;
+use core::sync::atomic::{AtomicU64, Ordering, compiler_fence};
 
 use aarch64::regs::*;
 use ahash::RandomState;
@@ -118,7 +119,20 @@ pub(crate) extern "C" fn do_fiq(state: &State) -> *mut usize {
 			}
 		}
 
-		core_scheduler().handle_waiting_tasks();
+		let mut network_wakeup_time = None;
+
+		#[cfg(any(feature = "tcp", feature = "udp"))]
+		if let Some(mut guard) = crate::executor::network::NIC.try_lock() {
+			if let crate::executor::network::NetworkState::Initialized(nic) = guard.deref_mut() {
+				let now = crate::executor::network::now();
+				nic.poll_common(now);
+				let time = crate::arch::processor::get_timer_ticks();
+				network_wakeup_time = nic.poll_delay(now).map(|d| d.total_micros() + time);
+			}
+		}
+
+		crate::executor::run();
+		core_scheduler().handle_waiting_tasks(network_wakeup_time);
 
 		GicV3::end_interrupt(irqid);
 
@@ -154,20 +168,37 @@ pub(crate) extern "C" fn do_irq(state: &State) -> *mut usize {
 			}
 		}
 
-		core_scheduler().handle_waiting_tasks();
+		let mut network_wakeup_time = None;
+
+		#[cfg(any(feature = "tcp", feature = "udp"))]
+		if let Some(mut guard) = crate::executor::network::NIC.try_lock() {
+			if let crate::executor::network::NetworkState::Initialized(nic) = guard.deref_mut() {
+				let now = crate::executor::network::now();
+				nic.poll_common(now);
+				let time = crate::arch::processor::get_timer_ticks();
+				network_wakeup_time = nic.poll_delay(now).map(|d| d.total_micros() + time);
+			}
+		}
+
+		crate::executor::run();
+		core_scheduler().handle_waiting_tasks(network_wakeup_time);
 
-		GicV3::end_interrupt(irqid);
 
 		if unsafe {
 			reschedule
 				|| vector == TIMER_INTERRUPT.try_into().unwrap()
 				|| vector == SGI_RESCHED.try_into().unwrap()
 		} {
 			// a timer interrupt may have caused unblocking of tasks
-			return core_scheduler()
-				.scheduler()
-				.unwrap_or(core::ptr::null_mut());
+			// run background tasks
+			crate::executor::run();
+			let ret = core_scheduler()
+				.scheduler();
+			compiler_fence(Ordering::SeqCst);
+				GicV3::end_interrupt(irqid);
+				return ret.unwrap_or(core::ptr::null_mut());
 		}
+		GicV3::end_interrupt(irqid);
 	}
 
 	core::ptr::null_mut()

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

@@ -1,5 +1,6 @@
 use alloc::boxed::Box;
 use alloc::vec::Vec;
+use interrupt_ref_cell::{InterruptRefCell, InterruptRefMut};
 use core::arch::asm;
 use core::cell::{Cell, RefCell, RefMut};
 use core::ptr;
@@ -24,7 +25,7 @@ pub(crate) struct CoreLocal {
 	/// Sequential ID of this CPU Core.
 	core_id: CoreId,
 	/// Scheduler for this CPU Core.
-	scheduler: Cell<*mut PerCoreScheduler>,
+	scheduler: InterruptRefCell<Option<PerCoreScheduler>>,
 	/// Task State Segment (TSS) allocated for this CPU Core.
 	pub tss: Cell<*mut TaskStateSegment>,
 	/// start address of the kernel stack
@@ -54,7 +55,7 @@ impl CoreLocal {
 		let this = Self {
 			this: ptr::null_mut(),
 			core_id,
-			scheduler: Cell::new(ptr::null_mut()),
+			scheduler: InterruptRefCell::new(None),
 			tss: Cell::new(ptr::null_mut()),
 			kernel_stack: Cell::new(0),
 			irq_statistics,
@@ -101,16 +102,23 @@ pub(crate) fn core_id() -> CoreId {
 	}
 }
 
-pub(crate) fn core_scheduler() -> &'static mut PerCoreScheduler {
-	unsafe { &mut *CoreLocal::get().scheduler.get() }
+#[track_caller]
+pub(crate) fn core_scheduler() -> InterruptRefMut<'static, PerCoreScheduler> {
+	println!("{}", core::panic::Location::caller());
+	if CoreLocal::get().scheduler.try_borrow().is_err() {
+		println!("Oh no");
+	}
+	InterruptRefMut::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/scheduler.rs

@@ -2,6 +2,7 @@
 
 use alloc::boxed::Box;
 use core::arch::asm;
+use core::ops::DerefMut;
 use core::{mem, ptr, slice};
 
 use align_address::Align;
@@ -365,7 +366,23 @@ impl TaskFrame for Task {
 
 extern "x86-interrupt" fn timer_handler(_stack_frame: interrupts::ExceptionStackFrame) {
 	increment_irq_counter(apic::TIMER_INTERRUPT_NUMBER);
-	core_scheduler().handle_waiting_tasks();
+
+
+	let mut network_wakeup_time = None;
+
+	#[cfg(any(feature = "tcp", feature = "udp"))]
+	if let Some(mut guard) = crate::executor::network::NIC.try_lock() {
+		if let crate::executor::network::NetworkState::Initialized(nic) = guard.deref_mut() {
+			let now = crate::executor::network::now();
+			nic.poll_common(now);
+			let time = crate::arch::processor::get_timer_ticks();
+			network_wakeup_time = nic.poll_delay(now).map(|d| d.total_micros() + time);
+		}
+	}
+
+	crate::executor::run();
+
+	core_scheduler().handle_waiting_tasks(network_wakeup_time);
 	apic::eoi();
 	core_scheduler().reschedule();
 }