refactor: Hoist panic out of virtio queue code

src/vmm/benches/block_request.rs

@@ -24,7 +24,7 @@ pub fn block_request_benchmark(c: &mut Criterion) {
     chain.set_header(request_header);
 
     let mut queue = virt_queue.create_queue();
-    let desc = queue.pop().unwrap();
+    let desc = queue.pop().unwrap().unwrap();
 
     c.bench_function("request_parse", |b| {
         b.iter(|| {

src/vmm/benches/queue.rs

@@ -61,7 +61,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
 
     set_dtable_one_chain(&rxq, 16);
     queue.next_avail = Wrapping(0);
-    let desc = queue.pop().unwrap();
+    let desc = queue.pop().unwrap().unwrap();
     c.bench_function("next_descriptor_16", |b| {
         b.iter(|| {
             let mut head = Some(desc);
@@ -75,7 +75,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
     c.bench_function("queue_pop_16", |b| {
         b.iter(|| {
             queue.next_avail = Wrapping(0);
-            while let Some(desc) = queue.pop() {
+            while let Some(desc) = queue.pop().unwrap() {
                 std::hint::black_box(desc);
             }
         })

src/vmm/src/device_manager/mmio.rs

@@ -461,7 +461,7 @@
                         // Stats queue doesn't need kicking as it is notified via a `timer_fd`.
                         if balloon.is_activated() {
                             info!("kick balloon {}.", id);
-                            balloon.process_virtio_queues();
+                            balloon.process_virtio_queues().unwrap();
                         }
                     }
                     TYPE_BLOCK => {
@@ -475,7 +475,7 @@
                             // any inflight `timer_fd` events can be safely discarded.
                             if block.is_activated() {
                                 info!("kick block {}.", id);
-                                block.process_virtio_queues();
+                                block.process_virtio_queues().unwrap()
                             }
                         }
                     }
@@ -487,7 +487,7 @@
                         // any inflight `timer_fd` events can be safely discarded.
                         if net.is_activated() {
                             info!("kick net {}.", id);
-                            net.process_virtio_queues();
+                            net.process_virtio_queues().unwrap();
                         }
                     }
                     TYPE_VSOCK => {
@@ -510,7 +510,7 @@
                         let entropy = virtio.as_mut_any().downcast_mut::<Entropy>().unwrap();
                         if entropy.is_activated() {
                             info!("kick entropy {id}.");
-                            entropy.process_virtio_queues();
+                            entropy.process_virtio_queues().unwrap();
                         }
                     }
                     _ => (),

src/vmm/src/devices/mod.rs

@@ -19,7 +19,7 @@
 use log::error;
 
 use crate::devices::virtio::net::metrics::NetDeviceMetrics;
-use crate::devices::virtio::queue::QueueError;
+use crate::devices::virtio::queue::{InvalidAvailIdx, QueueError};
 use crate::devices::virtio::vsock::VsockError;
 use crate::logger::IncMetric;
 
@@ -28,6 +28,9 @@
 // network metrics is reported per device so we need a handle to each net device's
 // metrics `net_iface_metrics` to report metrics for that device.
 pub(crate) fn report_net_event_fail(net_iface_metrics: &NetDeviceMetrics, err: DeviceError) {
+    if let DeviceError::InvalidAvailIdx(err) = err {
+        panic!("{}", err);
+    }
     error!("{:?}", err);
     net_iface_metrics.event_fails.inc();
 }
@@ -45,7 +48,9 @@
     /// Device received malformed descriptor.
     MalformedDescriptor,
     /// Error during queue processing: {0}
-    QueueError(QueueError),
+    QueueError(#[from] QueueError),
+    /// {0}
+    InvalidAvailIdx(#[from] InvalidAvailIdx),
     /// Vsock device error: {0}
-    VsockError(VsockError),
+    VsockError(#[from] VsockError),
 }

src/vmm/src/devices/virtio/balloon/device.rs

@@ -26,6 +26,7 @@
 use crate::devices::virtio::balloon::BalloonError;
 use crate::devices::virtio::device::{IrqTrigger, IrqType};
 use crate::devices::virtio::generated::virtio_config::VIRTIO_F_VERSION_1;
+use crate::devices::virtio::queue::InvalidAvailIdx;
 use crate::logger::IncMetric;
 use crate::utils::u64_to_usize;
 use crate::vstate::memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemoryMmap};
@@ -297,7 +298,7 @@
             // Internal loop processes descriptors and acummulates the pfns in `pfn_buffer`.
             // Breaks out when there is not enough space in `pfn_buffer` to completely process
             // the next descriptor.
-            while let Some(head) = queue.pop() {
+            while let Some(head) = queue.pop()? {
                 let len = head.len as usize;
                 let max_len = MAX_PAGES_IN_DESC * SIZE_OF_U32;
                 valid_descs_found = true;
@@ -339,7 +340,7 @@
 
                 // Acknowledge the receipt of the descriptor.
                 // 0 is number of bytes the device has written to memory.
-                queue.add_used(head.index, 0).map_err(BalloonError::Queue)?;
+                queue.add_used(head.index, 0)?;
                 needs_interrupt = true;
             }
 
@@ -375,8 +376,8 @@
         let queue = &mut self.queues[DEFLATE_INDEX];
         let mut needs_interrupt = false;
 
-        while let Some(head) = queue.pop() {
-            queue.add_used(head.index, 0).map_err(BalloonError::Queue)?;
+        while let Some(head) = queue.pop()? {
+            queue.add_used(head.index, 0)?;
             needs_interrupt = true;
         }
 
@@ -392,14 +393,12 @@
         let mem = self.device_state.mem().unwrap();
         METRICS.stats_updates_count.inc();
 
-        while let Some(head) = self.queues[STATS_INDEX].pop() {
+        while let Some(head) = self.queues[STATS_INDEX].pop()? {
             if let Some(prev_stats_desc) = self.stats_desc_index {
                 // We shouldn't ever have an extra buffer if the driver follows
                 // the protocol, but return it if we find one.
                 error!("balloon: driver is not compliant, more than one stats buffer received");
-                self.queues[STATS_INDEX]
-                    .add_used(prev_stats_desc, 0)
-                    .map_err(BalloonError::Queue)?;
+                self.queues[STATS_INDEX].add_used(prev_stats_desc, 0)?;
             }
             for index in (0..head.len).step_by(SIZE_OF_STAT) {
                 // Read the address at position `index`. The only case
@@ -433,9 +432,15 @@
     }
 
     /// Process device virtio queue(s).
-    pub fn process_virtio_queues(&mut self) {
-        let _ = self.process_inflate();
-        let _ = self.process_deflate_queue();
+    pub fn process_virtio_queues(&mut self) -> Result<(), InvalidAvailIdx> {
+        if let Err(BalloonError::InvalidAvailIdx(err)) = self.process_inflate() {
+            return Err(err);
+        }
+        if let Err(BalloonError::InvalidAvailIdx(err)) = self.process_deflate_queue() {
+            return Err(err);
+        }
+
+        Ok(())
     }
 
     /// Provides the ID of this balloon device.
@@ -447,9 +452,7 @@
         // The communication is driven by the device by using the buffer
         // and sending a used buffer notification
         if let Some(index) = self.stats_desc_index.take() {
-            self.queues[STATS_INDEX]
-                .add_used(index, 0)
-                .map_err(BalloonError::Queue)?;
+            self.queues[STATS_INDEX].add_used(index, 0)?;
             self.signal_used_queue()
         } else {
             error!("Failed to update balloon stats, missing descriptor.");
@@ -1084,7 +1087,7 @@
         balloon.set_queue(DEFLATE_INDEX, defq.create_queue());
 
         balloon.activate(mem).unwrap();
-        balloon.process_virtio_queues()
+        balloon.process_virtio_queues().unwrap();
     }
 
     #[test]

src/vmm/src/devices/virtio/balloon/mod.rs

@@ -11,10 +11,9 @@
 mod util;
 
 use log::error;
-use vm_memory::GuestMemoryError;
 
 pub use self::device::{Balloon, BalloonConfig, BalloonStats};
-use super::queue::QueueError;
+use super::queue::{InvalidAvailIdx, QueueError};
 use crate::devices::virtio::balloon::metrics::METRICS;
 use crate::devices::virtio::queue::FIRECRACKER_MAX_QUEUE_SIZE;
 use crate::logger::IncMetric;
@@ -68,16 +67,12 @@
 /// Balloon device related errors.
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum BalloonError {
-    /// Activation error: {0}
-    Activate(super::ActivateError),
     /// No balloon device found.
     DeviceNotFound,
     /// Device not activated yet.
     DeviceNotActive,
     /// EventFd error: {0}
     EventFd(std::io::Error),
-    /// Guest gave us bad memory addresses: {0}
-    GuestMemory(GuestMemoryError),
     /// Received error while sending an interrupt: {0}
     InterruptError(std::io::Error),
     /// Guest gave us a malformed descriptor.
@@ -93,9 +88,9 @@
     /// Amount of pages requested cannot fit in `u32`.
     TooManyPagesRequested,
     /// Error while processing the virt queues: {0}
-    Queue(QueueError),
-    /// Error removing a memory region at inflate time: {0}
-    RemoveMemoryRegion(RemoveRegionError),
+    Queue(#[from] QueueError),
+    /// {0}
+    InvalidAvailIdx(#[from] InvalidAvailIdx),
     /// Error creating the statistics timer: {0}
     Timer(std::io::Error),
 }
@@ -115,6 +110,9 @@
 }
 
 pub(super) fn report_balloon_event_fail(err: BalloonError) {
+    if let BalloonError::InvalidAvailIdx(err) = err {
+        panic!("{}", err);
+    }
     error!("{:?}", err);
     METRICS.event_fails.inc();
 }

src/vmm/src/devices/virtio/block/device.rs

@@ -9,7 +9,7 @@
 use super::vhost_user::device::{VhostUserBlock, VhostUserBlockConfig};
 use super::virtio::device::{VirtioBlock, VirtioBlockConfig};
 use crate::devices::virtio::device::{IrqTrigger, VirtioDevice};
-use crate::devices::virtio::queue::Queue;
+use crate::devices::virtio::queue::{InvalidAvailIdx, Queue};
 use crate::devices::virtio::{ActivateError, TYPE_BLOCK};
 use crate::rate_limiter::BucketUpdate;
 use crate::snapshot::Persist;
@@ -83,10 +83,10 @@
         }
     }
 
-    pub fn process_virtio_queues(&mut self) {
+    pub fn process_virtio_queues(&mut self) -> Result<(), InvalidAvailIdx> {
         match self {
             Self::Virtio(b) => b.process_virtio_queues(),
-            Self::VhostUser(_) => {}
+            Self::VhostUser(_) => Ok(()),
         }
     }
 

src/vmm/src/devices/virtio/block/virtio/device.rs

@@ -29,7 +29,7 @@ use crate::devices::virtio::generated::virtio_blk::{
 };
 use crate::devices::virtio::generated::virtio_config::VIRTIO_F_VERSION_1;
 use crate::devices::virtio::generated::virtio_ring::VIRTIO_RING_F_EVENT_IDX;
-use crate::devices::virtio::queue::Queue;
+use crate::devices::virtio::queue::{InvalidAvailIdx, Queue};
 use crate::devices::virtio::{ActivateError, TYPE_BLOCK};
 use crate::logger::{IncMetric, error, warn};
 use crate::rate_limiter::{BucketUpdate, RateLimiter};
@@ -366,21 +366,21 @@ impl VirtioBlock {
         } else if self.is_io_engine_throttled {
             self.metrics.io_engine_throttled_events.inc();
         } else {
-            self.process_virtio_queues();
+            self.process_virtio_queues().unwrap()
         }
     }
 
     /// Process device virtio queue(s).
-    pub fn process_virtio_queues(&mut self) {
-        self.process_queue(0);
+    pub fn process_virtio_queues(&mut self) -> Result<(), InvalidAvailIdx> {
+        self.process_queue(0)
     }
 
     pub(crate) fn process_rate_limiter_event(&mut self) {
         self.metrics.rate_limiter_event_count.inc();
         // Upon rate limiter event, call the rate limiter handler
         // and restart processing the queue.
         if self.rate_limiter.event_handler().is_ok() {
-            self.process_queue(0);
+            self.process_queue(0).unwrap()
         }
     }
 
@@ -403,14 +403,14 @@ impl VirtioBlock {
     }
 
     /// Device specific function for peaking inside a queue and processing descriptors.
-    pub fn process_queue(&mut self, queue_index: usize) {
+    pub fn process_queue(&mut self, queue_index: usize) -> Result<(), InvalidAvailIdx> {
         // This is safe since we checked in the event handler that the device is activated.
         let mem = self.device_state.mem().unwrap();
 
         let queue = &mut self.queues[queue_index];
         let mut used_any = false;
 
-        while let Some(head) = queue.pop_or_enable_notification() {
+        while let Some(head) = queue.pop_or_enable_notification()? {
             self.metrics.remaining_reqs_count.add(queue.len().into());
             let processing_result = match Request::parse(&head, mem, self.disk.nsectors) {
                 Ok(request) => {
@@ -463,6 +463,8 @@ impl VirtioBlock {
         if !used_any {
             self.metrics.no_avail_buffer.inc();
         }
+
+        Ok(())
     }
 
     fn process_async_completion_queue(&mut self) {
@@ -516,7 +518,7 @@ impl VirtioBlock {
 
             if self.is_io_engine_throttled {
                 self.is_io_engine_throttled = false;
-                self.process_queue(0);
+                self.process_queue(0).unwrap()
             }
         }
     }

src/vmm/src/devices/virtio/block/virtio/request.rs

@@ -484,15 +484,16 @@ mod tests {
             let memory = self.driver_queue.memory();
 
             assert!(matches!(
-                Request::parse(&q.pop().unwrap(), memory, NUM_DISK_SECTORS),
+                Request::parse(&q.pop().unwrap().unwrap(), memory, NUM_DISK_SECTORS),
                 Err(_e)
             ));
         }
 
         fn check_parse(&self, check_data: bool) {
             let mut q = self.driver_queue.create_queue();
             let memory = self.driver_queue.memory();
-            let request = Request::parse(&q.pop().unwrap(), memory, NUM_DISK_SECTORS).unwrap();
+            let request =
+                Request::parse(&q.pop().unwrap().unwrap(), memory, NUM_DISK_SECTORS).unwrap();
             let expected_header = self.header();
 
             assert_eq!(
@@ -959,7 +960,7 @@ mod tests {
     fn parse_random_requests() {
         let cfg = ProptestConfig::with_cases(1000);
         proptest!(cfg, |(mut request in random_request_parse())| {
-            let result = Request::parse(&request.2.pop().unwrap(), &request.1, NUM_DISK_SECTORS);
+            let result = Request::parse(&request.2.pop().unwrap().unwrap(), &request.1, NUM_DISK_SECTORS);
             match result {
                 Ok(r) => prop_assert!(r == request.0.unwrap()),
                 Err(err) => {