test: have uffd examples utilize page_size_kib information

Have the examples and tests utilize the page size information that
Firecracker sends as part of the initial UFFD payload. This allows us to
get rid of the separate uffd handler examples for the 4K and 2M test
cases.

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

Author: roypat

docs/snapshotting/handling-page-faults-on-snapshot-resume.md

@@ -162,7 +162,7 @@ connect/send data.
 ### Example
 
 An example of a handler process can be found
-[here](../../src/firecracker/examples/uffd/valid_4k_handler.rs). The process is
+[here](../../src/firecracker/examples/uffd/valid_handler.rs). The process is
 designed to tackle faults on a certain address by loading into memory the entire
 region that the address belongs to, but users can choose any other behavior that
 suits their use case best.

src/firecracker/Cargo.toml

@@ -50,16 +50,12 @@ serde_json = "1.0.113"
 tracing = ["log-instrument", "seccompiler/tracing", "utils/tracing", "vmm/tracing"]
 
 [[example]]
-name = "uffd_malicious_4k_handler"
-path = "examples/uffd/malicious_4k_handler.rs"
+name = "uffd_malicious_handler"
+path = "examples/uffd/malicious_handler.rs"
 
 [[example]]
-name = "uffd_valid_4k_handler"
-path = "examples/uffd/valid_4k_handler.rs"
-
-[[example]]
-name = "uffd_valid_2m_handler"
-path = "examples/uffd/valid_2m_handler.rs"
+name = "uffd_valid_handler"
+path = "examples/uffd/valid_handler.rs"
 
 [[example]]
 name = "uffd_fault_all_handler"

src/firecracker/examples/uffd/fault_all_handler.rs

@@ -11,7 +11,6 @@ use std::fs::File;
 use std::os::unix::net::UnixListener;
 
 use uffd_utils::{Runtime, UffdHandler};
-use utils::get_page_size;
 
 fn main() {
     let mut args = std::env::args();
@@ -24,13 +23,8 @@ fn main() {
     let listener = UnixListener::bind(uffd_sock_path).expect("Cannot bind to socket path");
     let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
 
-    // Populate a single page from backing memory file.
-    // This is just an example, probably, with the worst-case latency scenario,
-    // of how memory can be loaded in guest RAM.
-    let len = get_page_size().unwrap(); // page size does not matter, we fault in everything on the first fault
-
     let mut runtime = Runtime::new(stream, file);
-    runtime.run(len, |uffd_handler: &mut UffdHandler| {
+    runtime.run(|uffd_handler: &mut UffdHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
             .read_event()

src/firecracker/examples/uffd/malicious_4k_handler.rs renamed to src/firecracker/examples/uffd/malicious_handler.rs

@@ -23,7 +23,7 @@ fn main() {
     let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
 
     let mut runtime = Runtime::new(stream, file);
-    runtime.run(4096, |uffd_handler: &mut UffdHandler| {
+    runtime.run(|uffd_handler: &mut UffdHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
             .read_event()

src/firecracker/examples/uffd/uffd_utils.rs

@@ -30,6 +30,8 @@ pub struct GuestRegionUffdMapping {
     pub size: usize,
     /// Offset in the backend file/buffer where the region contents are.
     pub offset: u64,
+    /// The configured page size for this memory region.
+    pub page_size_kib: usize,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -49,18 +51,13 @@ pub struct MemRegion {
 #[derive(Debug)]
 pub struct UffdHandler {
     pub mem_regions: Vec<MemRegion>,
-    page_size: usize,
+    pub page_size: usize,
     backing_buffer: *const u8,
     uffd: Uffd,
 }
 
 impl UffdHandler {
-    pub fn from_unix_stream(
-        stream: &UnixStream,
-        page_size: usize,
-        backing_buffer: *const u8,
-        size: usize,
-    ) -> Self {
+    pub fn from_unix_stream(stream: &UnixStream, backing_buffer: *const u8, size: usize) -> Self {
         let mut message_buf = vec![0u8; 1024];
         let (bytes_read, file) = stream
             .recv_with_fd(&mut message_buf[..])
@@ -73,6 +70,8 @@ impl UffdHandler {
         let mappings = serde_json::from_str::<Vec<GuestRegionUffdMapping>>(&body)
             .expect("Cannot deserialize memory mappings.");
         let memsize: usize = mappings.iter().map(|r| r.size).sum();
+        // Page size is the same for all memory regions, so just grab the first one
+        let page_size = mappings.first().unwrap().page_size_kib;
 
         // Make sure memory size matches backing data size.
         assert_eq!(memsize, size);
@@ -214,7 +213,7 @@ impl Runtime {
     /// When uffd is polled, page fault is handled by
     /// calling `pf_event_dispatch` with corresponding
     /// uffd object passed in.
-    pub fn run(&mut self, page_size: usize, pf_event_dispatch: impl Fn(&mut UffdHandler)) {
+    pub fn run(&mut self, pf_event_dispatch: impl Fn(&mut UffdHandler)) {
         let mut pollfds = vec![];
 
         // Poll the stream for incoming uffds
@@ -249,7 +248,6 @@ impl Runtime {
                         // Handle new uffd from stream
                         let handler = UffdHandler::from_unix_stream(
                             &self.stream,
-                            page_size,
                             self.backing_memory,
                             self.backing_memory_size,
                         );
@@ -330,7 +328,7 @@ mod tests {
             let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
             // Update runtime with actual runtime
             let runtime = uninit_runtime.write(Runtime::new(stream, file));
-            runtime.run(4096, |_: &mut UffdHandler| {});
+            runtime.run(|_: &mut UffdHandler| {});
         });
 
         // wait for runtime thread to initialize itself
@@ -343,6 +341,7 @@ mod tests {
             base_host_virt_addr: 0,
             size: 0x1000,
             offset: 0,
+            page_size_kib: 4096,
         }];
         let dummy_memory_region_json = serde_json::to_string(&dummy_memory_region).unwrap();
 
@@ -375,6 +374,7 @@ mod tests {
             base_host_virt_addr: 0,
             size: 0,
             offset: 0,
+            page_size_kib: 4096,
         }];
         let error_memory_region_json = serde_json::to_string(&error_memory_region).unwrap();
         stream

src/firecracker/examples/uffd/valid_2m_handler.rs

@@ -11,7 +11,6 @@ use std::fs::File;
 use std::os::unix::net::UnixListener;
 
 use uffd_utils::{MemPageState, Runtime, UffdHandler};
-use utils::get_page_size;
 
 fn main() {
     let mut args = std::env::args();
@@ -24,13 +23,8 @@ fn main() {
     let listener = UnixListener::bind(uffd_sock_path).expect("Cannot bind to socket path");
     let (stream, _) = listener.accept().expect("Cannot listen on UDS socket");
 
-    // Populate a single page from backing memory file.
-    // This is just an example, probably, with the worst-case latency scenario,
-    // of how memory can be loaded in guest RAM.
-    let len = get_page_size().unwrap();
-
     let mut runtime = Runtime::new(stream, file);
-    runtime.run(len, |uffd_handler: &mut UffdHandler| {
+    runtime.run(|uffd_handler: &mut UffdHandler| {
         // Read an event from the userfaultfd.
         let event = uffd_handler
             .read_event()
@@ -40,7 +34,9 @@ fn main() {
         // We expect to receive either a Page Fault or Removed
         // event (if the balloon device is enabled).
         match event {
-            userfaultfd::Event::Pagefault { addr, .. } => uffd_handler.serve_pf(addr.cast(), len),
+            userfaultfd::Event::Pagefault { addr, .. } => {
+                uffd_handler.serve_pf(addr.cast(), uffd_handler.page_size)
+            }
             userfaultfd::Event::Remove { start, end } => uffd_handler.update_mem_state_mappings(
                 start as u64,
                 end as u64,

src/firecracker/examples/uffd/valid_4k_handler.rs renamed to src/firecracker/examples/uffd/valid_handler.rs

@@ -229,7 +229,7 @@ def uffd_handler_paths():
     """Build UFFD handler binaries."""
     handlers = {
         f"{handler}_handler": build_tools.get_example(f"uffd_{handler}_handler")
-        for handler in ["malicious_4k", "valid_4k", "valid_2m", "fault_all"]
+        for handler in ["malicious", "valid", "fault_all"]
     }
     yield handlers
 

tests/conftest.py

@@ -110,7 +110,7 @@ def test_valid_handler(uvm_plain, snapshot, uffd_handler_paths):
 
     # Spawn page fault handler process.
     _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["valid_4k_handler"], snapshot.mem
+        vm, uffd_handler_paths["valid_handler"], snapshot.mem
     )
 
     vm.restore_from_snapshot(snapshot, resume=True, uffd_path=SOCKET_PATH)
@@ -144,7 +144,7 @@ def test_malicious_handler(uvm_plain, snapshot, uffd_handler_paths):
 
     # Spawn page fault handler process.
     _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["malicious_4k_handler"], snapshot.mem
+        vm, uffd_handler_paths["malicious_handler"], snapshot.mem
     )
 
     # We expect Firecracker to freeze while resuming from a snapshot

tests/integration_tests/functional/test_uffd.py

@@ -109,7 +109,7 @@ def test_hugetlbfs_snapshot(
 
     # Spawn page fault handler process.
     _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["valid_2m_handler"], snapshot.mem
+        vm, uffd_handler_paths["valid_handler"], snapshot.mem
     )
 
     vm.restore_from_snapshot(snapshot, resume=True, uffd_path=SOCKET_PATH)
@@ -164,7 +164,7 @@ def test_hugetlbfs_diff_snapshot(microvm_factory, uvm_plain, uffd_handler_paths)
 
     # Spawn page fault handler process.
     _pf_handler = spawn_pf_handler(
-        vm, uffd_handler_paths["valid_2m_handler"], snapshot_merged.mem
+        vm, uffd_handler_paths["valid_handler"], snapshot_merged.mem
     )
 
     vm.restore_from_snapshot(snapshot_merged, resume=True, uffd_path=SOCKET_PATH)