VirtQueue should provide peek_avail

lib/propolis/src/hw/virtio/queue.rs

@@ -88,26 +88,30 @@ pub struct VqAvail {
 }
 
 impl VqAvail {
-    /// If there's a request ready, pop it off the queue and return the
-    /// corresponding descriptor and available ring indicies.
-    fn read_next_avail(&mut self, rsize: u16, mem: &MemCtx) -> Option<VqReq> {
+    /// Read the entry at the front of the avail ring without consuming it. The
+    /// same entry will remain visible on subsequent peeks and pops.
+    fn peek_next_avail(&self, rsize: u16, mem: &MemCtx) -> Option<VqReq> {
         if !self.valid {
             return None;
         }
-        if let Some(idx) = mem.read::<u16>(self.gpa_idx) {
-            let ndesc = Wrapping(*idx) - self.cur_avail_idx;
-            if ndesc.0 != 0 && ndesc.0 <= rsize {
-                let avail_idx = self.cur_avail_idx.0 & (rsize - 1);
-                self.cur_avail_idx += Wrapping(1);
-
-                fence(Ordering::Acquire);
-                let addr = self.gpa_ring.offset::<u16>(avail_idx as usize);
-                return mem
-                    .read(addr)
-                    .map(|desc_idx| VqReq { desc_idx: *desc_idx, avail_idx });
-            }
+        let idx = mem.read::<u16>(self.gpa_idx)?;
+        let ndesc = Wrapping(*idx) - self.cur_avail_idx;
+        if ndesc.0 == 0 || ndesc.0 > rsize {
+            return None;
         }
-        None
+        let avail_idx = self.cur_avail_idx.0 & (rsize - 1);
+
+        fence(Ordering::Acquire);
+        let addr = self.gpa_ring.offset::<u16>(avail_idx as usize);
+        mem.read(addr).map(|desc_idx| VqReq { desc_idx: *desc_idx, avail_idx })
+    }
+
+    /// If there's a request ready, pop it off the queue and return the
+    /// corresponding descriptor and available ring indicies.
+    fn read_next_avail(&mut self, rsize: u16, mem: &MemCtx) -> Option<VqReq> {
+        let req = self.peek_next_avail(rsize, mem)?;
+        self.cur_avail_idx += Wrapping(1);
+        Some(req)
     }
 
     fn read_ring_descr(
@@ -494,17 +498,53 @@ impl VirtQueue {
         assert!(chain.idx.is_none());
         let mut avail = self.avail.lock().unwrap();
         let req = avail.read_next_avail(self.size(), mem)?;
+        probes::virtio_vq_pop!(|| (
+            self as *const VirtQueue as u64,
+            req.desc_idx,
+            req.avail_idx,
+        ));
+        self.walk_chain(&avail, req, chain, mem)
+    }
 
+    /// Accumulates the descriptor chain at the front of the avail ring into
+    /// a `Chain`, returned as an [`AvailChain`], without consuming the queue
+    /// entry.
+    ///
+    /// The queue cursor is untouched: dropping the returned [`AvailChain`]
+    /// is a no-op and the same chain is returned by the next peek or pop.
+    /// Consuming it via [`AvailChain::push_used`] advances the cursor and
+    /// pushes the chain onto the used ring in one step.
+    pub fn peek_avail(&self, mem: &MemCtx) -> Option<AvailChain<'_>> {
+        let mut avail = self.avail.lock().unwrap();
+        let req = avail.peek_next_avail(self.size(), mem)?;
+        let pos = avail.cur_avail_idx;
+        let mut chain = Chain::with_capacity(10);
+        if self.walk_chain(&avail, req, &mut chain, mem).is_none() {
+            // TODO (#1190)
+            // If the chain is malformed we bump the index that would have
+            // happened in `avail.read_next_avail` so that we don't get stuck
+            // peeking the same malformed chain over and over. What we really
+            // need to do, is signal the error condition via DEVICE_NEEDS_RESET.
+            // Until then, we will leak these descriptors.
+            avail.cur_avail_idx += Wrapping(1);
+            return None;
+        }
+        Some(AvailChain { vq: self, req, cur_avail_idx_snap: pos, chain })
+    }
+
+    /// Given the [`VqReq`]'s `desc_idx`, fill the provided [`Chain`].
+    fn walk_chain(
+        &self,
+        avail: &VqAvail,
+        req: VqReq,
+        chain: &mut Chain,
+        mem: &MemCtx,
+    ) -> Option<(u16, u32)> {
         let mut desc = avail.read_ring_descr(req.desc_idx, self.size(), mem)?;
         let mut flags = DescFlag::from_bits_truncate(desc.flags);
         let mut count = 0;
         let mut len = 0;
         chain.idx = Some(req.desc_idx);
-        probes::virtio_vq_pop!(|| (
-            self as *const VirtQueue as u64,
-            req.desc_idx,
-            req.avail_idx,
-        ));
 
         // non-indirect descriptor(s)
         while !flags.contains(DescFlag::INDIRECT) {
@@ -873,6 +913,59 @@ impl VirtQueue {
     }
 }
 
+/// The current descriptor chain read from the front of the avail ring, but not
+/// yet consumed.
+///
+/// Dropping the [`AvailChain`] is a no-op and leaves the data where it is,
+/// while calling [`AvailChain::push_used`] will advance the cursor and push the
+/// chain onto the used ring.
+///
+/// NB: This is here mostly for our virtio-socket implementation so that we do
+/// hold onto cached descriptors for the poller, which ensures there is only a
+/// single consumer of the guest's virt queues.
+pub struct AvailChain<'a> {
+    vq: &'a VirtQueue,
+    req: VqReq,
+    // a snapshot of `cur_avail_idx` when calling [`VirtQueue::peek_avail`],
+    // used to validate the VirtQueue only has a single consumer.
+    cur_avail_idx_snap: Wrapping<u16>,
+    chain: Chain,
+}
+
+impl AvailChain<'_> {
+    pub fn chain(&self) -> &Chain {
+        &self.chain
+    }
+
+    pub fn chain_mut(&mut self) -> &mut Chain {
+        &mut self.chain
+    }
+
+    /// Consume the chain, and advance the avail cursor. This pushes it onto
+    /// the used ring.
+    pub fn push_used(mut self, mem: &MemCtx) {
+        {
+            let mut avail = self.vq.avail.lock().unwrap();
+            // Our implementation of virtio queues allow for multiple consumers,
+            // which means racing consumers can step on each other. This should
+            // only be called when a single consumer is polling a virt queue,
+            // such as with virtio-socket. We assert here to indicate programmer
+            // error if we ever find things out of sync.
+            assert_eq!(
+                avail.cur_avail_idx, self.cur_avail_idx_snap,
+                "avail cursor moved since peek"
+            );
+            avail.cur_avail_idx += Wrapping(1);
+        }
+        probes::virtio_vq_pop!(|| (
+            self.vq as *const VirtQueue as u64,
+            self.req.desc_idx,
+            self.req.avail_idx,
+        ));
+        self.vq.push_used(&mut self.chain, mem);
+    }
+}
+
 bitflags! {
     #[derive(Default)]
     pub struct DescFlag: u16 {

lib/propolis/src/hw/virtio/vsock.rs

@@ -6,15 +6,18 @@ use lazy_static::lazy_static;
 use slog::Logger;
 use std::sync::Arc;
 
+use crate::accessors::Guard;
 use crate::accessors::MemAccessor;
 use crate::common::*;
 use crate::hw::pci;
 use crate::hw::virtio;
+use crate::hw::virtio::queue::AvailChain;
 use crate::hw::virtio::queue::Chain;
 use crate::hw::virtio::queue::VirtQueue;
 use crate::hw::virtio::queue::VqSize;
 use crate::migrate::*;
 use crate::util::regmap::RegMap;
+use crate::vmm::mem::MemAccessed;
 use crate::vmm::MemCtx;
 use crate::vsock::packet::VsockPacket;
 use crate::vsock::packet::VsockPacketError;
@@ -40,36 +43,25 @@ pub const VSOCK_EVENT_QUEUE: u16 = 0x2;
 /// This guarantees we have space to send a packet to the guest before reading
 /// data from a host socket, preventing data loss if the queue is full.
 ///
-/// The permit holds a mutable reference to `VsockVq`, ensuring only one permit
-/// can exist at a time (enforced at compile time). If dropped without calling
-/// `write`, the chain is retained in `VsockVq` for reuse.
+/// NB: The underlying descriptor chain is not consumed until a call to
+/// [`RxPermit::write`] is made. Droopping the [`RxPermit`] is a no-op without a
+/// call to write is a no-op.
 pub struct RxPermit<'a> {
-    vq: &'a mut VsockVq,
+    avail: AvailChain<'a>,
+    mem: Guard<'a, MemAccessed>,
 }
 
 impl RxPermit<'_> {
     /// Returns the maximum data payload that can fit in this descriptor chain.
     pub fn available_data_space(&self) -> usize {
         let header_size = std::mem::size_of::<VsockPacketHeader>();
-        self.vq
-            .rx_chain
-            .as_ref()
-            .expect("has chain")
-            .remain_write_bytes()
-            .saturating_sub(header_size)
+        self.avail.chain().remain_write_bytes().saturating_sub(header_size)
     }
 
-    pub fn write(self, header: &VsockPacketHeader, data: &[u8]) {
-        // TODO: cannot access memory?
-        let mem = self.vq.acc_mem.access().expect("mem access for write");
-        let queue =
-            self.vq.queues.get(VSOCK_RX_QUEUE as usize).expect("rx queue");
-
-        // NB: `RxPermit` should only be created if the owning `VsockVq`
-        // actually has a `Some(Chain)`. Unfortuantely there doesn't seem to be
-        // a way to enforce this at compile time.
-        let mut chain = self.vq.rx_chain.take().expect("has chain");
-        chain.write(header, &mem);
+    pub fn write(mut self, header: &VsockPacketHeader, data: &[u8]) {
+        let mem = &*self.mem;
+        let chain = self.avail.chain_mut();
+        chain.write(header, mem);
 
         if !data.is_empty() {
             let mut done = 0;
@@ -86,29 +78,27 @@ impl RxPermit<'_> {
         }
 
         probes::vsock_pkt_rx!(|| header);
-        queue.push_used(&mut chain, &mem);
+        self.avail.push_used(mem);
     }
 }
 
 pub struct VsockVq {
     queues: Vec<Arc<VirtQueue>>,
     acc_mem: MemAccessor,
-    /// Cached rx chain for permit reuse when dropped without write
-    rx_chain: Option<Chain>,
 }
 
 impl VsockVq {
     pub(crate) fn new(
         queues: Vec<Arc<VirtQueue>>,
         acc_mem: MemAccessor,
     ) -> Self {
-        Self { queues, acc_mem, rx_chain: None }
+        Self { queues, acc_mem }
     }
 
     /// Try to acquire a permit for sending a packet to the guest.
     ///
-    /// Returns `Some(RxPermit)` if a descriptor chain is available,
-    /// `None` if the rx queue is full.
+    /// Returns `Some(RxPermit)` if a descriptor chain is available, `None`
+    /// if the rx queue is full.
     pub fn try_rx_permit(&mut self) -> Option<RxPermit<'_>> {
         let vq = self.queues.get(VSOCK_RX_QUEUE as usize)?;
         // See propolis#1110 & propolis#1115
@@ -121,22 +111,9 @@ impl VsockVq {
             return None;
         }
 
-        // Reuse cached chain or pop a new one
-        if self.rx_chain.is_none() {
-            // TODO: cannot access memory?
-            let mem = self.acc_mem.access().expect("mem access for write");
-            let mut chain = Chain::with_capacity(10);
-            if let Some(_) = vq.pop_avail(&mut chain, &mem) {
-                self.rx_chain = Some(chain);
-            }
-        }
-
-        // We only return a permit iff we know that we are holding onto a valid
-        // descriptor chain that can be used by the borrowing `RxPermit`
-        match self.rx_chain {
-            Some(_) => Some(RxPermit { vq: self }),
-            None => None,
-        }
+        let mem = self.acc_mem.access()?;
+        let avail = vq.peek_avail(&mem)?;
+        Some(RxPermit { avail, mem })
     }
 
     /// Receive all available packets from the TX queue.
@@ -161,15 +138,6 @@ impl VsockVq {
 
         Some(packet)
     }
-
-    /// Drop any cached descriptor chain.
-    ///
-    /// This MUST be called when reseting the virtio-socket device so
-    /// that we don't use stale `GuestAddr`s across device resets.
-    #[cfg(target_os = "illumos")]
-    pub(crate) fn clear_rx_chain(&mut self) {
-        self.rx_chain = None;
-    }
 }
 
 pub struct PciVirtioSock {

lib/propolis/src/vsock/poller.rs

@@ -984,7 +984,6 @@ impl VsockPoller {
         self.rx.clear();
         self.rx_blocked.clear();
         self.quiescing.clear();
-        self.queues.clear_rx_chain();
     }
 }
 
@@ -2187,9 +2186,6 @@ mod test {
         notify.wait_stopped();
         notify.reset_evloop();
 
-        // NOTE: We don't have a way to actually validate that the reset call
-        // above has not left behind a `GuestAddr` as apart of the `VsockVq`.
-
         // The host-side TCP socket should now be closed because the
         // poller dropped all connections during reset.
         let mut buf = [0u8; 1];