fixup! feat(mem): introduce KVM slots per GuestMemoryRegion

Author: Manciukic

src/vmm/src/vstate/memory.rs

@@ -117,6 +117,54 @@ impl From<&GuestMemorySlot<'_>> for kvm_userspace_memory_region {
 }
 
 impl<'a> GuestMemorySlot<'a> {
+    /// Dumps the dirty pages in this slot onto the writer
+    pub(crate) fn dump_dirty<T: WriteVolatile + std::io::Seek>(
+        &self,
+        writer: &mut T,
+        kvm_bitmap: &[u64],
+        page_size: usize,
+    ) -> Result<(), GuestMemoryError> {
+        let firecracker_bitmap = self.slice.bitmap();
+        let mut write_size = 0;
+        let mut skip_size = 0;
+        let mut dirty_batch_start = 0;
+
+        for (i, v) in kvm_bitmap.iter().enumerate() {
+            for j in 0..64 {
+                let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
+                let page_offset = ((i * 64) + j) * page_size;
+                let is_firecracker_page_dirty = firecracker_bitmap.dirty_at(page_offset);
+
+                if is_kvm_page_dirty || is_firecracker_page_dirty {
+                    // We are at the start of a new batch of dirty pages.
+                    if skip_size > 0 {
+                        // Seek forward over the unmodified pages.
+                        writer
+                            .seek(SeekFrom::Current(skip_size.try_into().unwrap()))
+                            .unwrap();
+                        dirty_batch_start = page_offset;
+                    }
+                    write_size += page_size;
+                    skip_size = 0;
+                } else if write_size > 0 {
+                    // We are at the end of a batch of dirty pages.
+                    writer
+                        .write_all_volatile(&self.slice.subslice(dirty_batch_start, write_size)?)?;
+
+                    write_size = 0;
+                } else {
+                    skip_size += page_size;
+                }
+            }
+        }
+
+        if write_size > 0 {
+            writer.write_all_volatile(&self.slice.subslice(dirty_batch_start, write_size)?)?;
+        }
+
+        Ok(())
+    }
+
     /// Makes the slot host memory PROT_NONE (true) or PROT_READ|PROT_WRITE (false)
     pub(crate) fn protect(&self, protected: bool) -> Result<(), MemoryError> {
         let prot = if protected {
@@ -201,20 +249,18 @@ impl GuestRegionMmapExt {
         u32::try_from(u64_to_usize(self.len()) / self.slot_size).unwrap()
     }
 
-    pub(crate) fn mem_slot(&self, slot: u32) -> Option<GuestMemorySlot<'_>> {
-        if slot < self.slot_from || slot >= self.slot_from + self.slot_cnt() {
-            None
-        } else {
-            let offset = ((slot - self.slot_from) as u64) * (self.slot_size as u64);
-
-            Some(GuestMemorySlot {
-                slot,
-                guest_addr: self.start_addr().checked_add(offset)?,
-                slice: self
-                    .inner
-                    .get_slice(MemoryRegionAddress(offset), self.slot_size)
-                    .ok()?,
-            })
+    pub(crate) fn mem_slot(&self, slot: u32) -> GuestMemorySlot<'_> {
+        assert!(slot >= self.slot_from && slot < self.slot_from + self.slot_cnt());
+
+        let offset = ((slot - self.slot_from) as u64) * (self.slot_size as u64);
+
+        GuestMemorySlot {
+            slot,
+            guest_addr: self.start_addr().unchecked_add(offset),
+            slice: self
+                .inner
+                .get_slice(MemoryRegionAddress(offset), self.slot_size)
+                .expect("slot range should be valid"),
         }
     }
 
@@ -229,8 +275,7 @@ impl GuestRegionMmapExt {
             .enumerate()
             .map(|(i, b)| {
                 (
-                    self.mem_slot(self.slot_from + u32::try_from(i).unwrap())
-                        .unwrap(),
+                    self.mem_slot(self.slot_from + u32::try_from(i).unwrap()),
                     *b,
                 )
             })
@@ -597,11 +642,8 @@ impl GuestMemoryExtension for GuestMemoryMmap {
             .flat_map(|region| region.slots())
             .try_for_each(|(mem_slot, plugged)| {
                 if !plugged {
-                    writer
-                        .seek(SeekFrom::Current(
-                            i64::try_from(mem_slot.slice.len()).unwrap(),
-                        ))
-                        .unwrap();
+                    let ilen = i64::try_from(mem_slot.slice.len()).unwrap();
+                    writer.seek(SeekFrom::Current(ilen)).unwrap();
                 } else {
                     writer.write_all_volatile(&mem_slot.slice)?;
                 }
@@ -616,57 +658,19 @@ impl GuestMemoryExtension for GuestMemoryMmap {
         writer: &mut T,
         dirty_bitmap: &DirtyBitmap,
     ) -> Result<(), MemoryError> {
-        let mut writer_offset = 0;
         let page_size = get_page_size().map_err(MemoryError::PageSize)?;
 
         let write_result =
             self.iter()
                 .flat_map(|region| region.slots())
                 .try_for_each(|(mem_slot, plugged)| {
                     if !plugged {
-                        writer_offset += mem_slot.slice.len() as u64;
-                        return Ok(());
-                    }
-
-                    let kvm_bitmap = dirty_bitmap.get(&mem_slot.slot).unwrap();
-                    let firecracker_bitmap = mem_slot.slice.bitmap();
-                    let mut write_size = 0;
-                    let mut dirty_batch_start = 0;
-
-                    for (i, v) in kvm_bitmap.iter().enumerate() {
-                        for j in 0..64 {
-                            let is_kvm_page_dirty = ((v >> j) & 1u64) != 0u64;
-                            let page_offset = ((i * 64) + j) * page_size;
-                            let is_firecracker_page_dirty =
-                                firecracker_bitmap.dirty_at(page_offset);
-
-                            if is_kvm_page_dirty || is_firecracker_page_dirty {
-                                // We are at the start of a new batch of dirty pages.
-                                if write_size == 0 {
-                                    // Seek forward over the unmodified pages.
-                                    writer
-                                        .seek(SeekFrom::Start(writer_offset + page_offset as u64))
-                                        .unwrap();
-                                    dirty_batch_start = page_offset;
-                                }
-                                write_size += page_size;
-                            } else if write_size > 0 {
-                                // We are at the end of a batch of dirty pages.
-                                writer.write_all_volatile(
-                                    &mem_slot.slice.subslice(dirty_batch_start, write_size)?,
-                                )?;
-
-                                write_size = 0;
-                            }
-                        }
-                    }
-
-                    if write_size > 0 {
-                        writer.write_all_volatile(
-                            &mem_slot.slice.subslice(dirty_batch_start, write_size)?,
-                        )?;
+                        let ilen = i64::try_from(mem_slot.slice.len()).unwrap();
+                        writer.seek(SeekFrom::Current(ilen)).unwrap();
+                    } else {
+                        let kvm_bitmap = dirty_bitmap.get(&mem_slot.slot).unwrap();
+                        mem_slot.dump_dirty(writer, kvm_bitmap, page_size)?;
                     }
-                    writer_offset += mem_slot.slice.len() as u64;
                     Ok(())
                 });