Don't race between enabling IRQs and halting CPU

kernel/src/lib.rs

@@ -111,8 +111,7 @@ fn ap_init() {
         );
 
         loop {
-            crate::thread::Thread::yield_now();
-            ostd::cpu::sleep_for_interrupt();
+            ostd::task::halt_cpu();
         }
     }
 
@@ -156,10 +155,10 @@ fn init_thread() {
         karg.get_initproc_envp().to_vec(),
     )
     .expect("Run init process failed.");
+
     // Wait till initproc become zombie.
     while !initproc.status().is_zombie() {
-        crate::thread::Thread::yield_now();
-        ostd::cpu::sleep_for_interrupt();
+        ostd::task::halt_cpu();
     }
 
     // TODO: exit via qemu isa debug device should not be the only way.

ostd/src/arch/riscv/cpu/mod.rs

@@ -7,17 +7,3 @@ pub mod extension;
 pub mod local;
 
 pub use extension::{has_extensions, IsaExtensions};
-
-/// Halts the CPU.
-///
-/// This function halts the CPU until the next interrupt is received. By
-/// halting, the CPU might consume less power. Internally it is implemented
-/// using the `wfi` instruction.
-///
-/// Since the function sleeps the CPU, it should not be used within an atomic
-/// mode ([`crate::task::atomic_mode`]).
-#[track_caller]
-pub fn sleep_for_interrupt() {
-    crate::task::atomic_mode::might_sleep();
-    riscv::asm::wfi();
-}

ostd/src/arch/riscv/irq.rs

@@ -31,7 +31,31 @@ impl IrqRemapping {
     }
 }
 
+// FIXME: Mark this as unsafe. See
+// <https://github.com/asterinas/asterinas/issues/1120#issuecomment-2748696592>.
 pub(crate) fn enable_local() {
+    // SAFETY: The safety is upheld by the caller.
+    unsafe { riscv::interrupt::enable() }
+}
+
+/// Enables local IRQs and halts the CPU to wait for interrupts.
+///
+/// This method guarantees that no interrupts can occur in the middle. In other words, IRQs must
+/// either have been processed before this method is called, or they must wake the CPU up from the
+/// halting state.
+//
+// FIXME: Mark this as unsafe. See
+// <https://github.com/asterinas/asterinas/issues/1120#issuecomment-2748696592>.
+pub(crate) fn enable_local_and_halt() {
+    // RISC-V Instruction Set Manual, Machine-Level ISA, Version 1.13 says:
+    // "The WFI instruction can also be executed when interrupts are disabled. The operation of WFI
+    // must be unaffected by the global interrupt bits in `mstatus` (MIE and SIE) [..]"
+    //
+    // So we can use `wfi` even if IRQs are disabled. Pending IRQs can still wake up the CPU, but
+    // they will only occur later when we enable local IRQs.
+    riscv::asm::wfi();
+
+    // SAFETY: The safety is upheld by the caller.
     unsafe { riscv::interrupt::enable() }
 }
 

ostd/src/arch/x86/cpu/mod.rs

@@ -4,16 +4,3 @@
 
 pub mod context;
 pub mod local;
-
-/// Halts the CPU.
-///
-/// This function halts the CPU until the next interrupt is received. By
-/// halting, the CPU will enter the C-0 state and consume less power.
-///
-/// Since the function sleeps the CPU, it should not be used within an atomic
-/// mode ([`crate::task::atomic_mode`]).
-#[track_caller]
-pub fn sleep_for_interrupt() {
-    crate::task::atomic_mode::might_sleep();
-    x86_64::instructions::hlt();
-}

ostd/src/arch/x86/irq.rs

@@ -50,6 +50,8 @@ impl IrqRemapping {
     }
 }
 
+// FIXME: Mark this as unsafe. See
+// <https://github.com/asterinas/asterinas/issues/1120#issuecomment-2748696592>.
 pub(crate) fn enable_local() {
     x86_64::instructions::interrupts::enable();
     // When emulated with QEMU, interrupts may not be delivered if a STI instruction is immediately
@@ -58,6 +60,29 @@ pub(crate) fn enable_local() {
     x86_64::instructions::nop();
 }
 
+/// Enables local IRQs and halts the CPU to wait for interrupts.
+///
+/// This method guarantees that no interrupts can occur in the middle. In other words, IRQs must
+/// either have been processed before this method is called, or they must wake the CPU up from the
+/// halting state.
+//
+// FIXME: Mark this as unsafe. See
+// <https://github.com/asterinas/asterinas/issues/1120#issuecomment-2748696592>.
+pub(crate) fn enable_local_and_halt() {
+    // SAFETY:
+    // 1. `sti` is safe to use because its safety requirement is upheld by the caller.
+    // 2. `hlt` is safe to use because it halts the CPU for interrupts.
+    unsafe {
+        // Intel(R) 64 and IA-32 Architectures Software Developer's Manual says:
+        // "If IF = 0, maskable hardware interrupts remain inhibited on the instruction boundary
+        // following an execution of STI."
+        //
+        // So interrupts will only occur at or after the HLT instruction, which guarantee that
+        // interrupts won't occur between enabling the local IRQs and halting the CPU.
+        core::arch::asm!("sti", "hlt", options(nomem, nostack, preserves_flags),)
+    };
+}
+
 pub(crate) fn disable_local() {
     x86_64::instructions::interrupts::disable();
 }

ostd/src/arch/x86/kernel/tsc.rs

@@ -80,11 +80,17 @@ pub fn determine_tsc_freq_via_pit() -> u64 {
     static FREQUENCY: AtomicU64 = AtomicU64::new(0);
 
     // Wait until `FREQUENCY` is ready
-    x86_64::instructions::interrupts::enable();
-    while !IS_FINISH.load(Ordering::Acquire) {
-        x86_64::instructions::hlt();
+    loop {
+        crate::arch::irq::enable_local_and_halt();
+
+        // Disable local IRQs so they won't come after checking `IS_FINISH`
+        // but before halting the CPU.
+        crate::arch::irq::disable_local();
+
+        if IS_FINISH.load(Ordering::Acquire) {
+            break;
+        }
     }
-    x86_64::instructions::interrupts::disable();
 
     // Disable PIT
     drop(irq);

ostd/src/arch/x86/timer/apic.rs

@@ -120,11 +120,17 @@ fn init_periodic_mode_config() {
     apic.set_timer_init_count(0xFFFF_FFFF);
 
     // Wait until `CONFIG` is ready
-    x86_64::instructions::interrupts::enable();
-    while !CONFIG.is_completed() {
-        x86_64::instructions::hlt();
+    loop {
+        crate::arch::irq::enable_local_and_halt();
+
+        // Disable local IRQs so they won't come after checking `CONFIG`
+        // but before halting the CPU.
+        crate::arch::irq::disable_local();
+
+        if CONFIG.is_completed() {
+            break;
+        }
     }
-    x86_64::instructions::interrupts::disable();
 
     // Disable PIT
     drop(irq);

ostd/src/task/mod.rs

@@ -24,7 +24,7 @@ use spin::Once;
 use utils::ForceSync;
 
 pub use self::{
-    preempt::{disable_preempt, DisabledPreemptGuard},
+    preempt::{disable_preempt, halt_cpu, DisabledPreemptGuard},
     scheduler::info::{AtomicCpuId, TaskScheduleInfo},
 };
 pub(crate) use crate::arch::task::{context_switch, TaskContext};

ostd/src/task/preempt/cpu_local.rs

@@ -27,7 +27,6 @@ pub(in crate::task) fn should_preempt() -> bool {
     PREEMPT_INFO.load() == 0
 }
 
-#[expect(dead_code)]
 pub(in crate::task) fn need_preempt() -> bool {
     PREEMPT_INFO.load() & NEED_PREEMPT_MASK == 0
 }

ostd/src/task/preempt/mod.rs

@@ -4,3 +4,34 @@ pub(super) mod cpu_local;
 mod guard;
 
 pub use self::guard::{disable_preempt, DisabledPreemptGuard};
+
+/// Halts the CPU until interrupts if no preemption is required.
+///
+/// This function will return if:
+///  - preemption is required when calling this function,
+///  - preemption is required during halting the CPU, or
+///  - interrupts occur during halting the CPU.
+///
+/// This function will perform preemption before returning if
+/// preemption is required.
+///
+/// # Panics
+///
+/// This function will panic if it is called in the atomic mode
+/// ([`crate::task::atomic_mode`]).
+#[track_caller]
+pub fn halt_cpu() {
+    crate::task::atomic_mode::might_sleep();
+
+    let irq_guard = crate::trap::irq::disable_local();
+
+    if cpu_local::need_preempt() {
+        drop(irq_guard);
+    } else {
+        core::mem::forget(irq_guard);
+        // IRQs were previously enabled (checked by `might_sleep`). So we can re-enable them now.
+        crate::arch::irq::enable_local_and_halt();
+    }
+
+    super::scheduler::might_preempt();
+}

ostd/src/task/processor.rs

@@ -72,7 +72,7 @@ pub(super) fn switch_to_task(next_task: Arc<Task>) {
     PREVIOUS_TASK_PTR.store(current_task_ptr);
 
     // We must disable IRQs when switching, see `after_switching_to`.
-    let _ = core::mem::ManuallyDrop::new(irq_guard);
+    core::mem::forget(irq_guard);
 
     // SAFETY:
     // 1. We have exclusive access to both the current context and the next context (see above).