add dirty bitmap tracking abstractions

Signed-off-by: Alexandru Agache <[email protected]>

Author: alexandruag

benches/guest_memory.rs

@@ -4,6 +4,8 @@
 #![cfg(feature = "backend-mmap")]
 
 pub use criterion::{black_box, Criterion};
+
+use vm_memory::bitmap::Bitmap;
 use vm_memory::{GuestAddress, GuestMemory, GuestMemoryMmap};
 
 const REGION_SIZE: usize = 0x10_0000;
@@ -13,7 +15,10 @@ pub fn benchmark_for_guest_memory(c: &mut Criterion) {
     benchmark_find_region(c);
 }
 
-fn find_region(mem: &GuestMemoryMmap) {
+fn find_region<B>(mem: &GuestMemoryMmap<B>)
+where
+    B: Bitmap + 'static,
+{
     for i in 0..REGIONS_COUNT {
         let _ = mem
             .find_region(black_box(GuestAddress(i * REGION_SIZE as u64)))

benches/main.rs

@@ -17,7 +17,7 @@ use volatile::benchmark_for_volatile;
 #[cfg(feature = "backend-mmap")]
 // Use this function with caution. It does not check against overflows
 // and `GuestMemoryMmap::from_ranges` errors.
-fn create_guest_memory_mmap(size: usize, count: u64) -> GuestMemoryMmap {
+fn create_guest_memory_mmap(size: usize, count: u64) -> GuestMemoryMmap<()> {
     let mut regions: Vec<(GuestAddress, usize)> = Vec::new();
     for i in 0..count {
         regions.push((GuestAddress(i * size as u64), size));

src/atomic.rs

@@ -142,11 +142,10 @@ impl<M: GuestMemory> GuestMemoryExclusiveGuard<'_, M> {
 #[cfg(feature = "backend-mmap")]
 mod tests {
     use super::*;
-    use crate::{
-        GuestAddress, GuestMemory, GuestMemoryMmap, GuestMemoryRegion, GuestRegionMmap, GuestUsize,
-        MmapRegion,
-    };
+    use crate::{GuestAddress, GuestMemory, GuestMemoryRegion, GuestUsize, MmapRegion};
 
+    type GuestMemoryMmap = crate::GuestMemoryMmap<()>;
+    type GuestRegionMmap = crate::GuestRegionMmap<()>;
     type GuestMemoryMmapAtomic = GuestMemoryAtomic<GuestMemoryMmap>;
 
     #[test]

src/bitmap/mod.rs

@@ -0,0 +1,105 @@
+// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
+
+//! This module holds abstractions that enable tracking the areas dirtied by writes of a specified
+//! length to a given offset. In particular, this is used to track write accesses within a
+//! `GuestMemoryRegion` object, and the resulting bitmaps can then be aggregated to build the
+//! global view for an entire `GuestMemory` object.
+
+use std::fmt::Debug;
+
+use crate::{GuestMemory, GuestMemoryRegion};
+
+/// Trait implemented by types that support creating `BitmapSlice` objects.
+pub trait WithBitmapSlice<'a> {
+    /// Type of the bitmap slice.
+    type S: BitmapSlice;
+}
+
+/// Trait used to represent that a `BitmapSlice` is a `Bitmap` itself, but also satisfies the
+/// restriction that slices created from it have the same type as `Self`.
+pub trait BitmapSlice:
+    Bitmap + Clone + Copy + Debug + for<'a> WithBitmapSlice<'a, S = Self>
+{
+}
+
+/// Common bitmap operations. Using Higher-Rank Trait Bounds (HRTBs) to effectively define
+/// an associated type that has a lifetime parameter, without tagging the `Bitmap` trait with
+/// a lifetime as well.
+///
+/// Using an associated type allows implementing the `Bitmap` and `BitmapSlice` functionality
+/// as a zero-cost abstraction when providing trivial implementations such as the one
+/// defined for `()`.
+// These methods represent the core functionality that's required by `vm-memory` abstractions
+// to implement generic tracking logic, as well as tests that can be reused by different backends.
+pub trait Bitmap: for<'a> WithBitmapSlice<'a> {
+    /// Mark the memory range specified by the given `offset` and `len` as dirtied.
+    fn mark_dirty(&self, offset: usize, len: usize);
+
+    /// Check whether the specified `offset` is marked as dirty.
+    fn dirty_at(&self, offset: usize) -> bool;
+
+    /// Return a `<Self as WithBitmapSlice>::S` slice of the current bitmap, starting at
+    /// the specified `offset`.
+    fn slice_at(&self, offset: usize) -> <Self as WithBitmapSlice>::S;
+}
+
+/// A no-op `Bitmap` implementation that can be provided for backends that do not actually
+/// require the tracking functionality.
+
+impl<'a> WithBitmapSlice<'a> for () {
+    type S = Self;
+}
+
+impl BitmapSlice for () {}
+
+impl Bitmap for () {
+    fn mark_dirty(&self, _offset: usize, _len: usize) {}
+
+    fn dirty_at(&self, _offset: usize) -> bool {
+        false
+    }
+
+    fn slice_at(&self, _offset: usize) -> Self {}
+}
+
+/// A `Bitmap` and `BitmapSlice` implementation for `Option<B>`.
+
+impl<'a, B> WithBitmapSlice<'a> for Option<B>
+where
+    B: WithBitmapSlice<'a>,
+{
+    type S = Option<B::S>;
+}
+
+impl<B: BitmapSlice> BitmapSlice for Option<B> {}
+
+impl<B: Bitmap> Bitmap for Option<B> {
+    fn mark_dirty(&self, offset: usize, len: usize) {
+        if let Some(inner) = self {
+            inner.mark_dirty(offset, len)
+        }
+    }
+
+    fn dirty_at(&self, offset: usize) -> bool {
+        if let Some(inner) = self {
+            return inner.dirty_at(offset);
+        }
+        false
+    }
+
+    fn slice_at(&self, offset: usize) -> Option<<B as WithBitmapSlice>::S> {
+        if let Some(inner) = self {
+            return Some(inner.slice_at(offset));
+        }
+        None
+    }
+}
+
+/// Helper type alias for referring to the `BitmapSlice` concrete type associated with
+/// an object `B: WithBitmapSlice<'a>`.
+pub type BS<'a, B> = <B as WithBitmapSlice<'a>>::S;
+
+/// Helper type alias for referring to the `BitmapSlice` concrete type associated with
+/// the memory regions of an object `M: GuestMemory`.
+pub type MS<'a, M> = BS<'a, <<M as GuestMemory>::R as GuestMemoryRegion>::B>;

src/bytes.rs

@@ -18,7 +18,7 @@ use std::slice::{from_raw_parts, from_raw_parts_mut};
 use std::sync::atomic::Ordering;
 
 use crate::atomic_integer::AtomicInteger;
-use crate::VolatileSlice;
+use crate::volatile_memory::VolatileSlice;
 
 /// Types for which it is safe to initialize from raw data.
 ///

src/guest_memory.rs

@@ -32,6 +32,14 @@
 //! Whenever a collection of `GuestMemoryRegion` objects is mutable,
 //! [`GuestAddressSpace`](trait.GuestAddressSpace.html) should be implemented
 //! for clients to obtain a [`GuestMemory`] reference or smart pointer.
+//!
+//! The `GuestMemoryRegion` trait has an associated `B: Bitmap` type which is used to handle
+//! dirty bitmap tracking. Backends are free to define the granularity (or whether tracking is
+//! actually performed at all). Those that do implement tracking functionality are expected to
+//! ensure the correctness of the underlying `Bytes` implementation. The user has to explicitly
+//! record (using the handle returned by `GuestRegionMmap::bitmap`) write accesses performed
+//! via pointers, references, or slices returned by methods of `GuestMemory`,`GuestMemoryRegion`,
+//! `VolatileSlice`, `VolatileRef`, or `VolatileArrayRef`.
 
 use std::convert::From;
 use std::fmt::{self, Display};
@@ -43,8 +51,9 @@ use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
 use crate::address::{Address, AddressValue};
+use crate::bitmap::{Bitmap, BS, MS};
 use crate::bytes::{AtomicAccess, Bytes};
-use crate::volatile_memory;
+use crate::volatile_memory::{self, VolatileSlice};
 
 static MAX_ACCESS_CHUNK: usize = 4096;
 
@@ -165,6 +174,9 @@ impl FileOffset {
 /// Represents a continuous region of guest physical memory.
 #[allow(clippy::len_without_is_empty)]
 pub trait GuestMemoryRegion: Bytes<MemoryRegionAddress, E = Error> {
+    /// Type used for dirty memory tracking.
+    type B: Bitmap;
+
     /// Returns the size of the region.
     fn len(&self) -> GuestUsize;
 
@@ -177,6 +189,9 @@ pub trait GuestMemoryRegion: Bytes<MemoryRegionAddress, E = Error> {
         self.start_addr().unchecked_add(self.len() - 1)
     }
 
+    /// Borrow the associated `Bitmap` object.
+    fn bitmap(&self) -> &Self::B;
+
     /// Returns the given address if it is within this region.
     fn check_address(&self, addr: MemoryRegionAddress) -> Option<MemoryRegionAddress> {
         if self.address_in_range(addr) {
@@ -245,7 +260,9 @@ pub trait GuestMemoryRegion: Bytes<MemoryRegionAddress, E = Error> {
     ///
     /// # Safety
     ///
-    /// Unsafe because of possible aliasing.
+    /// Unsafe because of possible aliasing. Mutable accesses performed through the
+    /// returned slice are not visible to the dirty bitmap tracking functionality of
+    /// the region, and must be manually recorded using the associated bitmap object.
     unsafe fn as_mut_slice(&self) -> Option<&mut [u8]> {
         None
     }
@@ -257,7 +274,7 @@ pub trait GuestMemoryRegion: Bytes<MemoryRegionAddress, E = Error> {
         &self,
         offset: MemoryRegionAddress,
         count: usize,
-    ) -> Result<volatile_memory::VolatileSlice> {
+    ) -> Result<VolatileSlice<BS<Self::B>>> {
         Err(Error::HostAddressNotAvailable)
     }
 
@@ -286,7 +303,7 @@ pub trait GuestMemoryRegion: Bytes<MemoryRegionAddress, E = Error> {
     /// assert_eq!(r.load(), v);
     /// # }
     /// ```
-    fn as_volatile_slice(&self) -> Result<volatile_memory::VolatileSlice> {
+    fn as_volatile_slice(&self) -> Result<VolatileSlice<BS<Self::B>>> {
         self.get_slice(MemoryRegionAddress(0), self.len() as usize)
     }
 
@@ -386,26 +403,26 @@ pub trait GuestAddressSpace {
 }
 
 impl<M: GuestMemory> GuestAddressSpace for &M {
-    type T = Self;
     type M = M;
+    type T = Self;
 
     fn memory(&self) -> Self {
         self
     }
 }
 
 impl<M: GuestMemory> GuestAddressSpace for Rc<M> {
-    type T = Self;
     type M = M;
+    type T = Self;
 
     fn memory(&self) -> Self {
         self.clone()
     }
 }
 
 impl<M: GuestMemory> GuestAddressSpace for Arc<M> {
-    type T = Self;
     type M = M;
+    type T = Self;
 
     fn memory(&self) -> Self {
         self.clone()
@@ -705,11 +722,7 @@ pub trait GuestMemory {
 
     /// Returns a [`VolatileSlice`](struct.VolatileSlice.html) of `count` bytes starting at
     /// `addr`.
-    fn get_slice(
-        &self,
-        addr: GuestAddress,
-        count: usize,
-    ) -> Result<volatile_memory::VolatileSlice> {
+    fn get_slice(&self, addr: GuestAddress, count: usize) -> Result<VolatileSlice<MS<Self>>> {
         self.to_region_addr(addr)
             .ok_or(Error::InvalidGuestAddress(addr))
             .and_then(|(r, addr)| r.get_slice(addr, count))
@@ -836,22 +849,29 @@ impl<T: GuestMemory> Bytes<GuestAddress> for T {
             // Check if something bad happened before doing unsafe things.
             assert!(offset <= count);
             if let Some(dst) = unsafe { region.as_mut_slice() } {
-                // This is safe cause `start` and `len` are within the `region`.
+                // This is safe cause `start` and `len` are within the `region`, and we manually
+                // record the dirty status of the written range below.
                 let start = caddr.raw_value() as usize;
                 let end = start + len;
-                loop {
+                let result = loop {
                     match src.read(&mut dst[start..end]) {
                         Ok(n) => break Ok(n),
                         Err(ref e) if e.kind() == std::io::ErrorKind::Interrupted => continue,
                         Err(e) => break Err(Error::IOError(e)),
                     }
-                }
+                };
+
+                region.bitmap().mark_dirty(start, len);
+                result
             } else {
                 let len = std::cmp::min(len, MAX_ACCESS_CHUNK);
                 let mut buf = vec![0u8; len].into_boxed_slice();
                 loop {
                     match src.read(&mut buf[..]) {
                         Ok(bytes_read) => {
+                            // We don't need to update the dirty bitmap manually here because it's
+                            // expected to be handled by the logic within the `Bytes`
+                            // implementation for the region object.
                             let bytes_written = region.write(&buf[0..bytes_read], caddr)?;
                             assert_eq!(bytes_written, bytes_read);
                             break Ok(bytes_read);
@@ -1008,14 +1028,17 @@ mod tests {
     #[cfg(feature = "backend-mmap")]
     use crate::bytes::ByteValued;
     #[cfg(feature = "backend-mmap")]
-    use crate::{GuestAddress, GuestMemoryMmap};
+    use crate::GuestAddress;
     #[cfg(feature = "backend-mmap")]
     use std::io::Cursor;
     #[cfg(feature = "backend-mmap")]
     use std::time::{Duration, Instant};
 
     use vmm_sys_util::tempfile::TempFile;
 
+    #[cfg(feature = "backend-mmap")]
+    type GuestMemoryMmap = crate::GuestMemoryMmap<()>;
+
     #[cfg(feature = "backend-mmap")]
     fn make_image(size: u8) -> Vec<u8> {
         let mut image: Vec<u8> = Vec::with_capacity(size as usize);

src/lib.rs

@@ -31,6 +31,8 @@ pub use atomic::{GuestMemoryAtomic, GuestMemoryLoadGuard};
 mod atomic_integer;
 pub use atomic_integer::AtomicInteger;
 
+pub mod bitmap;
+
 pub mod bytes;
 pub use bytes::{AtomicAccess, ByteValued, Bytes};