Introduce the semantics of the atomic-mode guard

ostd/src/cpu/mod.rs

@@ -16,10 +16,7 @@ cfg_if::cfg_if! {
 pub use set::{AtomicCpuSet, CpuSet};
 use spin::Once;
 
-use crate::{
-    arch::boot::smp::get_num_processors, cpu_local_cell, task::DisabledPreemptGuard,
-    trap::DisabledLocalIrqGuard,
-};
+use crate::{arch::boot::smp::get_num_processors, cpu_local_cell, task::atomic_mode::InAtomicMode};
 
 /// The ID of a CPU in the system.
 ///
@@ -125,13 +122,10 @@ pub fn current_cpu_racy() -> CpuId {
     CpuId(id)
 }
 
-// SAFETY: When IRQs are disabled, the task cannot be passively preempted and
-// migrates to another CPU. If the task actively calls `yield`, it will not be
-// successful either.
-unsafe impl PinCurrentCpu for DisabledLocalIrqGuard {}
-// SAFETY: When preemption is disabled, the task cannot be preempted and migrates
-// to another CPU.
-unsafe impl PinCurrentCpu for DisabledPreemptGuard {}
+// SAFETY: A guard that enforces the atomic mode requires disabling any
+// context switching. So naturally, the current task is pinned on the CPU.
+unsafe impl<T: InAtomicMode> PinCurrentCpu for T {}
+unsafe impl PinCurrentCpu for dyn InAtomicMode + '_ {}
 
 cpu_local_cell! {
     /// The number of the current CPU.

ostd/src/task/atomic_mode.rs

@@ -45,3 +45,31 @@ pub fn might_sleep() {
         );
     }
 }
+
+/// A marker trait for guard types that enforce the atomic mode.
+///
+/// Key kernel primitives such as `SpinLock` and `Rcu` rely on
+/// [the atomic mode](crate::task::atomic_mode) for correctness or soundness.
+/// The existence of such a guard guarantees that
+/// the current task is executing in the atomic mode.
+///
+/// # Safety
+///
+/// The implementer must ensure that the atomic mode is maintained while
+/// the guard type is alive.
+///
+/// [`DisabledLocalIrqGuard`]: crate::task::DisabledPreemptGuard
+/// [`DisabledPreemptGuard`]: crate::trap::DisabledLocalIrqGuard
+pub unsafe trait InAtomicMode {}
+
+/// Abstracts any type from which one can obtain a reference to an atomic-mode guard.
+pub trait AsAtomicModeGuard {
+    /// Returns a guard for the atomic mode.
+    fn as_atomic_mode_guard(&self) -> &dyn InAtomicMode;
+}
+
+impl<G: InAtomicMode> AsAtomicModeGuard for G {
+    fn as_atomic_mode_guard(&self) -> &dyn InAtomicMode {
+        self
+    }
+}

ostd/src/task/preempt/guard.rs

@@ -1,6 +1,6 @@
 // SPDX-License-Identifier: MPL-2.0
 
-use crate::sync::GuardTransfer;
+use crate::{sync::GuardTransfer, task::atomic_mode::InAtomicMode};
 
 /// A guard for disable preempt.
 #[clippy::has_significant_drop]
@@ -13,6 +13,10 @@ pub struct DisabledPreemptGuard {
 
 impl !Send for DisabledPreemptGuard {}
 
+// SAFETY: The guard disables preemptions, which meets the second
+// sufficient condition for atomic mode.
+unsafe impl InAtomicMode for DisabledPreemptGuard {}
+
 impl DisabledPreemptGuard {
     fn new() -> Self {
         super::cpu_local::inc_guard_count();

ostd/src/trap/irq.rs

@@ -6,6 +6,7 @@ use crate::{
     arch::irq::{self, IrqCallbackHandle, IRQ_ALLOCATOR},
     prelude::*,
     sync::GuardTransfer,
+    task::atomic_mode::InAtomicMode,
     trap::TrapFrame,
     Error,
 };
@@ -140,6 +141,10 @@ pub struct DisabledLocalIrqGuard {
 
 impl !Send for DisabledLocalIrqGuard {}
 
+// SAFETY: The guard disables local IRQs, which meets the first
+// sufficient condition for atomic mode.
+unsafe impl InAtomicMode for DisabledLocalIrqGuard {}
+
 impl DisabledLocalIrqGuard {
     fn new() -> Self {
         let was_enabled = irq::is_local_enabled();