Add guard which disables bottom half

kernel/comps/softirq/src/lib.rs

@@ -10,12 +10,18 @@ use alloc::boxed::Box;
 use core::sync::atomic::{AtomicU8, Ordering};
 
 use component::{init_component, ComponentInitError};
-use ostd::{cpu_local_cell, trap::register_bottom_half_handler};
+use lock::is_softirq_enabled;
+use ostd::{
+    cpu_local_cell,
+    trap::{disable_local, register_bottom_half_handler, DisabledLocalIrqGuard},
+};
 use spin::Once;
 
+mod lock;
 pub mod softirq_id;
 mod taskless;
 
+pub use lock::{BottomHalfDisabled, DisableLocalBottomHalfGuard};
 pub use taskless::Taskless;
 
 /// A representation of a software interrupt (softirq) line.
@@ -122,10 +128,19 @@ cpu_local_cell! {
 }
 
 /// Processes pending softirqs.
+fn process_pending(irq_guard: DisabledLocalIrqGuard) -> DisabledLocalIrqGuard {
+    if !is_softirq_enabled() {
+        return irq_guard;
+    }
+
+    process_all_pending(irq_guard)
+}
+
+/// Processes all pending softirqs regardless of whether softirqs are disabled.
 ///
 /// The processing instructions will iterate for `SOFTIRQ_RUN_TIMES` times. If any softirq
 /// is raised during the iteration, it will be processed.
-fn process_pending() {
+fn process_all_pending(mut irq_guard: DisabledLocalIrqGuard) -> DisabledLocalIrqGuard {
     const SOFTIRQ_RUN_TIMES: u8 = 5;
 
     for _i in 0..SOFTIRQ_RUN_TIMES {
@@ -138,10 +153,19 @@ fn process_pending() {
         if action_mask == 0 {
             break;
         }
+
+        drop(irq_guard);
+
         while action_mask > 0 {
             let action_id = u8::trailing_zeros(action_mask) as u8;
             SoftIrqLine::get(action_id).callback.get().unwrap()();
             action_mask &= action_mask - 1;
         }
+
+        irq_guard = disable_local();
     }
+
+    // TODO: Wake up ksoftirqd if some softirqs are still pending.
+
+    irq_guard
 }

kernel/comps/softirq/src/lock.rs

@@ -0,0 +1,176 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use ostd::{
+    cpu_local_cell,
+    sync::{GuardTransfer, SpinGuardian},
+    task::{
+        atomic_mode::{AsAtomicModeGuard, InAtomicMode},
+        disable_preempt, DisabledPreemptGuard,
+    },
+    trap::{disable_local, in_interrupt_context},
+};
+
+use crate::process_all_pending;
+
+cpu_local_cell! {
+    static DISABLE_SOFTIRQ_COUNT: u8 = 0;
+}
+
+/// Returns whether softirq is enabled on local CPU.
+pub(super) fn is_softirq_enabled() -> bool {
+    DISABLE_SOFTIRQ_COUNT.load() == 0
+}
+
+/// A guardian that disables bottom half while holding a lock.
+pub enum BottomHalfDisabled {}
+
+/// A guard for disabled local softirqs.
+pub struct DisableLocalBottomHalfGuard {
+    preempt: DisabledPreemptGuard,
+}
+
+impl AsAtomicModeGuard for DisableLocalBottomHalfGuard {
+    fn as_atomic_mode_guard(&self) -> &dyn InAtomicMode {
+        &self.preempt
+    }
+}
+
+impl Drop for DisableLocalBottomHalfGuard {
+    fn drop(&mut self) {
+        // Once the guard is dropped, we will process pending items within
+        // the current thread's context if softirq is going to be enabled.
+        // This behavior is similar to how Linux handles pending softirqs.
+        if DISABLE_SOFTIRQ_COUNT.load() == 1 && !in_interrupt_context() {
+            // Preemption and softirq are not really enabled at the moment,
+            // so we can guarantee that we'll process any pending softirqs for the current CPU.
+            let irq_guard = disable_local();
+            let irq_guard = process_all_pending(irq_guard);
+
+            // To avoid race conditions, we should decrease the softirq count first,
+            // then drop the IRQ guard.
+            DISABLE_SOFTIRQ_COUNT.sub_assign(1);
+            drop(irq_guard);
+            return;
+        }
+
+        DISABLE_SOFTIRQ_COUNT.sub_assign(1);
+    }
+}
+
+#[must_use]
+fn disable_local_bottom_half() -> DisableLocalBottomHalfGuard {
+    // When disabling softirq, we must also disable preemption
+    // to avoid the task to be scheduled to other CPUs.
+    let preempt = disable_preempt();
+    DISABLE_SOFTIRQ_COUNT.add_assign(1);
+    DisableLocalBottomHalfGuard { preempt }
+}
+
+impl GuardTransfer for DisableLocalBottomHalfGuard {
+    fn transfer_to(&mut self) -> Self {
+        disable_local_bottom_half()
+    }
+}
+
+impl SpinGuardian for BottomHalfDisabled {
+    type Guard = DisableLocalBottomHalfGuard;
+    type ReadGuard = DisableLocalBottomHalfGuard;
+
+    fn read_guard() -> Self::ReadGuard {
+        disable_local_bottom_half()
+    }
+
+    fn guard() -> Self::Guard {
+        disable_local_bottom_half()
+    }
+}
+
+#[cfg(ktest)]
+mod test {
+    use ostd::{
+        prelude::*,
+        sync::{RwLock, SpinLock},
+    };
+
+    use super::*;
+
+    #[ktest]
+    fn spinlock_disable_bh() {
+        let lock = SpinLock::<(), BottomHalfDisabled>::new(());
+
+        assert!(is_softirq_enabled());
+
+        let guard = lock.lock();
+        assert!(!is_softirq_enabled());
+
+        drop(guard);
+        assert!(is_softirq_enabled());
+    }
+
+    #[ktest]
+    fn nested_spin_lock_disable_bh() {
+        let lock1 = SpinLock::<(), BottomHalfDisabled>::new(());
+        let lock2 = SpinLock::<(), BottomHalfDisabled>::new(());
+
+        assert!(is_softirq_enabled());
+
+        let guard1 = lock1.lock();
+        let guard2 = lock2.lock();
+        assert!(!is_softirq_enabled());
+
+        drop(guard1);
+        assert!(!is_softirq_enabled());
+
+        drop(guard2);
+        assert!(is_softirq_enabled());
+    }
+
+    #[ktest]
+    fn rwlock_disable_bh() {
+        let rwlock: RwLock<(), BottomHalfDisabled> = RwLock::new(());
+
+        assert!(is_softirq_enabled());
+
+        let write_guard = rwlock.write();
+        assert!(!is_softirq_enabled());
+
+        drop(write_guard);
+        assert!(is_softirq_enabled());
+    }
+
+    #[ktest]
+    fn nested_rwlock_disable_bh() {
+        let rwlock: RwLock<(), BottomHalfDisabled> = RwLock::new(());
+
+        assert!(is_softirq_enabled());
+
+        let read_guard1 = rwlock.read();
+        let read_guard2 = rwlock.read();
+        assert!(!is_softirq_enabled());
+
+        drop(read_guard1);
+        assert!(!is_softirq_enabled());
+
+        drop(read_guard2);
+        assert!(is_softirq_enabled());
+    }
+
+    #[test]
+    fn upgradable_rwlock_disable_bh() {
+        let rwlock: RwLock<(), BottomHalfDisabled> = RwLock::new(());
+
+        assert!(is_softirq_enabled());
+
+        let upgrade_guard = rwlock.upread();
+        assert!(!is_softirq_enabled());
+
+        let write_guard = upgrade_guard.upgrade();
+        assert!(!is_softirq_enabled());
+
+        let read_guard = write_guard.downgrade();
+        assert!(!is_softirq_enabled());
+
+        drop(read_guard);
+        assert!(is_softirq_enabled());
+    }
+}