virtio: build descriptor chains without a per-packet allocation

Cargo.lock

@@ -9642,6 +9642,7 @@ dependencies = [
  "parking_lot",
  "pci_core",
  "pci_resources",
+ "smallvec",
  "task_control",
  "test_with_tracing",
  "thiserror 2.0.16",

vm/devices/virtio/virtio/Cargo.toml

@@ -28,6 +28,7 @@ anyhow.workspace = true
 async-trait.workspace = true
 futures.workspace = true
 parking_lot.workspace = true
+smallvec.workspace = true
 thiserror.workspace = true
 zerocopy.workspace = true
 [dev-dependencies]

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

@@ -8,6 +8,7 @@ use crate::queue::QueueError;
 use crate::queue::QueueParams;
 use crate::queue::QueueState;
 use crate::queue::VirtioQueuePayload;
+use crate::queue::VirtioQueuePayloadVec;
 use crate::queue::new_queue;
 use crate::spec::VirtioDeviceFeatures;
 use crate::spec::VirtioDeviceType;
@@ -105,14 +106,11 @@ fn read_from_payload_at_offset(
 #[must_use]
 pub struct VirtioQueueCallbackWork {
     completion: QueueCompletion,
-    pub payload: Vec<VirtioQueuePayload>,
+    pub payload: VirtioQueuePayloadVec,
 }
 
 impl VirtioQueueCallbackWork {
-    pub(crate) fn from_parts(
-        completion: QueueCompletion,
-        payload: Vec<VirtioQueuePayload>,
-    ) -> Self {
+    pub(crate) fn from_parts(completion: QueueCompletion, payload: VirtioQueuePayloadVec) -> Self {
         Self {
             completion,
             payload,

vm/devices/virtio/virtio/src/queue.rs

@@ -15,6 +15,7 @@ use inspect::Inspect;
 use packed::PackedQueueCompleteWork;
 pub use packed::PackedQueueCompletionContext;
 use packed::PackedQueueGetWork;
+use smallvec::SmallVec;
 use spec::DescriptorFlags;
 use spec::PackedDescriptor;
 use spec::SplitDescriptor;
@@ -405,11 +406,10 @@ impl QueueCoreGetWork {
     }
 
     fn work_from_index(&mut self, index: u16) -> Result<VirtioQueueCallbackWork, QueueError> {
+        let mut payload = VirtioQueuePayloadVec::new();
         if let QueueGetWorkInner::Split(split) = &mut self.inner {
             let descriptor_index = split.get_available_descriptor_index(index)?;
-            let payload = self
-                .reader(descriptor_index)
-                .collect::<Result<Vec<_>, _>>()?;
+            self.read_chain(descriptor_index, &mut payload)?;
             Ok(VirtioQueueCallbackWork::from_parts(
                 QueueCompletion {
                     descriptor_index,
@@ -418,21 +418,19 @@ impl QueueCoreGetWork {
                 payload,
             ))
         } else {
-            let (payload, last_primary_desc_index) = {
-                let mut reader = self.reader(index);
-                (
-                    (&mut reader).collect::<Result<Vec<_>, _>>()?,
-                    reader.last_primary_desc_index(),
-                )
-            };
-            let last = self.descriptor(&self.queue_desc, last_primary_desc_index, None)?;
-            let count = if last_primary_desc_index >= index {
-                last_primary_desc_index - index + 1
+            let tail = self.read_chain(index, &mut payload)?;
+            let count = if tail.last_primary_desc_index >= index {
+                tail.last_primary_desc_index - index + 1
             } else {
                 // Wrapped around the end of the queue.
-                self.queue_size - index + last_primary_desc_index + 1
+                self.queue_size - index + tail.last_primary_desc_index + 1
             };
-            let completion_context = PackedQueueCompletionContext::new(&last, count);
+            // The packed path always reads at least one primary descriptor, and
+            // `descriptor` always sets a buffer id there.
+            let buffer_id = tail
+                .last_primary_buffer_id
+                .expect("packed descriptors have buffer id");
+            let completion_context = PackedQueueCompletionContext::new(buffer_id, count);
             Ok(VirtioQueueCallbackWork::from_parts(
                 QueueCompletion {
                     context: QueueCompletionContext::Packed(completion_context),
@@ -443,10 +441,23 @@ impl QueueCoreGetWork {
         }
     }
 
-    fn reader(&mut self, descriptor_index: u16) -> DescriptorReader<'_> {
-        DescriptorReader {
-            chain: DescriptorChain::new(self, self.features.ring_indirect_desc(), descriptor_index),
+    /// Walks the descriptor chain at `descriptor_index`, appending each buffer
+    /// to `payload`, and reports what the completion needs from its tail.
+    fn read_chain(
+        &self,
+        descriptor_index: u16,
+        payload: &mut VirtioQueuePayloadVec,
+    ) -> Result<ChainTail, QueueError> {
+        let mut chain =
+            DescriptorChain::new(self, self.features.ring_indirect_desc(), descriptor_index);
+        while let Some(descriptor) = chain.next_descriptor()? {
+            payload.push(VirtioQueuePayload {
+                writeable: descriptor.flags.write(),
+                address: descriptor.address,
+                length: descriptor.length,
+            });
         }
+        Ok(chain.tail())
     }
 
     fn descriptor(
@@ -585,34 +596,33 @@ pub(crate) fn new_queue(
     Ok((get_work, complete_work))
 }
 
-struct DescriptorReader<'a> {
-    chain: DescriptorChain<'a>,
-}
-
-impl DescriptorReader<'_> {
-    pub fn last_primary_desc_index(&self) -> u16 {
-        self.chain.last_primary_desc_index()
-    }
-}
-
 pub struct VirtioQueuePayload {
     pub writeable: bool,
     pub address: u64,
     pub length: u32,
 }
 
-impl Iterator for DescriptorReader<'_> {
-    type Item = Result<VirtioQueuePayload, QueueError>;
+/// Number of payload buffers a descriptor chain can hold without allocating.
+///
+/// Sized for the short chains, such as a header, one data buffer and a status
+/// byte. Longer chains still work and spill to the heap.
+const PAYLOAD_INLINE_CAPACITY: usize = 4;
 
-    fn next(&mut self) -> Option<Self::Item> {
-        self.chain.next().map(|descriptor| {
-            descriptor.map(|descriptor| VirtioQueuePayload {
-                writeable: descriptor.flags.write(),
-                address: descriptor.address,
-                length: descriptor.length,
-            })
-        })
-    }
+/// The payload buffers of one descriptor chain.
+pub type VirtioQueuePayloadVec = SmallVec<[VirtioQueuePayload; PAYLOAD_INLINE_CAPACITY]>;
+
+/// What a completion needs from the end of a descriptor chain.
+///
+/// Packed rings post against the buffer id of the last descriptor in the primary
+/// ring, so the walk records it instead of re-reading that descriptor. The
+/// captured id identifies the buffer the guest offered even if the guest rewrites
+/// the ring afterwards.
+struct ChainTail {
+    /// Index of the last descriptor read from the primary ring.
+    last_primary_desc_index: u16,
+    /// Buffer id of that descriptor. Always set for packed rings; split rings
+    /// have no buffer ids and ignore this.
+    last_primary_buffer_id: Option<u16>,
 }
 
 struct DescriptorChain<'a> {
@@ -625,6 +635,8 @@ struct DescriptorChain<'a> {
     indirect_table_len: Option<u16>,
     descriptor_index: Option<u16>,
     last_primary_desc_index: u16,
+    /// Buffer id of the descriptor at `last_primary_desc_index`.
+    last_primary_buffer_id: Option<u16>,
     num_read: u16,
 }
 
@@ -638,10 +650,19 @@ impl<'a> DescriptorChain<'a> {
             indirect_table_len: None,
             descriptor_index: Some(descriptor_index),
             last_primary_desc_index: descriptor_index,
+            last_primary_buffer_id: None,
             num_read: 0,
         }
     }
 
+    /// What the completion needs from the end of this chain, once the walk is done.
+    fn tail(&self) -> ChainTail {
+        ChainTail {
+            last_primary_desc_index: self.last_primary_desc_index,
+            last_primary_buffer_id: self.last_primary_buffer_id,
+        }
+    }
+
     fn next_descriptor(&mut self) -> Result<Option<QueueDescriptor>, QueueError> {
         let Some(descriptor_index) = self.descriptor_index else {
             return Ok(None);
@@ -653,6 +674,12 @@ impl<'a> DescriptorChain<'a> {
             descriptor_index,
             self.indirect_table_len,
         )?;
+        if self.indirect_queue.is_none() {
+            // Read from the primary ring, so this is the id a packed completion
+            // posts against. An indirect chain's only primary descriptor is its
+            // head, which is what `last_primary_desc_index` keeps.
+            self.last_primary_buffer_id = descriptor.buffer_id;
+        }
         let descriptor = if !self.indirect_support || !descriptor.flags.indirect() {
             if self.indirect_queue.is_none() {
                 self.last_primary_desc_index = descriptor_index;
@@ -685,18 +712,6 @@ impl<'a> DescriptorChain<'a> {
         }
         Ok(Some(descriptor))
     }
-
-    pub fn last_primary_desc_index(&self) -> u16 {
-        self.last_primary_desc_index
-    }
-}
-
-impl Iterator for DescriptorChain<'_> {
-    type Item = Result<QueueDescriptor, QueueError>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        self.next_descriptor().transpose()
-    }
 }
 
 #[cfg(test)]

vm/devices/virtio/virtio/src/queue/packed.rs

@@ -3,7 +3,6 @@
 
 //! Virtio packed queue implementation.
 
-use crate::queue::QueueDescriptor;
 use crate::queue::QueueError;
 use crate::queue::QueueParams;
 use crate::queue::descriptor_offset;
@@ -23,11 +22,9 @@ pub struct PackedQueueCompletionContext {
 }
 
 impl PackedQueueCompletionContext {
-    pub(super) fn new(last_descriptor: &QueueDescriptor, descriptor_count: u16) -> Self {
+    pub(super) fn new(buffer_id: u16, descriptor_count: u16) -> Self {
         Self {
-            buffer_id: last_descriptor
-                .buffer_id
-                .expect("packed descriptors have buffer id"),
+            buffer_id,
             descriptor_count,
         }
     }

vm/devices/virtio/virtio/src/tests.rs

@@ -1000,9 +1000,15 @@ impl VirtioTestGuest {
                 0
             };
             if matches!(self.info, VirtioQueueInfo::Packed(_)) {
-                // buffer id (ignored for indirect descriptors)
+                // Buffer id. A packed completion is posted against the id in the
+                // PRIMARY ring descriptor; the ids inside an indirect table are
+                // reserved and must be ignored (spec 2.8.7). Stamp a sentinel
+                // rather than 0 so that reading one by mistake cannot coincide
+                // with the head's id, which is what the completion assertions
+                // check. With 0 here the two were indistinguishable and the
+                // distinction went untested.
                 self.test_mem
-                    .modify_memory_map(base + 12, &0_u16.to_le_bytes(), false);
+                    .modify_memory_map(base + 12, &(0xD1E0 + i).to_le_bytes(), false);
                 // flags
                 self.test_mem
                     .modify_memory_map(base + 14, &flags.to_le_bytes(), false);
@@ -2706,6 +2712,70 @@ async fn verify_packed_queue_linked(driver: DefaultDriver) {
     let test_mem = VirtioTestMemoryAccess::new();
     verify_queue_linked_inner(VirtioTestGuest::new_packed(&driver, &test_mem, 1, 8, true)).await;
 }
+
+/// A chain longer than the payload's inline capacity, so it spills to the heap.
+///
+/// Every other chain test here stays at or under the inline bound, so this is the
+/// only one that would catch a spill dropping, duplicating or reordering buffers.
+async fn verify_queue_chain_longer_than_inline_capacity_inner(mut guest: VirtioTestGuest) {
+    // Comfortably past the inline bound while still fitting the size-8 ring.
+    const CHAIN_LEN: usize = 6;
+
+    let (tx, mut rx) = mesh::mpsc_channel();
+    let base_address = guest.get_queue_descriptor_backing_memory_address(0);
+    let event = Event::new();
+    let mut queues = guest.create_direct_queues(|i| {
+        let tx = tx.clone();
+        CreateDirectQueueParams {
+            process_work: Box::new(
+                move |queue: &mut VirtioQueue, work: VirtioQueueCallbackWork| {
+                    assert_eq!(
+                        work.payload.len(),
+                        CHAIN_LEN,
+                        "spilled chain lost or gained buffers"
+                    );
+                    // Order matters as much as count: the buffers are one packet.
+                    for (i, payload) in work.payload.iter().enumerate() {
+                        assert_eq!(payload.address, base_address + 0x1000 * i as u64);
+                        assert_eq!(payload.length, 0x1000);
+                    }
+                    queue.complete(work, 789);
+                },
+            ),
+            notify: Interrupt::from_fn(move || {
+                tx.send(i as usize);
+            }),
+            event: event.clone(),
+        }
+    });
+
+    guest.add_linked_to_avail_queue(0, CHAIN_LEN as u16);
+    event.signal();
+    must_recv_in_timeout(&mut rx, Duration::from_millis(100)).await;
+    let (desc, len) = guest.get_next_completed(0).unwrap();
+    assert_eq!(desc, 0u16);
+    assert_eq!(len, 789);
+    assert_eq!(guest.get_next_completed(0).is_none(), true);
+
+    queues[0].stop().await;
+}
+
+#[async_test]
+async fn verify_split_queue_chain_longer_than_inline_capacity(driver: DefaultDriver) {
+    let test_mem = VirtioTestMemoryAccess::new();
+    verify_queue_chain_longer_than_inline_capacity_inner(VirtioTestGuest::new_split(
+        &driver, &test_mem, 1, 8, true,
+    ))
+    .await;
+}
+#[async_test]
+async fn verify_packed_queue_chain_longer_than_inline_capacity(driver: DefaultDriver) {
+    let test_mem = VirtioTestMemoryAccess::new();
+    verify_queue_chain_longer_than_inline_capacity_inner(VirtioTestGuest::new_packed(
+        &driver, &test_mem, 1, 8, true,
+    ))
+    .await;
+}
 /// A packed linked chain that starts near the end of the ring and wraps to
 /// the beginning (indices 2 → 3 → 0 in a queue_size=4 ring). Without
 /// correct index wrapping in `descriptor()`, the NEXT index after 3 would