Enable check of atomic mode

kernel/src/lib.rs

@@ -30,6 +30,8 @@
 #![feature(linked_list_retain)]
 #![register_tool(component_access_control)]
 
+use core::sync::atomic::Ordering;
+
 use ostd::{
     arch::qemu::{exit_qemu, QemuExitCode},
     boot,
@@ -80,6 +82,7 @@ pub fn main() {
     ostd::early_println!("[kernel] OSTD initialized. Preparing components.");
     component::init_all(component::parse_metadata!()).unwrap();
     init();
+    ostd::IN_BOOTSTRAP_CONTEXT.store(false, Ordering::Relaxed);
     Thread::spawn_kernel_thread(ThreadOptions::new(init_thread));
     unreachable!()
 }

ostd/src/lib.rs

@@ -44,6 +44,8 @@ pub mod task;
 pub mod trap;
 pub mod user;
 
+use core::sync::atomic::AtomicBool;
+
 pub use ostd_macros::main;
 pub use ostd_pod::Pod;
 
@@ -96,6 +98,9 @@ pub unsafe fn init() {
     invoke_ffi_init_funcs();
 }
 
+/// Indicates whether the kernel is in bootstrap context.
+pub static IN_BOOTSTRAP_CONTEXT: AtomicBool = AtomicBool::new(true);
+
 /// Invoke the initialization functions defined in the FFI.
 /// The component system uses this function to call the initialization functions of
 /// the components.

ostd/src/task/atomic_mode.rs

@@ -0,0 +1,46 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! Atomic Mode
+//!
+//! Multitasking, while powerful, can sometimes lead to undesirable
+//! or catastrophic consequences if being misused.
+//! For instance, a user of OSTD might accidentally write an IRQ handler
+//! that relies on mutexes,
+//! which could attempt to sleep within an interrupt context---something that must be avoided.
+//! Another common mistake is
+//! acquiring a spinlock in a task context and then attempting to yield or sleep,
+//! which can easily lead to deadlocks.
+//!
+//! To mitigate the risks associated with improper multitasking,
+//! we introduce the concept of atomic mode.
+//! Kernel code is considered to be running in atomic mode
+//! if one of the following conditions is met:
+//!
+//! 1. Task preemption is disabled, such as when a spinlock is held.
+//! 2. Local IRQs are disabled, such as during interrupt context.
+//!
+//! While in atomic mode,
+//! any attempt to perform "sleep-like" actions will trigger a panic:
+//!
+//! 1. Switching to another task.
+//! 2. Switching to user space.
+//!
+//! This module provides API to detect such "sleep-like" actions.
+
+use core::sync::atomic::Ordering;
+
+/// Marks a function as one that might sleep.
+///
+/// This function will panic if it is executed in atomic mode.
+pub fn might_sleep() {
+    let preempt_count = super::preempt::cpu_local::get_guard_count();
+    let is_local_irq_enabled = crate::arch::irq::is_local_enabled();
+    if (preempt_count != 0 || !is_local_irq_enabled)
+        && !crate::IN_BOOTSTRAP_CONTEXT.load(Ordering::Relaxed)
+    {
+        panic!(
+            "This function might break atomic mode (preempt_count = {}, is_local_irq_enabled = {})",
+            preempt_count, is_local_irq_enabled
+        );
+    }
+}

ostd/src/task/mod.rs

@@ -2,6 +2,7 @@
 
 //! Tasks are the unit of code execution.
 
+pub(crate) mod atomic_mode;
 mod preempt;
 mod processor;
 pub mod scheduler;

ostd/src/task/processor.rs

@@ -2,10 +2,7 @@
 
 use alloc::sync::Arc;
 
-use super::{
-    preempt::cpu_local,
-    task::{context_switch, Task, TaskContext},
-};
+use super::task::{context_switch, Task, TaskContext};
 use crate::cpu_local_cell;
 
 cpu_local_cell! {
@@ -46,18 +43,7 @@ pub(super) fn current_task() -> Option<Arc<Task>> {
 /// This function will panic if called while holding preemption locks or with
 /// local IRQ disabled.
 pub(super) fn switch_to_task(next_task: Arc<Task>) {
-    let guard_count = cpu_local::get_guard_count();
-    if guard_count != 0 {
-        panic!(
-            "Switching task with preemption disabled (nesting depth: {})",
-            guard_count
-        );
-    }
-
-    assert!(
-        crate::arch::irq::is_local_enabled(),
-        "Switching task with local IRQ disabled"
-    );
+    super::atomic_mode::might_sleep();
 
     let irq_guard = crate::trap::disable_local();
 

ostd/src/user.rs

@@ -140,6 +140,7 @@ impl<'a> UserMode<'a> {
         F: FnMut() -> bool,
     {
         debug_assert!(Arc::ptr_eq(&self.current, &Task::current().unwrap()));
+        crate::task::atomic_mode::might_sleep();
         self.context.execute(has_kernel_event)
     }