Less memory checks in the queue

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(&mem).unwrap();
+    let desc = queue.pop().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, 1);
     queue.next_avail = Wrapping(0);
-    let desc = queue.pop(&mem).unwrap();
+    let desc = queue.pop().unwrap();
     c.bench_function("next_descriptor_1", |b| {
         b.iter(|| {
             let mut head = Some(desc.clone());
@@ -73,7 +73,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
 
     set_dtable_one_chain(&rxq, 2);
     queue.next_avail = Wrapping(0);
-    let desc = queue.pop(&mem).unwrap();
+    let desc = queue.pop().unwrap();
     c.bench_function("next_descriptor_2", |b| {
         b.iter(|| {
             let mut head = Some(desc.clone());
@@ -85,7 +85,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
 
     set_dtable_one_chain(&rxq, 4);
     queue.next_avail = Wrapping(0);
-    let desc = queue.pop(&mem).unwrap();
+    let desc = queue.pop().unwrap();
     c.bench_function("next_descriptor_4", |b| {
         b.iter(|| {
             let mut head = Some(desc.clone());
@@ -97,7 +97,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
 
     set_dtable_one_chain(&rxq, 16);
     queue.next_avail = Wrapping(0);
-    let desc = queue.pop(&mem).unwrap();
+    let desc = queue.pop().unwrap();
     c.bench_function("next_descriptor_16", |b| {
         b.iter(|| {
             let mut head = Some(desc.clone());
@@ -113,7 +113,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
     c.bench_function("queue_pop_1", |b| {
         b.iter(|| {
             queue.next_avail = Wrapping(0);
-            while let Some(desc) = queue.pop(&mem) {
+            while let Some(desc) = queue.pop() {
                 std::hint::black_box(desc);
             }
         })
@@ -123,7 +123,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
     c.bench_function("queue_pop_4", |b| {
         b.iter(|| {
             queue.next_avail = Wrapping(0);
-            while let Some(desc) = queue.pop(&mem) {
+            while let Some(desc) = queue.pop() {
                 std::hint::black_box(desc);
             }
         })
@@ -133,7 +133,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(&mem) {
+            while let Some(desc) = queue.pop() {
                 std::hint::black_box(desc);
             }
         })
@@ -146,7 +146,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
             for i in 0_u16..1_u16 {
                 let index = std::hint::black_box(i);
                 let len = std::hint::black_box(i + 1);
-                _ = queue.add_used(&mem, index as u16, len as u32);
+                _ = queue.add_used(index as u16, len as u32);
             }
         })
     });
@@ -158,7 +158,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
             for i in 0_u16..16_u16 {
                 let index = std::hint::black_box(i);
                 let len = std::hint::black_box(i + 1);
-                _ = queue.add_used(&mem, index as u16, len as u32);
+                _ = queue.add_used(index as u16, len as u32);
             }
         })
     });
@@ -170,7 +170,7 @@ pub fn queue_benchmark(c: &mut Criterion) {
             for i in 0_u16..256_u16 {
                 let index = std::hint::black_box(i);
                 let len = std::hint::black_box(i + 1);
-                _ = queue.add_used(&mem, index as u16, len as u32);
+                _ = queue.add_used(index as u16, len as u32);
             }
         })
     });

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

@@ -297,7 +297,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(mem) {
+            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,9 +339,7 @@
 
                 // Acknowledge the receipt of the descriptor.
                 // 0 is number of bytes the device has written to memory.
-                queue
-                    .add_used(mem, head.index, 0)
-                    .map_err(BalloonError::Queue)?;
+                queue.add_used(head.index, 0).map_err(BalloonError::Queue)?;
                 needs_interrupt = true;
             }
 
@@ -372,17 +370,13 @@
     }
 
     pub(crate) fn process_deflate_queue(&mut self) -> Result<(), BalloonError> {
-        // This is safe since we checked in the event handler that the device is activated.
-        let mem = self.device_state.mem().unwrap();
         METRICS.deflate_count.inc();
 
         let queue = &mut self.queues[DEFLATE_INDEX];
         let mut needs_interrupt = false;
 
-        while let Some(head) = queue.pop(mem) {
-            queue
-                .add_used(mem, head.index, 0)
-                .map_err(BalloonError::Queue)?;
+        while let Some(head) = queue.pop() {
+            queue.add_used(head.index, 0).map_err(BalloonError::Queue)?;
             needs_interrupt = true;
         }
 
@@ -398,13 +392,13 @@
         let mem = self.device_state.mem().unwrap();
         METRICS.stats_updates_count.inc();
 
-        while let Some(head) = self.queues[STATS_INDEX].pop(mem) {
+        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(mem, prev_stats_desc, 0)
+                    .add_used(prev_stats_desc, 0)
                     .map_err(BalloonError::Queue)?;
             }
             for index in (0..head.len).step_by(SIZE_OF_STAT) {
@@ -450,14 +444,11 @@
     }
 
     fn trigger_stats_update(&mut self) -> Result<(), BalloonError> {
-        // This is safe since we checked in the event handler that the device is activated.
-        let mem = self.device_state.mem().unwrap();
-
         // 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(mem, index, 0)
+                .add_used(index, 0)
                 .map_err(BalloonError::Queue)?;
             self.signal_used_queue()
         } else {
@@ -611,6 +602,11 @@
     }
 
     fn activate(&mut self, mem: GuestMemoryMmap) -> Result<(), ActivateError> {
+        for q in self.queues.iter_mut() {
+            q.initialize(&mem)
+                .map_err(ActivateError::QueueMemoryError)?;
+        }
+
         self.device_state = DeviceState::Activated(mem);
         if self.activate_evt.write(1).is_err() {
             METRICS.activate_fails.inc();

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

@@ -331,6 +331,11 @@ impl<T: VhostUserHandleBackend + Send + 'static> VirtioDevice for VhostUserBlock
     }
 
     fn activate(&mut self, mem: GuestMemoryMmap) -> Result<(), ActivateError> {
+        for q in self.queues.iter_mut() {
+            q.initialize(&mem)
+                .map_err(ActivateError::QueueMemoryError)?;
+        }
+
         let start_time = utils::time::get_time_us(utils::time::ClockType::Monotonic);
         // Setting features again, because now we negotiated them
         // with guest driver as well.

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

@@ -388,15 +388,14 @@ impl VirtioBlock {
         queue: &mut Queue,
         index: u16,
         len: u32,
-        mem: &GuestMemoryMmap,
         irq_trigger: &IrqTrigger,
         block_metrics: &BlockDeviceMetrics,
     ) {
-        queue.add_used(mem, index, len).unwrap_or_else(|err| {
+        queue.add_used(index, len).unwrap_or_else(|err| {
             error!("Failed to add available descriptor head {}: {}", index, err)
         });
 
-        if queue.prepare_kick(mem) {
+        if queue.prepare_kick() {
             irq_trigger.trigger_irq(IrqType::Vring).unwrap_or_else(|_| {
                 block_metrics.event_fails.inc();
             });
@@ -411,8 +410,8 @@ impl VirtioBlock {
         let queue = &mut self.queues[queue_index];
         let mut used_any = false;
 
-        while let Some(head) = queue.pop_or_enable_notification(mem) {
-            self.metrics.remaining_reqs_count.add(queue.len(mem).into());
+        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) => {
                     if request.rate_limit(&mut self.rate_limiter) {
@@ -448,7 +447,6 @@ impl VirtioBlock {
                         queue,
                         head.index,
                         finished.num_bytes_to_mem,
-                        mem,
                         &self.irq_trigger,
                         &self.metrics,
                     );
@@ -500,7 +498,6 @@ impl VirtioBlock {
                         queue,
                         finished.desc_idx,
                         finished.num_bytes_to_mem,
-                        mem,
                         &self.irq_trigger,
                         &self.metrics,
                     );
@@ -633,6 +630,11 @@ impl VirtioDevice for VirtioBlock {
     }
 
     fn activate(&mut self, mem: GuestMemoryMmap) -> Result<(), ActivateError> {
+        for q in self.queues.iter_mut() {
+            q.initialize(&mem)
+                .map_err(ActivateError::QueueMemoryError)?;
+        }
+
         let event_idx = self.has_feature(u64::from(VIRTIO_RING_F_EVENT_IDX));
         if event_idx {
             for queue in &mut self.queues {

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

@@ -473,16 +473,15 @@ mod tests {
             let memory = self.driver_queue.memory();
 
             assert!(matches!(
-                Request::parse(&q.pop(memory).unwrap(), memory, NUM_DISK_SECTORS),
+                Request::parse(&q.pop().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(memory).unwrap(), memory, NUM_DISK_SECTORS).unwrap();
+            let request = Request::parse(&q.pop().unwrap(), memory, NUM_DISK_SECTORS).unwrap();
             let expected_header = self.header();
 
             assert_eq!(
@@ -949,7 +948,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(&request.1).unwrap(), &request.1, NUM_DISK_SECTORS);
+            let result = Request::parse(&request.2.pop().unwrap(), &request.1, NUM_DISK_SECTORS);
             match result {
                 Ok(r) => prop_assert!(r == request.0.unwrap()),
                 Err(err) => {

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

@@ -12,7 +12,7 @@
 use utils::eventfd::EventFd;
 
 use super::mmio::{VIRTIO_MMIO_INT_CONFIG, VIRTIO_MMIO_INT_VRING};
-use super::queue::Queue;
+use super::queue::{Queue, QueueError};
 use super::ActivateError;
 use crate::devices::virtio::AsAny;
 use crate::logger::{error, warn};
@@ -180,6 +180,14 @@
     fn reset(&mut self) -> Option<(EventFd, Vec<EventFd>)> {
         None
     }
+
+    /// Mark pages used by queues as dirty.
+    fn mark_queue_memory_dirty(&self, mem: &GuestMemoryMmap) -> Result<(), QueueError> {
+        for queue in self.queues() {
+            queue.mark_memory_dirty(mem)?
+        }
+        Ok(())
+    }
 }
 
 impl fmt::Debug for dyn VirtioDevice {