Refactor the API of spinlocks

kernel/comps/network/src/buffer.rs

@@ -74,7 +74,8 @@ impl HasDaddr for TxBuffer {
 impl Drop for TxBuffer {
     fn drop(&mut self) {
         self.pool
-            .lock_irq_disabled()
+            .disable_irq()
+            .lock()
             .push_back(self.dma_stream.clone());
     }
 }
@@ -145,7 +146,7 @@ fn get_tx_stream_from_pool(
     nbytes: usize,
     tx_buffer_pool: &'static SpinLock<LinkedList<DmaStream>>,
 ) -> Option<DmaStream> {
-    let mut pool = tx_buffer_pool.lock_irq_disabled();
+    let mut pool = tx_buffer_pool.disable_irq().lock();
     let mut cursor = pool.cursor_front_mut();
     while let Some(current) = cursor.current() {
         if current.nbytes() >= nbytes {

kernel/comps/network/src/dma_pool.rs

@@ -97,7 +97,7 @@ impl DmaPool {
     pub fn alloc_segment(self: &Arc<Self>) -> Result<DmaSegment, ostd::Error> {
         // Lock order: pool.avail_pages -> pool.all_pages
         //             pool.avail_pages -> page.allocated_segments
-        let mut avail_pages = self.avail_pages.lock_irq_disabled();
+        let mut avail_pages = self.avail_pages.disable_irq().lock();
         if avail_pages.is_empty() {
             /// Allocate a new page
             let new_page = {
@@ -109,7 +109,7 @@ impl DmaPool {
                     pool,
                 )?)
             };
-            let mut all_pages = self.all_pages.lock_irq_disabled();
+            let mut all_pages = self.all_pages.disable_irq().lock();
             avail_pages.push_back(new_page.clone());
             all_pages.push_back(new_page);
         }
@@ -124,7 +124,7 @@ impl DmaPool {
 
     /// Returns the number of pages in pool
     fn num_pages(&self) -> usize {
-        self.all_pages.lock_irq_disabled().len()
+        self.all_pages.disable_irq().lock().len()
     }
 
     /// Return segment size in pool
@@ -166,7 +166,7 @@ impl DmaPage {
     }
 
     fn alloc_segment(self: &Arc<Self>) -> Option<DmaSegment> {
-        let mut segments = self.allocated_segments.lock_irq_disabled();
+        let mut segments = self.allocated_segments.disable_irq().lock();
         let free_segment_index = get_next_free_index(&segments, self.nr_blocks_per_page())?;
         segments.set(free_segment_index, true);
 
@@ -189,7 +189,7 @@ impl DmaPage {
     }
 
     fn is_full(&self) -> bool {
-        let segments = self.allocated_segments.lock_irq_disabled();
+        let segments = self.allocated_segments.disable_irq().lock();
         get_next_free_index(&segments, self.nr_blocks_per_page()).is_none()
     }
 }
@@ -257,10 +257,10 @@ impl Drop for DmaSegment {
 
         // Keep the same lock order as `pool.alloc_segment`
         // Lock order: pool.avail_pages -> pool.all_pages -> page.allocated_segments
-        let mut avail_pages = pool.avail_pages.lock_irq_disabled();
-        let mut all_pages = pool.all_pages.lock_irq_disabled();
+        let mut avail_pages = pool.avail_pages.disable_irq().lock();
+        let mut all_pages = pool.all_pages.disable_irq().lock();
 
-        let mut allocated_segments = page.allocated_segments.lock_irq_disabled();
+        let mut allocated_segments = page.allocated_segments.disable_irq().lock();
 
         let nr_blocks_per_page = PAGE_SIZE / self.size;
         let became_avail = get_next_free_index(&allocated_segments, nr_blocks_per_page).is_none();

kernel/comps/network/src/lib.rs

@@ -18,7 +18,10 @@ use core::{any::Any, fmt::Debug};
 pub use buffer::{RxBuffer, TxBuffer, RX_BUFFER_POOL, TX_BUFFER_POOL};
 use component::{init_component, ComponentInitError};
 pub use dma_pool::DmaSegment;
-use ostd::{sync::SpinLock, Pod};
+use ostd::{
+    sync::{PreemptDisabled, SpinLock},
+    Pod,
+};
 use smoltcp::phy;
 use spin::Once;
 
@@ -52,21 +55,23 @@ pub trait AnyNetworkDevice: Send + Sync + Any + Debug {
 
 pub trait NetDeviceIrqHandler = Fn() + Send + Sync + 'static;
 
-pub fn register_device(name: String, device: Arc<SpinLock<dyn AnyNetworkDevice>>) {
+pub fn register_device(name: String, device: Arc<SpinLock<dyn AnyNetworkDevice, PreemptDisabled>>) {
     COMPONENT
         .get()
         .unwrap()
         .network_device_table
-        .lock_irq_disabled()
+        .disable_irq()
+        .lock()
         .insert(name, (Arc::new(SpinLock::new(Vec::new())), device));
 }
 
-pub fn get_device(str: &str) -> Option<Arc<SpinLock<dyn AnyNetworkDevice>>> {
+pub fn get_device(str: &str) -> Option<Arc<SpinLock<dyn AnyNetworkDevice, PreemptDisabled>>> {
     let table = COMPONENT
         .get()
         .unwrap()
         .network_device_table
-        .lock_irq_disabled();
+        .disable_irq()
+        .lock();
     let (_, device) = table.get(str)?;
     Some(device.clone())
 }
@@ -80,11 +85,12 @@ pub fn register_recv_callback(name: &str, callback: impl NetDeviceIrqHandler) {
         .get()
         .unwrap()
         .network_device_table
-        .lock_irq_disabled();
+        .disable_irq()
+        .lock();
     let Some((callbacks, _)) = device_table.get(name) else {
         return;
     };
-    callbacks.lock_irq_disabled().push(Arc::new(callback));
+    callbacks.disable_irq().lock().push(Arc::new(callback));
 }
 
 pub fn handle_recv_irq(name: &str) {
@@ -92,11 +98,12 @@ pub fn handle_recv_irq(name: &str) {
         .get()
         .unwrap()
         .network_device_table
-        .lock_irq_disabled();
+        .disable_irq()
+        .lock();
     let Some((callbacks, _)) = device_table.get(name) else {
         return;
     };
-    let callbacks = callbacks.lock_irq_disabled();
+    let callbacks = callbacks.disable_irq().lock();
     for callback in callbacks.iter() {
         callback();
     }
@@ -107,7 +114,8 @@ pub fn all_devices() -> Vec<(String, NetworkDeviceRef)> {
         .get()
         .unwrap()
         .network_device_table
-        .lock_irq_disabled();
+        .disable_irq()
+        .lock();
     network_devs
         .iter()
         .map(|(name, (_, device))| (name.clone(), device.clone()))
@@ -115,8 +123,9 @@ pub fn all_devices() -> Vec<(String, NetworkDeviceRef)> {
 }
 
 static COMPONENT: Once<Component> = Once::new();
-pub(crate) static NETWORK_IRQ_HANDLERS: Once<SpinLock<Vec<Arc<dyn NetDeviceIrqHandler>>>> =
-    Once::new();
+pub(crate) static NETWORK_IRQ_HANDLERS: Once<
+    SpinLock<Vec<Arc<dyn NetDeviceIrqHandler>>, PreemptDisabled>,
+> = Once::new();
 
 #[init_component]
 fn init() -> Result<(), ComponentInitError> {
@@ -127,13 +136,13 @@ fn init() -> Result<(), ComponentInitError> {
     Ok(())
 }
 
-type NetDeviceIrqHandlerListRef = Arc<SpinLock<Vec<Arc<dyn NetDeviceIrqHandler>>>>;
-type NetworkDeviceRef = Arc<SpinLock<dyn AnyNetworkDevice>>;
+type NetDeviceIrqHandlerListRef = Arc<SpinLock<Vec<Arc<dyn NetDeviceIrqHandler>>, PreemptDisabled>>;
+type NetworkDeviceRef = Arc<SpinLock<dyn AnyNetworkDevice, PreemptDisabled>>;
 
 struct Component {
     /// Device list, the key is device name, value is (callbacks, device);
     network_device_table:
-        SpinLock<BTreeMap<String, (NetDeviceIrqHandlerListRef, NetworkDeviceRef)>>,
+        SpinLock<BTreeMap<String, (NetDeviceIrqHandlerListRef, NetworkDeviceRef)>, PreemptDisabled>,
 }
 
 impl Component {