Improve vqueue abstractions: use iterators for VirtioDevice.queues()

Signed-off-by: Kornel Lesiński <[email protected]>

src/vmm/src/device_manager/mmio.rs

@@ -467,7 +467,7 @@ mod tests {
 
     use super::*;
     use crate::devices::virtio::device::VirtioDevice;
-    use crate::devices::virtio::queue::Queue;
+    use crate::devices::virtio::queue::{Queue, QueueIter, QueueIterMut};
     use crate::devices::virtio::ActivateError;
     use crate::vstate::memory::{GuestAddress, GuestMemoryExtension, GuestMemoryMmap};
     use crate::{builder, Vm};
@@ -535,12 +535,12 @@ mod tests {
             0
         }
 
-        fn queues(&self) -> &[Queue] {
-            &self.queues
+        fn queues(&self) -> QueueIter {
+            self.queues.iter()
         }
 
-        fn queues_mut(&mut self) -> &mut [Queue] {
-            &mut self.queues
+        fn queues_mut(&mut self) -> QueueIterMut {
+            self.queues.iter_mut()
         }
 
         fn queue_events(&self) -> &[EventFd] {

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

@@ -14,7 +14,6 @@ use utils::eventfd::EventFd;
 use utils::u64_to_usize;
 
 use super::super::device::{DeviceState, VirtioDevice};
-use super::super::queue::Queue;
 use super::super::{ActivateError, TYPE_BALLOON};
 use super::metrics::METRICS;
 use super::util::{compact_page_frame_numbers, remove_range};
@@ -30,6 +29,7 @@ use super::{
 use crate::devices::virtio::balloon::BalloonError;
 use crate::devices::virtio::device::{IrqTrigger, IrqType};
 use crate::devices::virtio::gen::virtio_blk::VIRTIO_F_VERSION_1;
+use crate::devices::virtio::queue::{Queue, QueueIter, QueueIterMut};
 use crate::logger::IncMetric;
 use crate::vstate::memory::{Address, ByteValued, Bytes, GuestAddress, GuestMemoryMmap};
 
@@ -573,12 +573,12 @@ impl VirtioDevice for Balloon {
         TYPE_BALLOON
     }
 
-    fn queues(&self) -> &[Queue] {
-        &self.queues
+    fn queues(&self) -> QueueIter {
+        self.queues.iter()
     }
 
-    fn queues_mut(&mut self) -> &mut [Queue] {
-        &mut self.queues
+    fn queues_mut(&mut self) -> QueueIterMut {
+        self.queues.iter_mut()
     }
 
     fn queue_events(&self) -> &[EventFd] {

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

@@ -209,7 +209,7 @@ mod tests {
         assert_eq!(restored_balloon.acked_features, balloon.acked_features);
         assert_eq!(restored_balloon.avail_features, balloon.avail_features);
         assert_eq!(restored_balloon.config_space, balloon.config_space);
-        assert_eq!(restored_balloon.queues(), balloon.queues());
+        assert!(restored_balloon.queues().eq(balloon.queues()));
         assert_eq!(
             restored_balloon.interrupt_status().load(Ordering::Relaxed),
             balloon.interrupt_status().load(Ordering::Relaxed)

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

@@ -12,8 +12,8 @@ use std::sync::Arc;
 use utils::eventfd::EventFd;
 
 use super::mmio::{VIRTIO_MMIO_INT_CONFIG, VIRTIO_MMIO_INT_VRING};
-use super::queue::Queue;
 use super::ActivateError;
+use crate::devices::virtio::queue::{QueueIter, QueueIterMut};
 use crate::devices::virtio::AsAny;
 use crate::logger::{error, warn};
 use crate::vstate::memory::GuestMemoryMmap;
@@ -111,10 +111,10 @@ pub trait VirtioDevice: AsAny + Send {
     fn device_type(&self) -> u32;
 
     /// Returns the device queues.
-    fn queues(&self) -> &[Queue];
+    fn queues(&self) -> QueueIter;
 
     /// Returns a mutable reference to the device queues.
-    fn queues_mut(&mut self) -> &mut [Queue];
+    fn queues_mut(&mut self) -> QueueIterMut;
 
     /// Returns the device queues event fds.
     fn queue_events(&self) -> &[EventFd];
@@ -190,6 +190,7 @@ impl fmt::Debug for dyn VirtioDevice {
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;
+    use crate::devices::virtio::queue::{QueueIter, QueueIterMut};
 
     impl IrqTrigger {
         pub fn has_pending_irq(&self, irq_type: IrqType) -> bool {
@@ -254,11 +255,11 @@ pub(crate) mod tests {
             todo!()
         }
 
-        fn queues(&self) -> &[Queue] {
+        fn queues(&self) -> QueueIter {
             todo!()
         }
 
-        fn queues_mut(&mut self) -> &mut [Queue] {
+        fn queues_mut(&mut self) -> QueueIterMut {
             todo!()
         }
 

src/vmm/src/devices/virtio/mmio.rs

@@ -97,10 +97,7 @@ impl MmioTransport {
     }
 
     fn are_queues_valid(&self) -> bool {
-        self.locked_device()
-            .queues()
-            .iter()
-            .all(|q| q.is_valid(&self.mem))
+        self.locked_device().queues().all(|q| q.is_valid(&self.mem))
     }
 
     fn with_queue<U, F>(&self, d: U, f: F) -> U
@@ -111,7 +108,7 @@ impl MmioTransport {
         match self
             .locked_device()
             .queues()
-            .get(self.queue_select as usize)
+            .nth(self.queue_select as usize)
         {
             Some(queue) => f(queue),
             None => d,
@@ -122,7 +119,7 @@ impl MmioTransport {
         if let Some(queue) = self
             .locked_device()
             .queues_mut()
-            .get_mut(self.queue_select as usize)
+            .nth(self.queue_select as usize)
         {
             f(queue);
             true
@@ -363,6 +360,7 @@ pub(crate) mod tests {
     use utils::u64_to_usize;
 
     use super::*;
+    use crate::devices::virtio::queue::{Queue, QueueIter, QueueIterMut};
     use crate::devices::virtio::ActivateError;
     use crate::vstate::memory::{GuestMemoryExtension, GuestMemoryMmap};
 
@@ -417,12 +415,12 @@ pub(crate) mod tests {
             123
         }
 
-        fn queues(&self) -> &[Queue] {
-            &self.queues
+        fn queues(&self) -> QueueIter {
+            self.queues.iter()
         }
 
-        fn queues_mut(&mut self) -> &mut [Queue] {
-            &mut self.queues
+        fn queues_mut(&mut self) -> QueueIterMut {
+            self.queues.iter_mut()
         }
 
         fn queue_events(&self) -> &[EventFd] {
@@ -482,15 +480,23 @@ pub(crate) mod tests {
         assert_eq!(d.with_queue(0, |q| q.max_size()), 16);
         assert!(d.with_queue_mut(|q| q.set_size(16)));
         assert_eq!(
-            d.locked_device().queues()[d.queue_select as usize].size(),
+            d.locked_device()
+                .queues()
+                .nth(d.queue_select as usize)
+                .unwrap()
+                .size(),
             16
         );
 
         d.queue_select = 1;
         assert_eq!(d.with_queue(0, |q| q.max_size()), 32);
         assert!(d.with_queue_mut(|q| q.set_size(16)));
         assert_eq!(
-            d.locked_device().queues()[d.queue_select as usize].size(),
+            d.locked_device()
+                .queues()
+                .nth(d.queue_select as usize)
+                .unwrap()
+                .size(),
             16
         );
 
@@ -673,43 +679,52 @@ pub(crate) mod tests {
         assert_eq!(d.queue_select, 3);
 
         d.queue_select = 0;
-        assert_eq!(d.locked_device().queues()[0].size(), 0);
+        assert_eq!(d.locked_device().queues().nth(0).unwrap().size(), 0);
         write_le_u32(&mut buf[..], 16);
         d.bus_write(0x38, &buf[..]);
-        assert_eq!(d.locked_device().queues()[0].size(), 16);
+        assert_eq!(d.locked_device().queues().nth(0).unwrap().size(), 16);
 
-        assert!(!d.locked_device().queues()[0].ready());
+        assert!(!d.locked_device().queues().nth(0).unwrap().ready());
         write_le_u32(&mut buf[..], 1);
         d.bus_write(0x44, &buf[..]);
-        assert!(d.locked_device().queues()[0].ready());
+        assert!(d.locked_device().queues().nth(0).unwrap().ready());
 
-        assert_eq!(d.locked_device().queues()[0].desc_table().0, 0);
+        assert_eq!(d.locked_device().queues().nth(0).unwrap().desc_table().0, 0);
         write_le_u32(&mut buf[..], 123);
         d.bus_write(0x80, &buf[..]);
-        assert_eq!(d.locked_device().queues()[0].desc_table().0, 123);
+        assert_eq!(
+            d.locked_device().queues().nth(0).unwrap().desc_table().0,
+            123
+        );
         d.bus_write(0x84, &buf[..]);
         assert_eq!(
-            d.locked_device().queues()[0].desc_table().0,
+            d.locked_device().queues().nth(0).unwrap().desc_table().0,
             123 + (123 << 32)
         );
 
-        assert_eq!(d.locked_device().queues()[0].avail_ring().0, 0);
+        assert_eq!(d.locked_device().queues().nth(0).unwrap().avail_ring().0, 0);
         write_le_u32(&mut buf[..], 124);
         d.bus_write(0x90, &buf[..]);
-        assert_eq!(d.locked_device().queues()[0].avail_ring().0, 124);
+        assert_eq!(
+            d.locked_device().queues().nth(0).unwrap().avail_ring().0,
+            124
+        );
         d.bus_write(0x94, &buf[..]);
         assert_eq!(
-            d.locked_device().queues()[0].avail_ring().0,
+            d.locked_device().queues().nth(0).unwrap().avail_ring().0,
             124 + (124 << 32)
         );
 
-        assert_eq!(d.locked_device().queues()[0].used_ring().0, 0);
+        assert_eq!(d.locked_device().queues().nth(0).unwrap().used_ring().0, 0);
         write_le_u32(&mut buf[..], 125);
         d.bus_write(0xa0, &buf[..]);
-        assert_eq!(d.locked_device().queues()[0].used_ring().0, 125);
+        assert_eq!(
+            d.locked_device().queues().nth(0).unwrap().used_ring().0,
+            125
+        );
         d.bus_write(0xa4, &buf[..]);
         assert_eq!(
-            d.locked_device().queues()[0].used_ring().0,
+            d.locked_device().queues().nth(0).unwrap().used_ring().0,
             125 + (125 << 32)
         );
 

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

@@ -33,7 +33,7 @@ use crate::devices::virtio::net::tap::Tap;
 use crate::devices::virtio::net::{
     gen, NetError, NetQueue, MAX_BUFFER_SIZE, NET_QUEUE_SIZES, RX_INDEX, TX_INDEX,
 };
-use crate::devices::virtio::queue::{DescriptorChain, Queue};
+use crate::devices::virtio::queue::{DescriptorChain, Queue, QueueIter, QueueIterMut};
 use crate::devices::virtio::{ActivateError, TYPE_NET};
 use crate::devices::{report_net_event_fail, DeviceError};
 use crate::dumbo::pdu::arp::ETH_IPV4_FRAME_LEN;
@@ -787,12 +787,12 @@ impl VirtioDevice for Net {
         TYPE_NET
     }
 
-    fn queues(&self) -> &[Queue] {
-        &self.queues
+    fn queues(&self) -> QueueIter {
+        self.queues.iter()
     }
 
-    fn queues_mut(&mut self) -> &mut [Queue] {
-        &mut self.queues
+    fn queues_mut(&mut self) -> QueueIterMut {
+        self.queues.iter_mut()
     }
 
     fn queue_events(&self) -> &[EventFd] {
@@ -845,7 +845,7 @@ impl VirtioDevice for Net {
     fn activate(&mut self, mem: GuestMemoryMmap) -> Result<(), ActivateError> {
         let event_idx = self.has_feature(u64::from(VIRTIO_RING_F_EVENT_IDX));
         if event_idx {
-            for queue in &mut self.queues {
+            for queue in self.queues_mut() {
                 queue.enable_notif_suppression();
             }
         }
@@ -2007,10 +2007,9 @@ pub mod tests {
         let net = th.net.lock().unwrap();
 
         // Test queues count (TX and RX).
-        let queues = net.queues();
-        assert_eq!(queues.len(), NET_QUEUE_SIZES.len());
-        assert_eq!(queues[RX_INDEX].size(), th.rxq.size());
-        assert_eq!(queues[TX_INDEX].size(), th.txq.size());
+        assert_eq!(net.queues().count(), NET_QUEUE_SIZES.len());
+        assert_eq!(net.queues().nth(RX_INDEX).unwrap().size(), th.rxq.size());
+        assert_eq!(net.queues().nth(TX_INDEX).unwrap().size(), th.txq.size());
 
         // Test corresponding queues events.
         assert_eq!(net.queue_events().len(), NET_QUEUE_SIZES.len());
@@ -2028,8 +2027,7 @@ pub mod tests {
         th.activate_net();
 
         let net = th.net();
-        let queues = net.queues();
-        assert!(queues[RX_INDEX].uses_notif_suppression());
-        assert!(queues[TX_INDEX].uses_notif_suppression());
+        assert!(net.queues().nth(RX_INDEX).unwrap().uses_notif_suppression());
+        assert!(net.queues().nth(TX_INDEX).unwrap().uses_notif_suppression());
     }
 }

src/vmm/src/devices/virtio/persist.rs

@@ -118,7 +118,7 @@ impl VirtioDeviceState {
             device_type: device.device_type(),
             avail_features: device.avail_features(),
             acked_features: device.acked_features(),
-            queues: device.queues().iter().map(Persist::save).collect(),
+            queues: device.queues().map(Persist::save).collect(),
             interrupt_status: device.interrupt_status().load(Ordering::Relaxed),
             interrupt_status_old: device.interrupt_status().load(Ordering::Relaxed) as usize,
             activated: device.is_activated(),

src/vmm/src/devices/virtio/queue.rs

@@ -36,6 +36,9 @@ pub enum QueueError {
     UsedRing(#[from] vm_memory::GuestMemoryError),
 }
 
+pub type QueueIter<'a> = std::slice::Iter<'a, Queue>;
+pub type QueueIterMut<'a> = std::slice::IterMut<'a, Queue>;
+
 /// A virtio descriptor constraints with C representative.
 #[repr(C)]
 #[derive(Default, Clone, Copy)]

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

@@ -15,6 +15,7 @@ use crate::devices::virtio::device::{DeviceState, IrqTrigger, IrqType, VirtioDev
 use crate::devices::virtio::gen::virtio_rng::VIRTIO_F_VERSION_1;
 use crate::devices::virtio::iovec::IoVecBufferMut;
 use crate::devices::virtio::queue::{Queue, FIRECRACKER_MAX_QUEUE_SIZE};
+use crate::devices::virtio::queue::{QueueIter, QueueIterMut};
 use crate::devices::virtio::{ActivateError, TYPE_RNG};
 use crate::devices::DeviceError;
 use crate::logger::{debug, error, IncMetric};
@@ -246,12 +247,12 @@ impl VirtioDevice for Entropy {
         TYPE_RNG
     }
 
-    fn queues(&self) -> &[Queue] {
-        &self.queues
+    fn queues(&self) -> QueueIter {
+        self.queues.iter()
     }
 
-    fn queues_mut(&mut self) -> &mut [Queue] {
-        &mut self.queues
+    fn queues_mut(&mut self) -> QueueIterMut {
+        self.queues.iter_mut()
     }
 
     fn queue_events(&self) -> &[EventFd] {
@@ -430,14 +431,24 @@ mod tests {
         let mut entropy_dev = th.device();
 
         // This should succeed, we just added two descriptors
-        let desc = entropy_dev.queues_mut()[RNG_QUEUE].pop(&mem).unwrap();
+        let desc = entropy_dev
+            .queues_mut()
+            .nth(RNG_QUEUE)
+            .unwrap()
+            .pop(&mem)
+            .unwrap();
         assert!(matches!(
             IoVecBufferMut::from_descriptor_chain(&mem, desc,),
             Err(crate::devices::virtio::iovec::IoVecError::ReadOnlyDescriptor)
         ));
 
         // This should succeed, we should have one more descriptor
-        let desc = entropy_dev.queues_mut()[RNG_QUEUE].pop(&mem).unwrap();
+        let desc = entropy_dev
+            .queues_mut()
+            .nth(RNG_QUEUE)
+            .unwrap()
+            .pop(&mem)
+            .unwrap();
         let mut iovec = IoVecBufferMut::from_descriptor_chain(&mem, desc).unwrap();
         assert!(entropy_dev.handle_one(&mut iovec).is_ok());
     }