Let running virtio devices follow guest memory layout changes

alioth/src/cpu/cpu_amd64/sev.rs

@@ -57,7 +57,7 @@ where
     fn update_snp_desc(&self, desc: &SevMetadataDesc) -> Result<()> {
         let mut cpuid_table = SnpCpuidInfo::new_zeroed();
         let ram_bus = self.ctx.board.memory.ram_bus();
-        let ram = ram_bus.lock_layout();
+        let ram = ram_bus.load();
         let page_type = match desc.type_ {
             SevDescType::SNP_DESC_MEM => SnpPageType::UNMEASURED,
             SevDescType::SNP_SECRETS => SnpPageType::SECRETS,

alioth/src/mem/addressable.rs

@@ -20,7 +20,7 @@ pub trait SlotBackend {
     fn size(&self) -> u64;
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 struct Slot<B>
 where
     B: SlotBackend,
@@ -52,7 +52,7 @@ where
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct Addressable<B>
 where
     B: SlotBackend,

alioth/src/mem/mapped.rs

@@ -24,14 +24,15 @@ use std::os::fd::FromRawFd;
 use std::os::fd::{AsFd, AsRawFd, BorrowedFd};
 use std::ptr::{NonNull, null_mut};
 use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
 
 #[cfg(target_os = "linux")]
 use libc::{MADV_HUGEPAGE, MFD_CLOEXEC};
 use libc::{
     MAP_ANONYMOUS, MAP_FAILED, MAP_PRIVATE, MAP_SHARED, MS_ASYNC, PROT_READ, PROT_WRITE, c_void,
     madvise, mmap, msync, munmap,
 };
-use parking_lot::{RwLock, RwLockReadGuard};
+use parking_lot::RwLock;
 use snafu::ResultExt;
 use zerocopy::{FromBytes, Immutable, IntoBytes};
 
@@ -180,14 +181,26 @@ impl ArcMemPages {
     }
 }
 
-#[derive(Debug)]
+/// A snapshot of the guest RAM layout.
+///
+/// A `Ram` does not change once a [`RamBus`] publishes it. Its slots hold
+/// reference-counted [`ArcMemPages`], so a snapshot keeps the host mappings
+/// it refers to alive, even after they are removed from the bus.
+#[derive(Debug, Clone, Default)]
 pub struct Ram {
     inner: Addressable<ArcMemPages>,
+    generation: u64,
 }
 
-#[derive(Debug)]
+/// The guest RAM layout, updated in the read-copy-update (RCU) way.
+///
+/// A reader takes a snapshot with [`RamBus::load`] and keeps it as long as
+/// it needs, while a writer publishes a new layout with [`RamBus::update`]
+/// without waiting for readers.
+#[derive(Debug, Default)]
 pub struct RamBus {
-    ram: RwLock<Ram>,
+    ram: RwLock<Arc<Ram>>,
+    generation: AtomicU64,
 }
 
 struct Iter<'m> {
@@ -397,74 +410,91 @@ impl Ram {
         }
         Ok(())
     }
-}
 
-impl Default for RamBus {
-    fn default() -> Self {
-        Self::new()
+    pub(crate) fn add(&mut self, gpa: u64, user_mem: ArcMemPages) -> Result<()> {
+        self.inner.add(gpa, user_mem)?;
+        Ok(())
+    }
+
+    pub(crate) fn remove(&mut self, gpa: u64) -> Result<ArcMemPages> {
+        self.inner.remove(gpa)
+    }
+
+    /// Returns the generation of the layout, see [`RamBus::generation`].
+    pub fn generation(&self) -> u64 {
+        self.generation
     }
 }
 
 impl RamBus {
-    pub fn lock_layout(&self) -> RwLockReadGuard<'_, Ram> {
-        self.ram.read()
+    pub fn new() -> Self {
+        Self::default()
     }
 
-    pub fn new() -> Self {
-        Self {
-            ram: RwLock::new(Ram {
-                inner: Addressable::default(),
-            }),
-        }
+    /// Returns a snapshot of the current layout.
+    ///
+    /// The snapshot stays valid for as long as the caller holds it, but does
+    /// not see later updates.
+    pub fn load(&self) -> Arc<Ram> {
+        self.ram.read().clone()
     }
 
-    pub(crate) fn add(&self, gpa: u64, user_mem: ArcMemPages) -> Result<(), Error> {
-        let mut ram = self.ram.write();
-        ram.inner.add(gpa, user_mem)?;
-        Ok(())
+    /// Returns the generation of the current layout.
+    ///
+    /// Every update increments it, so a reader can tell whether its snapshot
+    /// is the latest one without taking the lock.
+    pub fn generation(&self) -> u64 {
+        self.generation.load(Ordering::Acquire)
     }
 
-    pub(crate) fn remove(&self, gpa: u64) -> Result<ArcMemPages, Error> {
-        let mut ram = self.ram.write();
-        ram.inner.remove(gpa)
+    /// Copies the current layout, changes the copy with `f` and publishes it.
+    ///
+    /// If `f` fails, nothing is published, so a change of several slots is
+    /// either fully visible or not visible at all.
+    pub(crate) fn update<T>(&self, f: impl FnOnce(&mut Ram) -> Result<T>) -> Result<T> {
+        let mut current = self.ram.write();
+        let mut ram = Ram::clone(&current);
+        let ret = f(&mut ram)?;
+        ram.generation = current.generation + 1;
+        // A reader that sees the new generation before the new layout is
+        // published waits for it in load().
+        self.generation.store(ram.generation, Ordering::Release);
+        *current = Arc::new(ram);
+        Ok(ret)
     }
 
     pub fn read(&self, gpa: u64, buf: &mut [u8]) -> Result<()> {
-        let ram = self.ram.read();
-        ram.read(gpa, buf)
+        self.load().read(gpa, buf)
     }
 
     pub fn write(&self, gpa: u64, buf: &[u8]) -> Result<()> {
-        let ram = self.ram.read();
-        ram.write(gpa, buf)
+        self.load().write(gpa, buf)
     }
 
     pub fn read_t<T>(&self, gpa: u64) -> Result<T, Error>
     where
         T: FromBytes + IntoBytes,
     {
-        let ram = self.ram.read();
-        ram.read_t(gpa)
+        self.load().read_t(gpa)
     }
 
     pub fn write_t<T>(&self, gpa: u64, val: &T) -> Result<(), Error>
     where
         T: IntoBytes + Immutable,
     {
-        let ram = self.ram.read();
-        ram.write_t(gpa, val)
+        self.load().write_t(gpa, val)
     }
 
     pub fn read_range(&self, gpa: u64, len: u64, dst: &mut impl Write) -> Result<()> {
-        let ram = self.ram.read();
+        let ram = self.load();
         for r in ram.slice_iter(gpa, len) {
             dst.write_all(r?).context(error::Write)?;
         }
         Ok(())
     }
 
     pub fn write_range(&self, gpa: u64, len: u64, mut src: impl Read) -> Result<()> {
-        let ram = self.ram.read();
+        let ram = self.load();
         for r in ram.slice_iter_mut(gpa, len) {
             src.read_exact(r?).context(error::Read)?;
         }

alioth/src/mem/mapped_test.rs

@@ -38,11 +38,11 @@ fn test_ram_bus_read() {
     let mem2 = ArcMemPages::from_anonymous(PAGE_SIZE as usize, Some(prot), None).unwrap();
 
     if mem1.addr > mem2.addr {
-        bus.add(0x0, mem1).unwrap();
-        bus.add(PAGE_SIZE, mem2).unwrap();
+        bus.update(|ram| ram.add(0x0, mem1)).unwrap();
+        bus.update(|ram| ram.add(PAGE_SIZE, mem2)).unwrap();
     } else {
-        bus.add(0x0, mem2).unwrap();
-        bus.add(PAGE_SIZE, mem1).unwrap();
+        bus.update(|ram| ram.add(0x0, mem2)).unwrap();
+        bus.update(|ram| ram.add(PAGE_SIZE, mem1)).unwrap();
     }
 
     let data = MyStruct {
@@ -75,19 +75,49 @@ fn test_ram_bus_read() {
 
     let guest_iov = [(0, 16), (PAGE_SIZE - 16, 32), (2 * PAGE_SIZE - 16, 16)];
 
-    let locked_bus = bus.lock_layout();
-    let mut iov = locked_bus.translate_iov_mut(&guest_iov).unwrap();
+    let ram = bus.load();
+    let mut iov = ram.translate_iov_mut(&guest_iov).unwrap();
     assert_eq!(iov.len(), 4);
     assert_matches!((&*data).read_vectored(&mut iov), Ok(64));
     drop(iov);
-    let iov = locked_bus.translate_iov(&guest_iov).unwrap();
+    let iov = ram.translate_iov(&guest_iov).unwrap();
     assert_eq!(iov.len(), 4);
     let mut buf_read = Vec::new();
     assert_matches!(buf_read.write_vectored(&iov), Ok(64));
     assert_eq!(buf_read, data);
     drop(iov);
-    drop(locked_bus);
-    bus.remove(0x0).unwrap();
+    drop(ram);
+    bus.update(|ram| ram.remove(0x0)).unwrap();
+}
+
+#[test]
+fn test_ram_bus_snapshot() {
+    let bus = RamBus::new();
+    let empty = bus.load();
+    assert_eq!((bus.generation(), empty.generation()), (0, 0));
+
+    // Holding a snapshot does not block updates.
+    let mem = ArcMemPages::from_anonymous(PAGE_SIZE as usize, None, None).unwrap();
+    bus.update(|ram| ram.add(0x0, mem.clone())).unwrap();
+    assert_matches!(empty.read_t::<u32>(0x0), Err(_));
+    let snapshot = bus.load();
+    assert_eq!((bus.generation(), snapshot.generation()), (1, 1));
+    snapshot.write_t(0x0, &0xdeadbeef_u32).unwrap();
+
+    // A failed update publishes nothing.
+    let ret = bus.update(|ram| {
+        ram.add(PAGE_SIZE, mem.clone())?;
+        ram.add(0x0, mem.clone())
+    });
+    assert_matches!(ret, Err(_));
+    assert_matches!(bus.read_t::<u32>(PAGE_SIZE), Err(_));
+    assert_eq!(bus.generation(), 1);
+
+    // A removed slot stays accessible through older snapshots.
+    bus.update(|ram| ram.remove(0x0)).unwrap();
+    assert_eq!((bus.generation(), bus.load().generation()), (2, 2));
+    assert_matches!(bus.read_t::<u32>(0x0), Err(_));
+    assert_matches!(snapshot.read_t::<u32>(0x0), Ok(0xdeadbeef));
 }
 
 #[test]
@@ -97,9 +127,9 @@ fn test_ram_get_array_ptr() {
     // Two adjacent guest pages, backed by separate host mappings.
     for gpa in [0, PAGE_SIZE] {
         let pages = ArcMemPages::from_anonymous(PAGE_SIZE as usize, Some(prot), None).unwrap();
-        bus.add(gpa, pages).unwrap();
+        bus.update(|ram| ram.add(gpa, pages)).unwrap();
     }
-    let ram = bus.lock_layout();
+    let ram = bus.load();
 
     let ptr = ram.get_array_ptr::<u64>(PAGE_SIZE - 16, 2).unwrap();
     assert_eq!(ptr, ram.get_ptr::<u64>(PAGE_SIZE - 16).unwrap());

alioth/src/mem/mem.rs

@@ -266,6 +266,19 @@ struct LayoutCallbacks {
     updated: Vec<Box<dyn LayoutUpdated>>,
 }
 
+/// Iterates over the RAM-backed ranges of `region` placed at `addr`.
+fn ram_ranges(addr: u64, region: &MemRegion) -> impl Iterator<Item = (u64, &ArcMemPages)> {
+    let ranges = region.ranges.iter().scan(addr, |gpa, range| {
+        let start = *gpa;
+        *gpa += range.size();
+        Some((start, range))
+    });
+    ranges.filter_map(|(gpa, range)| match range {
+        MemRange::Ram(pages) | MemRange::DevMem { pages, .. } => Some((gpa, pages)),
+        MemRange::Emulated(_) | MemRange::Span(_) => None,
+    })
+}
+
 // lock order: region -> callbacks -> bus
 #[derive(Debug, Default)]
 pub struct Memory {
@@ -308,7 +321,7 @@ impl Memory {
     pub fn register_update_callback(&self, callback: Box<dyn LayoutUpdated>) -> Result<()> {
         let _regions = self.regions.lock();
         let mut callbacks = self.callbacks.lock();
-        let ram = self.ram_bus.lock_layout();
+        let ram = self.ram_bus.load();
         callback.ram_updated(&ram)?;
         callbacks.updated.push(callback);
         Ok(())
@@ -332,15 +345,20 @@ impl Memory {
         region.validate()?;
         let mut regions = self.regions.lock();
         regions.add(addr, region.clone())?;
-        let mut offset = 0;
         let callbacks = self.callbacks.lock();
-        let mut ram_updated = false;
+        let ram_ranges: Vec<_> = ram_ranges(addr, &region).collect();
+        let ram_updated = !ram_ranges.is_empty();
+        if ram_updated {
+            self.ram_bus.update(|ram| {
+                for (gpa, pages) in ram_ranges {
+                    ram.add(gpa, pages.clone())?;
+                }
+                Ok(())
+            })?;
+        }
+        let mut offset = 0;
         for range in &region.ranges {
             let gpa = addr + offset;
-            if let MemRange::Ram(pages) | MemRange::DevMem { pages, .. } = range {
-                self.ram_bus.add(gpa, pages.clone())?;
-                ram_updated = true;
-            }
             match range {
                 MemRange::Emulated(r) => {
                     let mut mmio_bus = self.mmio_bus.write();
@@ -361,7 +379,7 @@ impl Memory {
             offset += range.size();
         }
         if ram_updated {
-            let ram = self.ram_bus.lock_layout();
+            let ram = self.ram_bus.load();
             for update_callback in &callbacks.updated {
                 update_callback.ram_updated(&ram)?;
             }
@@ -374,15 +392,20 @@ impl Memory {
     }
 
     fn unmap_region(&self, addr: u64, region: &MemRegion) -> Result<()> {
-        let mut offset = 0;
         let callbacks = self.callbacks.lock();
-        let mut ram_updated = false;
+        let ram_ranges: Vec<_> = ram_ranges(addr, region).collect();
+        let ram_updated = !ram_ranges.is_empty();
+        if ram_updated {
+            self.ram_bus.update(|ram| {
+                for (gpa, _) in ram_ranges {
+                    ram.remove(gpa)?;
+                }
+                Ok(())
+            })?;
+        }
+        let mut offset = 0;
         for range in &region.ranges {
             let gpa = addr + offset;
-            if let MemRange::Ram(_) | MemRange::DevMem { .. } = range {
-                self.ram_bus.remove(gpa)?;
-                ram_updated = true;
-            }
             match range {
                 MemRange::Emulated(_) => {
                     let mut mmio_bus = self.mmio_bus.write();
@@ -403,7 +426,7 @@ impl Memory {
             offset += range.size();
         }
         if ram_updated {
-            let ram = self.ram_bus.lock_layout();
+            let ram = self.ram_bus.load();
             for callback in &callbacks.updated {
                 callback.ram_updated(&ram)?;
             }