Fix virtio-block request ID exhaustion panics with SyncIdAlloc

kernel/comps/virtio/src/device/block/device.rs

@@ -20,7 +20,6 @@ use aster_block::{
 };
 use aster_util::mem_obj_slice::Slice;
 use device_id::{DeviceId, MinorId};
-use id_alloc::IdAlloc;
 use log::{debug, info};
 use ostd::{
     arch::trap::TrapFrame,
@@ -35,6 +34,7 @@ use crate::{
         block::{ReqType, RespStatus},
         VirtioDeviceError,
     },
+    id_alloc::SyncIdAlloc,
     queue::VirtQueue,
     transport::{ConfigManager, VirtioTransport},
     VIRTIO_BLOCK_MAJOR_ID,
@@ -206,7 +206,7 @@ struct DeviceInner {
     transport: SpinLock<Box<dyn VirtioTransport>>,
     block_requests: Arc<DmaStream>,
     block_responses: Arc<DmaStream>,
-    id_allocator: SpinLock<IdAlloc>,
+    id_allocator: SyncIdAlloc,
     submitted_requests: SpinLock<BTreeMap<u16, SubmittedRequest>>,
 }
 
@@ -254,7 +254,7 @@ impl DeviceInner {
             transport: SpinLock::new(transport),
             block_requests,
             block_responses,
-            id_allocator: SpinLock::new(IdAlloc::with_capacity(Self::QUEUE_SIZE as usize)),
+            id_allocator: SyncIdAlloc::with_capacity(Self::QUEUE_SIZE as usize),
             submitted_requests: SpinLock::new(BTreeMap::new()),
         });
 
@@ -304,7 +304,7 @@ impl DeviceInner {
                 Slice::new(&self.block_responses, id * RESP_SIZE..(id + 1) * RESP_SIZE);
             resp_slice.sync().unwrap();
             let resp: BlockResp = resp_slice.read_val(0).unwrap();
-            self.id_allocator.lock().free(id);
+            self.id_allocator.dealloc(id);
             match RespStatus::try_from(resp.status).unwrap() {
                 RespStatus::Ok => {}
                 // FIXME: Return an error instead of triggering a kernel panic
@@ -337,7 +337,7 @@ impl DeviceInner {
 
     /// Reads data from the device, this function is non-blocking.
     fn read(&self, bio_request: BioRequest) {
-        let id = self.id_allocator.disable_irq().lock().alloc().unwrap();
+        let id = self.id_allocator.alloc();
         let req_slice = {
             let req_slice = Slice::new(
                 self.block_requests.clone(),
@@ -404,7 +404,7 @@ impl DeviceInner {
 
     /// Writes data to the device, this function is non-blocking.
     fn write(&self, bio_request: BioRequest) {
-        let id = self.id_allocator.disable_irq().lock().alloc().unwrap();
+        let id = self.id_allocator.alloc();
         let req_slice = {
             let req_slice = Slice::new(
                 self.block_requests.clone(),
@@ -478,7 +478,7 @@ impl DeviceInner {
             return;
         }
 
-        let id = self.id_allocator.disable_irq().lock().alloc().unwrap();
+        let id = self.id_allocator.alloc();
         let req_slice = {
             let req_slice = Slice::new(&self.block_requests, id * REQ_SIZE..(id + 1) * REQ_SIZE);
             let req = BlockReq {

kernel/comps/virtio/src/id_alloc.rs

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: MPL-2.0
+
+use id_alloc::IdAlloc;
+use ostd::sync::{SpinLock, WaitQueue};
+
+/// A synchronized ID allocator that manages a fixed-size pool of IDs.
+///
+/// We say this ID allocator "synchronized" because
+/// its `alloc` method may be used by multiple threads to allocate IDs concurrently
+/// and it would block if no free IDs are available at the moment.
+/// Once an ID is no longer in use, the `dealloc` method may be called
+/// to return the ID to the ID allocator.
+///
+/// This ID allocator is designed for use by device drivers.
+/// The `alloc` method should only be called in the task context,
+/// whereas the `dealloc` method should only be used in the IRQ handler of a device driver.
+/// Failing to conform with the above requirement may result in deadlock.
+pub struct SyncIdAlloc {
+    wait_queue: WaitQueue,
+    id_allocator: SpinLock<IdAlloc>,
+}
+
+impl SyncIdAlloc {
+    /// Creates an allocator that may return IDs in the range of `0..capacity`.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            wait_queue: WaitQueue::new(),
+            id_allocator: SpinLock::new(IdAlloc::with_capacity(capacity)),
+        }
+    }
+
+    /// Allocates a new ID.
+    ///
+    /// This method must only be called in the task context.
+    /// It will block until an ID is free to allocate.
+    pub fn alloc(&self) -> usize {
+        self.wait_queue
+            .wait_until(|| self.id_allocator.disable_irq().lock().alloc())
+    }
+
+    /// Deallocates an ID.
+    ///
+    /// This method assumes that the caller is in the context of an IRQ handler.
+    ///
+    /// # Panics
+    ///
+    /// This method would panic if `id` is greater than or equal to the capacity of this ID allocator.
+    pub fn dealloc(&self, id: usize) {
+        self.id_allocator.disable_irq().lock().free(id);
+        self.wait_queue.wake_all();
+    }
+}
+
+impl core::fmt::Debug for SyncIdAlloc {
+    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
+        f.debug_struct("SyncIdAlloc")
+            .field("id_allocator", &self.id_allocator.lock())
+            .finish_non_exhaustive()
+    }
+}

kernel/comps/virtio/src/lib.rs

@@ -30,6 +30,7 @@ use crate::transport::VirtioTransport;
 
 pub mod device;
 mod dma_buf;
+mod id_alloc;
 pub mod queue;
 mod transport;