allow persistent bounce buffers in MaybeBounce

This is particularly useful for virtio devices, as we know that a virtio
descriptor chain is at most u16::MAX bytes long, meaning we can easily
just allocate a persistent bounce buffer of that size on device creation
and then never reallocate on the hotpath.

Signed-off-by: Patrick Roy <[email protected]>

Author: roypat

src/vmm/src/builder.rs

@@ -274,7 +274,7 @@ pub fn build_microvm_for_boot(
     }
 
     let entry_point = load_kernel(
-        MaybeBounce(&boot_config.kernel_file, secret_free),
+        MaybeBounce::new(&boot_config.kernel_file, secret_free),
         vmm.vm.guest_memory(),
     )?;
     let initrd = match &boot_config.initrd_file {
@@ -286,7 +286,7 @@ pub fn build_microvm_for_boot(
 
             Some(InitrdConfig::from_reader(
                 vmm.vm.guest_memory(),
-                MaybeBounce(initrd_file, secret_free),
+                MaybeBounce::new(initrd_file, secret_free),
                 u64_to_usize(size),
             )?)
         }

src/vmm/src/vstate/memory.rs

@@ -64,54 +64,123 @@ pub enum MemoryError {
 /// `Write` respectively, by reading/writing using a bounce buffer, and memcpy-ing into the
 /// [`VolatileSlice`].
 #[derive(Debug)]
-pub struct MaybeBounce<T>(pub T, pub bool);
+pub struct MaybeBounce<T, const N: usize = 0> {
+    pub(crate) target: T,
+    pub(crate) should_bounce: bool,
+    persistent_buffer: [u8; N],
+}
+
+impl<T> MaybeBounce<T, 0> {
+    /// Creates a new `MaybeBounce` that always allocates a bounce
+    /// buffer on-demand
+    pub fn new(target: T, should_bounce: bool) -> Self {
+        MaybeBounce::new_persistent(target, should_bounce)
+    }
+}
+
+impl<T, const N: usize> MaybeBounce<T, N> {
+    /// Creates a new `MaybeBounce` that uses a persistent, fixed size bounce buffer
+    /// of size `N`. If a read/write request exceeds the size of this bounce buffer, it
+    /// is split into multiple, `<= N`-size read/writes.
+    pub fn new_persistent(target: T, should_bounce: bool) -> Self {
+        MaybeBounce {
+            target,
+            should_bounce,
+            persistent_buffer: [0u8; N],
+        }
+    }
+}
 
 // FIXME: replace AsFd with ReadVolatile once &File: ReadVolatile in vm-memory.
-impl<T: Read + AsFd> ReadVolatile for MaybeBounce<T> {
+impl<T: Read + AsFd, const N: usize> ReadVolatile for MaybeBounce<T, N> {
     fn read_volatile<B: BitmapSlice>(
         &mut self,
         buf: &mut VolatileSlice<B>,
     ) -> Result<usize, VolatileMemoryError> {
-        if self.1 {
-            let mut bbuf = vec![0; buf.len()];
-            let n = self
-                .0
-                .read(bbuf.as_mut_slice())
-                .map_err(VolatileMemoryError::IOError)?;
-            buf.copy_from(&bbuf[..n]);
-            Ok(n)
+        if self.should_bounce {
+            let mut bbuf = (N == 0).then(|| vec![0u8; buf.len()]);
+            let bbuf = bbuf
+                .as_deref_mut()
+                .unwrap_or(self.persistent_buffer.as_mut_slice());
+
+            let mut buf = buf.offset(0)?;
+            let mut total = 0;
+            while !buf.is_empty() {
+                let how_much = buf.len().min(bbuf.len());
+                let n = self
+                    .target
+                    .read(&mut bbuf[..how_much])
+                    .map_err(VolatileMemoryError::IOError)?;
+                buf.copy_from(&bbuf[..n]);
+
+                buf = buf.offset(n)?;
+                total += n;
+
+                if n < how_much {
+                    break;
+                }
+            }
+
+            Ok(total)
         } else {
-            self.0.as_fd().read_volatile(buf)
+            self.target.as_fd().read_volatile(buf)
         }
     }
 }
 
-impl<T: Write + AsFd> WriteVolatile for MaybeBounce<T> {
+impl<T: Write + AsFd, const N: usize> WriteVolatile for MaybeBounce<T, N> {
     fn write_volatile<B: BitmapSlice>(
         &mut self,
         buf: &VolatileSlice<B>,
     ) -> Result<usize, VolatileMemoryError> {
-        if self.1 {
-            let mut bbuf = vec![0; buf.len()];
-            buf.copy_to(bbuf.as_mut_slice());
-            self.0
-                .write(bbuf.as_slice())
-                .map_err(VolatileMemoryError::IOError)
+        if self.should_bounce {
+            let mut bbuf = (N == 0).then(|| vec![0u8; buf.len()]);
+            let bbuf = bbuf
+                .as_deref_mut()
+                .unwrap_or(self.persistent_buffer.as_mut_slice());
+
+            let mut buf = buf.offset(0)?;
+            let mut total = 0;
+            while !buf.is_empty() {
+                let how_much = buf.copy_to(bbuf);
+                let n = self
+                    .target
+                    .write(&bbuf[..how_much])
+                    .map_err(VolatileMemoryError::IOError)?;
+                buf = buf.offset(n)?;
+                total += n;
+
+                if n < how_much {
+                    break;
+                }
+            }
+
+            Ok(total)
         } else {
-            self.0.as_fd().write_volatile(buf)
+            self.target.as_fd().write_volatile(buf)
         }
     }
 }
 
-impl<R: Read> Read for MaybeBounce<R> {
+impl<R: Read, const N: usize> Read for MaybeBounce<R, N> {
     fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
-        self.0.read(buf)
+        self.target.read(buf)
+    }
+}
+
+impl<W: Write, const N: usize> Write for MaybeBounce<W, N> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.target.write(buf)
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.target.flush()
     }
 }
 
-impl<S: Seek> Seek for MaybeBounce<S> {
+impl<S: Seek, const N: usize> Seek for MaybeBounce<S, N> {
     fn seek(&mut self, pos: SeekFrom) -> std::io::Result<u64> {
-        self.0.seek(pos)
+        self.target.seek(pos)
     }
 }
 
@@ -947,30 +1016,45 @@ mod tests {
     fn test_bounce() {
         let file_direct = TempFile::new().unwrap();
         let file_bounced = TempFile::new().unwrap();
+        let file_persistent_bounced = TempFile::new().unwrap();
 
         let mut data = (0..=255).collect_vec();
 
-        MaybeBounce(file_direct.as_file(), false)
+        MaybeBounce::new(file_direct.as_file(), false)
+            .write_all_volatile(&VolatileSlice::from(data.as_mut_slice()))
+            .unwrap();
+        MaybeBounce::new(file_bounced.as_file(), true)
             .write_all_volatile(&VolatileSlice::from(data.as_mut_slice()))
             .unwrap();
-        MaybeBounce(file_bounced.as_file(), true)
+        MaybeBounce::<_, 7>::new_persistent(file_persistent_bounced.as_file(), true)
             .write_all_volatile(&VolatileSlice::from(data.as_mut_slice()))
             .unwrap();
 
         let mut data_direct = vec![0u8; 256];
         let mut data_bounced = vec![0u8; 256];
+        let mut data_persistent_bounced = vec![0u8; 256];
 
         file_direct.as_file().seek(SeekFrom::Start(0)).unwrap();
         file_bounced.as_file().seek(SeekFrom::Start(0)).unwrap();
+        file_persistent_bounced
+            .as_file()
+            .seek(SeekFrom::Start(0))
+            .unwrap();
 
-        MaybeBounce(file_direct.as_file(), false)
+        MaybeBounce::new(file_direct.as_file(), false)
             .read_exact_volatile(&mut VolatileSlice::from(data_direct.as_mut_slice()))
             .unwrap();
-        MaybeBounce(file_bounced.as_file(), true)
+        MaybeBounce::new(file_bounced.as_file(), true)
             .read_exact_volatile(&mut VolatileSlice::from(data_bounced.as_mut_slice()))
             .unwrap();
+        MaybeBounce::<_, 7>::new_persistent(file_persistent_bounced.as_file(), true)
+            .read_exact_volatile(&mut VolatileSlice::from(
+                data_persistent_bounced.as_mut_slice(),
+            ))
+            .unwrap();
 
         assert_eq!(data_direct, data_bounced);
         assert_eq!(data_direct, data);
+        assert_eq!(data_persistent_bounced, data);
     }
 }

src/vmm/src/vstate/vm.rs

@@ -466,7 +466,7 @@ impl Vm {
                     .iter()
                     .any(|r| r.inner().guest_memfd != 0);
                 self.guest_memory()
-                    .dump(&mut MaybeBounce(&file, secret_hidden))
+                    .dump(&mut MaybeBounce::new(&file, secret_hidden))
                     .and_then(|_| self.swiotlb_regions().dump(&mut file))?;
                 self.reset_dirty_bitmap();
                 self.guest_memory().reset_dirty();