virtio: remove per-work Arc<Mutex> clone from queue completion path (#3173)

Every `VirtioQueueCallbackWork` currently clones an
`Arc<Mutex<VirtioQueueUsedHandler>>` so that it can independently call
`complete()` on itself. This means every descriptor processed by every
virtio device pays for an atomic ref-count increment/decrement pair plus
a mutex lock on the completion path—overhead that adds up on
high-throughput devices like virtio-net and virtio-blk.

This PR moves completion responsibility from the work item to the queue:
callers now invoke `VirtioQueue::complete(work, bytes_written)` instead
of `work.complete(bytes_written)`. Because the queue already owns the
used-ring handler, the `Arc<Mutex>` clone per work item, the `completed`
tracking flag, and the `Drop` guard on `VirtioQueue` are all eliminated.

All in-tree virtio device crates are updated to the new pattern. For
virtio-vsock, the refactor also fixes a latent bug where `write_packet`
could early-return on error without completing the work item; the caller
now always completes it.

Author: jstarks

vm/devices/virtio/virtio/src/common.rs

@@ -19,7 +19,6 @@ use guestmem::GuestMemoryError;
 use inspect::Inspect;
 use pal_async::wait::PolledWait;
 use pal_event::Event;
-use parking_lot::Mutex;
 use std::io::Error;
 use std::pin::Pin;
 use std::sync::Arc;
@@ -98,67 +97,21 @@ fn read_from_payload_at_offset(
     Ok(read_bytes)
 }
 
-#[derive(Debug)]
-pub(crate) struct VirtioQueueUsedHandler {
-    core: QueueCoreCompleteWork,
-    notify_guest: Interrupt,
-}
-
-impl VirtioQueueUsedHandler {
-    pub(crate) fn new(core: QueueCoreCompleteWork, notify_guest: Interrupt) -> Self {
-        Self { core, notify_guest }
-    }
-
-    pub(crate) fn complete_descriptor(&mut self, work: &QueueWork, bytes_written: u32) {
-        match self.core.complete_descriptor(work, bytes_written) {
-            Ok(true) => {
-                self.notify_guest.deliver();
-            }
-            Ok(false) => {}
-            Err(err) => {
-                tracelimit::error_ratelimited!(
-                    error = &err as &dyn std::error::Error,
-                    "failed to complete descriptor"
-                );
-            }
-        }
-    }
-}
-
 /// A descriptor chain popped from a [`VirtioQueue`].
 ///
-/// The device must call [`complete`](Self::complete) exactly once to post a
+/// The device must call [`VirtioQueue::complete`] exactly once to post a
 /// completion to the guest's used ring. Dropping without completing is a bug
 /// and will not automatically post a completion.
 #[must_use]
 pub struct VirtioQueueCallbackWork {
-    used_queue_handler: Arc<Mutex<VirtioQueueUsedHandler>>,
     work: QueueWork,
     pub payload: Vec<VirtioQueuePayload>,
-    completed: bool,
 }
 
 impl VirtioQueueCallbackWork {
-    pub(crate) fn new(
-        mut work: QueueWork,
-        used_queue_handler: &Arc<Mutex<VirtioQueueUsedHandler>>,
-    ) -> Self {
-        let used_queue_handler = used_queue_handler.clone();
+    pub(crate) fn new(mut work: QueueWork) -> Self {
         let payload = std::mem::take(&mut work.payload);
-        Self {
-            work,
-            payload,
-            used_queue_handler,
-            completed: false,
-        }
-    }
-
-    pub fn complete(&mut self, bytes_written: u32) {
-        assert!(!self.completed);
-        self.used_queue_handler
-            .lock()
-            .complete_descriptor(&self.work, bytes_written);
-        self.completed = true;
+        Self { work, payload }
     }
 
     pub fn descriptor_index(&self) -> u16 {
@@ -290,10 +243,10 @@ impl<'a> PeekedWork<'a> {
     /// Consume this peeked work, advancing the queue's available index.
     ///
     /// Returns a [`VirtioQueueCallbackWork`] that must be explicitly
-    /// completed via [`VirtioQueueCallbackWork::complete`].
+    /// completed via [`VirtioQueue::complete`].
     pub fn consume(self) -> VirtioQueueCallbackWork {
         self.queue.core.advance(&self.work);
-        VirtioQueueCallbackWork::new(self.work, &self.queue.used_handler)
+        VirtioQueueCallbackWork::new(self.work)
     }
 }
 
@@ -302,7 +255,9 @@ pub struct VirtioQueue {
     #[inspect(flatten)]
     core: QueueCoreGetWork,
     #[inspect(skip)]
-    used_handler: Arc<Mutex<VirtioQueueUsedHandler>>,
+    complete: QueueCoreCompleteWork,
+    #[inspect(skip)]
+    notify_guest: Interrupt,
     #[inspect(skip)]
     queue_event: PolledWait<Event>,
 }
@@ -317,13 +272,10 @@ impl VirtioQueue {
         initial_state: Option<QueueState>,
     ) -> Result<Self, QueueError> {
         let (get_work, complete_work) = new_queue(features, mem, params, initial_state)?;
-        let used_handler = Arc::new(Mutex::new(VirtioQueueUsedHandler::new(
-            complete_work,
-            notify,
-        )));
         Ok(Self {
             core: get_work,
-            used_handler,
+            complete: complete_work,
+            notify_guest: notify,
             queue_event,
         })
     }
@@ -332,7 +284,7 @@ impl VirtioQueue {
     pub fn queue_state(&self) -> QueueState {
         QueueState {
             avail_index: self.core.avail_index(),
-            used_index: self.used_handler.lock().core.used_index(),
+            used_index: self.complete.used_index(),
         }
     }
 
@@ -362,7 +314,7 @@ impl VirtioQueue {
             .core
             .try_next_work()
             .map_err(Error::other)?
-            .map(|work| VirtioQueueCallbackWork::new(work, &self.used_handler)))
+            .map(VirtioQueueCallbackWork::new))
     }
 
     /// Peek at the next available descriptor without advancing the available
@@ -402,6 +354,28 @@ impl VirtioQueue {
         Ok(PeekedWork::new(self, work))
     }
 
+    /// Complete a descriptor previously obtained from this queue.
+    ///
+    /// Writes `bytes_written` to the used ring and delivers an interrupt
+    /// to the guest (unless interrupt suppression is active).
+    ///
+    /// Takes ownership of the work item, ensuring it can only be completed
+    /// once.
+    pub fn complete(&mut self, work: VirtioQueueCallbackWork, bytes_written: u32) {
+        match self.complete.complete_descriptor(&work.work, bytes_written) {
+            Ok(true) => {
+                self.notify_guest.deliver();
+            }
+            Ok(false) => {}
+            Err(err) => {
+                tracelimit::error_ratelimited!(
+                    error = &err as &dyn std::error::Error,
+                    "failed to complete descriptor"
+                );
+            }
+        }
+    }
+
     fn poll_next_buffer(
         &mut self,
         cx: &mut Context<'_>,
@@ -415,14 +389,6 @@ impl VirtioQueue {
     }
 }
 
-impl Drop for VirtioQueue {
-    fn drop(&mut self) {
-        if Arc::get_mut(&mut self.used_handler).is_none() {
-            tracing::error!("Virtio queue dropped with outstanding work pending")
-        }
-    }
-}
-
 impl Stream for VirtioQueue {
     type Item = Result<VirtioQueueCallbackWork, Error>;