feat(iovec): add ability to reload IoVecBufferMut

Now IoVecBufferMut can be reloaded from DescriptorChain
same way as IoVecBuffer does it.
This is helpful to avoid unnecessary allocations/deallocations
when reusing same buffer.

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

Author: ShadowCurse

src/vmm/src/devices/virtio/iovec.rs

@@ -217,7 +217,7 @@ impl IoVecBuffer {
 /// It describes a write-only buffer passed to us by the guest that is scattered across multiple
 /// memory regions. Additionally, this wrapper provides methods that allow reading arbitrary ranges
 /// of data from that buffer.
-#[derive(Debug)]
+#[derive(Debug, Default, Clone)]
 pub struct IoVecBufferMut {
     // container of the memory regions included in this IO vector
     vecs: IoVecVec,
@@ -226,12 +226,19 @@ pub struct IoVecBufferMut {
 }
 
 impl IoVecBufferMut {
-    /// Create an `IoVecBufferMut` from a `DescriptorChain`
-    pub fn from_descriptor_chain(head: DescriptorChain) -> Result<Self, IoVecError> {
-        let mut vecs = IoVecVec::new();
-        let mut len = 0u32;
+    /// Create an `IoVecBuffer` from a `DescriptorChain`
+    ///
+    /// # Safety
+    ///
+    /// The descriptor chain cannot be referencing the same memory location as another chain
+    pub unsafe fn load_descriptor_chain(
+        &mut self,
+        head: DescriptorChain,
+    ) -> Result<(), IoVecError> {
+        self.clear();
 
-        for desc in head {
+        let mut next_descriptor = Some(head);
+        while let Some(desc) = next_descriptor {
             if !desc.is_write_only() {
                 return Err(IoVecError::ReadOnlyDescriptor);
             }
@@ -247,23 +254,45 @@ impl IoVecBufferMut {
             slice.bitmap().mark_dirty(0, desc.len as usize);
 
             let iov_base = slice.ptr_guard_mut().as_ptr().cast::<c_void>();
-            vecs.push(iovec {
+            self.vecs.push(iovec {
                 iov_base,
                 iov_len: desc.len as size_t,
             });
-            len = len
+            self.len = self
+                .len
                 .checked_add(desc.len)
                 .ok_or(IoVecError::OverflowedDescriptor)?;
+
+            next_descriptor = desc.next_descriptor();
         }
 
-        Ok(Self { vecs, len })
+        Ok(())
+    }
+
+    /// Create an `IoVecBuffer` from a `DescriptorChain`
+    ///
+    /// # Safety
+    ///
+    /// The descriptor chain cannot be referencing the same memory location as another chain
+    pub unsafe fn from_descriptor_chain(head: DescriptorChain) -> Result<Self, IoVecError> {
+        let mut new_buffer: Self = Default::default();
+
+        new_buffer.load_descriptor_chain(head)?;
+
+        Ok(new_buffer)
     }
 
     /// Get the total length of the memory regions covered by this `IoVecBuffer`
     pub(crate) fn len(&self) -> u32 {
         self.len
     }
 
+    /// Clears the `iovec` array
+    pub fn clear(&mut self) {
+        self.vecs.clear();
+        self.len = 0u32;
+    }
+
     /// Writes a number of bytes into the `IoVecBufferMut` starting at a given offset.
     ///
     /// This will try to fill `IoVecBufferMut` writing bytes from the `buf` starting from
@@ -468,11 +497,13 @@ mod tests {
 
         let (mut q, _) = read_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
-        IoVecBufferMut::from_descriptor_chain(head).unwrap_err();
+        // SAFETY: This descriptor chain is only loaded into one buffer
+        unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap_err() };
 
         let (mut q, _) = write_only_chain(&mem);
         let head = q.pop(&mem).unwrap();
-        IoVecBufferMut::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer
+        unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap() };
     }
 
     #[test]
@@ -493,7 +524,7 @@ mod tests {
         let head = q.pop(&mem).unwrap();
 
         // SAFETY: This descriptor chain is only loaded once in this test
-        let iovec = IoVecBufferMut::from_descriptor_chain(head).unwrap();
+        let iovec = unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap() };
         assert_eq!(iovec.len(), 4 * 64);
     }
 
@@ -558,7 +589,8 @@ mod tests {
         // This is a descriptor chain with 4 elements 64 bytes long each.
         let head = q.pop(&mem).unwrap();
 
-        let mut iovec = IoVecBufferMut::from_descriptor_chain(head).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer
+        let mut iovec = unsafe { IoVecBufferMut::from_descriptor_chain(head).unwrap() };
         let buf = vec![0u8, 1, 2, 3, 4];
 
         // One test vector for each part of the chain

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

@@ -132,7 +132,10 @@ impl Entropy {
             let index = desc.index;
             METRICS.entropy_event_count.inc();
 
-            let bytes = match IoVecBufferMut::from_descriptor_chain(desc) {
+            // SAFETY: This descriptor chain is only loaded once
+            // virtio requests are handled sequentially so no two IoVecBuffers
+            // are live at the same time, meaning this has exclusive ownership over the memory
+            let bytes = match unsafe { IoVecBufferMut::from_descriptor_chain(desc) } {
                 Ok(mut iovec) => {
                     debug!(
                         "entropy: guest request for {} bytes of entropy",
@@ -428,13 +431,15 @@ mod tests {
         // This should succeed, we just added two descriptors
         let desc = entropy_dev.queues_mut()[RNG_QUEUE].pop(&mem).unwrap();
         assert!(matches!(
-            IoVecBufferMut::from_descriptor_chain(desc),
+            // SAFETY: This descriptor chain is only loaded into one buffer
+            unsafe { IoVecBufferMut::from_descriptor_chain(desc) },
             Err(crate::devices::virtio::iovec::IoVecError::ReadOnlyDescriptor)
         ));
 
         // This should succeed, we should have one more descriptor
         let desc = entropy_dev.queues_mut()[RNG_QUEUE].pop(&mem).unwrap();
-        let mut iovec = IoVecBufferMut::from_descriptor_chain(desc).unwrap();
+        // SAFETY: This descriptor chain is only loaded into one buffer
+        let mut iovec = unsafe { IoVecBufferMut::from_descriptor_chain(desc).unwrap() };
         entropy_dev.handle_one(&mut iovec).unwrap();
     }
 

src/vmm/src/devices/virtio/vsock/packet.rs

@@ -161,7 +161,10 @@ impl VsockPacket {
     /// Returns [`VsockError::DescChainTooShortForHeader`] if the descriptor chain's total buffer
     /// length is insufficient to hold the 44 byte vsock header
     pub fn from_rx_virtq_head(chain: DescriptorChain) -> Result<Self, VsockError> {
-        let buffer = IoVecBufferMut::from_descriptor_chain(chain)?;
+        // SAFETY: This descriptor chain is only loaded once
+        // virtio requests are handled sequentially so no two IoVecBuffers
+        // are live at the same time, meaning this has exclusive ownership over the memory
+        let buffer = unsafe { IoVecBufferMut::from_descriptor_chain(chain)? };
 
         if buffer.len() < VSOCK_PKT_HDR_SIZE {
             return Err(VsockError::DescChainTooShortForHeader(buffer.len() as usize));