TEST: use queue opt for block device

Signed-off-by: Egor Lazarchuk <[email protected]>

Author: ShadowCurse

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_opt(index, len).unwrap_or_else(|err| {
             error!("Failed to add available descriptor head {}: {}", index, err)
         });
 
-        if queue.prepare_kick(mem) {
+        if queue.prepare_kick_opt() {
             irq_trigger.trigger_irq(IrqType::Vring).unwrap_or_else(|_| {
                 block_metrics.event_fails.inc();
             });
@@ -411,9 +410,11 @@ 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());
-            let processing_result = match Request::parse(&head, mem, self.disk.nsectors) {
+        while let Some(head) = queue.pop_or_enable_notification_opt() {
+            self.metrics
+                .remaining_reqs_count
+                .add(queue.len_opt().into());
+            let processing_result = match Request::parse_opt(&head, mem, self.disk.nsectors) {
                 Ok(request) => {
                     if request.rate_limit(&mut self.rate_limiter) {
                         // Stop processing the queue and return this descriptor chain to the
@@ -448,7 +449,6 @@ impl VirtioBlock {
                         queue,
                         head.index,
                         finished.num_bytes_to_mem,
-                        mem,
                         &self.irq_trigger,
                         &self.metrics,
                     );
@@ -500,7 +500,6 @@ impl VirtioBlock {
                         queue,
                         finished.desc_idx,
                         finished.num_bytes_to_mem,
-                        mem,
                         &self.irq_trigger,
                         &self.metrics,
                     );
@@ -633,6 +632,7 @@ impl VirtioDevice for VirtioBlock {
     }
 
     fn activate(&mut self, mem: GuestMemoryMmap) -> Result<(), ActivateError> {
+        self.queues[0].set_pointers(&mem);
         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

@@ -16,7 +16,7 @@ pub use crate::devices::virtio::gen::virtio_blk::{
     VIRTIO_BLK_ID_BYTES, VIRTIO_BLK_S_IOERR, VIRTIO_BLK_S_OK, VIRTIO_BLK_S_UNSUPP,
     VIRTIO_BLK_T_FLUSH, VIRTIO_BLK_T_GET_ID, VIRTIO_BLK_T_IN, VIRTIO_BLK_T_OUT,
 };
-use crate::devices::virtio::queue::DescriptorChain;
+use crate::devices::virtio::queue::{DescriptorChain, DescriptorChainOpt};
 use crate::logger::{error, IncMetric};
 use crate::rate_limiter::{RateLimiter, TokenType};
 use crate::vstate::memory::{ByteValued, Bytes, GuestAddress, GuestMemoryMmap};
@@ -330,6 +330,95 @@ impl Request {
         Ok(req)
     }
 
+    pub fn parse_opt(
+        avail_desc: &DescriptorChainOpt,
+        mem: &GuestMemoryMmap,
+        num_disk_sectors: u64,
+    ) -> Result<Request, VirtioBlockError> {
+        // The head contains the request type which MUST be readable.
+        if avail_desc.is_write_only() {
+            return Err(VirtioBlockError::UnexpectedWriteOnlyDescriptor);
+        }
+
+        let request_header = RequestHeader::read_from(mem, avail_desc.addr)?;
+        let mut req = Request {
+            r#type: RequestType::from(request_header.request_type),
+            sector: request_header.sector,
+            data_addr: GuestAddress(0),
+            data_len: 0,
+            status_addr: GuestAddress(0),
+        };
+
+        let data_desc;
+        let status_desc;
+        let desc = avail_desc
+            .next_descriptor()
+            .ok_or(VirtioBlockError::DescriptorChainTooShort)?;
+
+        if !desc.has_next() {
+            status_desc = desc;
+            // Only flush requests are allowed to skip the data descriptor.
+            if req.r#type != RequestType::Flush {
+                return Err(VirtioBlockError::DescriptorChainTooShort);
+            }
+        } else {
+            data_desc = desc;
+            status_desc = data_desc
+                .next_descriptor()
+                .ok_or(VirtioBlockError::DescriptorChainTooShort)?;
+
+            if data_desc.is_write_only() && req.r#type == RequestType::Out {
+                return Err(VirtioBlockError::UnexpectedWriteOnlyDescriptor);
+            }
+            if !data_desc.is_write_only() && req.r#type == RequestType::In {
+                return Err(VirtioBlockError::UnexpectedReadOnlyDescriptor);
+            }
+            if !data_desc.is_write_only() && req.r#type == RequestType::GetDeviceID {
+                return Err(VirtioBlockError::UnexpectedReadOnlyDescriptor);
+            }
+
+            req.data_addr = data_desc.addr;
+            req.data_len = data_desc.len;
+        }
+
+        // check request validity
+        match req.r#type {
+            RequestType::In | RequestType::Out => {
+                // Check that the data length is a multiple of 512 as specified in the virtio
+                // standard.
+                if req.data_len % SECTOR_SIZE != 0 {
+                    return Err(VirtioBlockError::InvalidDataLength);
+                }
+                let top_sector = req
+                    .sector
+                    .checked_add(u64::from(req.data_len) >> SECTOR_SHIFT)
+                    .ok_or(VirtioBlockError::InvalidOffset)?;
+                if top_sector > num_disk_sectors {
+                    return Err(VirtioBlockError::InvalidOffset);
+                }
+            }
+            RequestType::GetDeviceID => {
+                if req.data_len < VIRTIO_BLK_ID_BYTES {
+                    return Err(VirtioBlockError::InvalidDataLength);
+                }
+            }
+            _ => {}
+        }
+
+        // The status MUST always be writable.
+        if !status_desc.is_write_only() {
+            return Err(VirtioBlockError::UnexpectedReadOnlyDescriptor);
+        }
+
+        if status_desc.len < 1 {
+            return Err(VirtioBlockError::DescriptorLengthTooSmall);
+        }
+
+        req.status_addr = status_desc.addr;
+
+        Ok(req)
+    }
+
     pub(crate) fn rate_limit(&self, rate_limiter: &mut RateLimiter) -> bool {
         // If limiter.consume() fails it means there is no more TokenType::Ops
         // budget and rate limiting is in effect.

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

@@ -620,6 +620,21 @@ impl Queue {
         self.do_pop_unchecked(mem)
     }
 
+    pub fn pop_opt(&mut self) -> Option<DescriptorChainOpt> {
+        let len = self.len_opt();
+        if len > self.actual_size() {
+            panic!(
+                "The number of available virtio descriptors {len} is greater than queue size: {}!",
+                self.actual_size()
+            );
+        }
+        if len == 0 {
+            return None;
+        }
+
+        self.do_pop_unchecked_opt()
+    }
+
     /// Try to pop the first available descriptor chain from the avail ring.
     /// If no descriptor is available, enable notifications.
     pub fn pop_or_enable_notification<'b, M: GuestMemory>(
@@ -637,6 +652,18 @@ impl Queue {
         self.do_pop_unchecked(mem)
     }
 
+    pub fn pop_or_enable_notification_opt(&mut self) -> Option<DescriptorChainOpt> {
+        if !self.uses_notif_suppression {
+            return self.pop_opt();
+        }
+
+        if self.try_enable_notification_opt() {
+            return None;
+        }
+
+        self.do_pop_unchecked_opt()
+    }
+
     /// Pop the first available descriptor chain from the avail ring.
     ///
     /// # Important
@@ -723,6 +750,16 @@ impl Queue {
         Ok(())
     }
 
+    pub fn add_used_opt(&mut self, desc_index: u16, len: u32) -> Result<(), QueueError> {
+        let used_element = UsedElement {
+            id: u32::from(desc_index),
+            len,
+        };
+        self.write_used_ring_opt(self.next_used.0, used_element)?;
+        self.advance_used_ring_opt(1);
+        Ok(())
+    }
+
     /// Advance number of used descriptor heads by `n`.
     pub fn advance_used_ring<M: GuestMemory>(&mut self, mem: &M, n: u16) {
         self.num_added += Wrapping(n);
@@ -963,6 +1000,33 @@ impl Queue {
         self.next_avail.0 == self.avail_idx(mem).0
     }
 
+    pub fn try_enable_notification_opt(&mut self) -> bool {
+        if !self.uses_notif_suppression {
+            return true;
+        }
+
+        let len = self.len_opt();
+        if len != 0 {
+            if len > self.actual_size() {
+                panic!(
+                    "The number of available virtio descriptors {len} is greater than queue size: \
+                     {}!",
+                    self.actual_size()
+                );
+            }
+            return false;
+        }
+
+        *self.ur.avail_event() = self.next_avail.0;
+
+        // Make sure all subsequent reads are performed after `set_used_ring_avail_event`.
+        fence(Ordering::SeqCst);
+
+        // If the actual avail_idx is different than next_avail one or more descriptors can still
+        // be consumed from the available ring.
+        self.next_avail.0 == *self.ar.idx()
+    }
+
     /// Enable notification suppression.
     pub fn enable_notif_suppression(&mut self) {
         self.uses_notif_suppression = true;
@@ -994,6 +1058,24 @@ impl Queue {
 
         new - used_event - Wrapping(1) < new - old
     }
+
+    pub fn prepare_kick_opt(&mut self) -> bool {
+        // If the device doesn't use notification suppression, always return true
+        if !self.uses_notif_suppression {
+            return true;
+        }
+
+        // We need to expose used array entries before checking the used_event.
+        fence(Ordering::SeqCst);
+
+        let new = self.next_used;
+        let old = self.next_used - self.num_added;
+        let used_event = Wrapping(*self.ar.used_event());
+
+        self.num_added = Wrapping(0);
+
+        new - used_event - Wrapping(1) < new - old
+    }
 }
 
 #[cfg(kani)]