[DO-NOT-MERGE] Virtio mem rebase preview

.buildkite/pipeline_perf.py

@@ -80,6 +80,11 @@
         "tests": "integration_tests/performance/test_mmds.py",
         "devtool_opts": "-c 1-10 -m 0",
     },
+    "memory-hotplug": {
+        "label": "memory-hotplug",
+        "tests": "integration_tests/performance/test_hotplug_memory.py",
+        "devtool_opts": "-c 1-10 -m 0",
+    },
 }
 
 REVISION_A = os.environ.get("REVISION_A")

resources/seccomp/aarch64-unknown-linux-musl.json

@@ -445,6 +445,18 @@
                     }
                 ]
             },
+            {
+                "syscall": "ioctl",
+                "args": [
+                    {
+                        "index": 1,
+                        "type": "dword",
+                        "op": "eq",
+                        "val": 1075883590,
+                        "comment": "KVM_SET_USER_MEMORY_REGION, used to (un)plug memory for the virtio-mem device"
+                    }
+                ]
+            },
             {
                 "syscall": "sched_yield",
                 "comment": "Used by the rust standard library in std::sync::mpmc. Firecracker uses mpsc channels from this module for inter-thread communication"
@@ -460,6 +472,10 @@
             {
                 "syscall": "restart_syscall",
                 "comment": "automatically issued by the kernel when specific timing-related syscalls (e.g. nanosleep) get interrupted by SIGSTOP"
+            },
+            {
+                "syscall": "mprotect",
+                "comment": "Used by memory hotplug to protect access to underlying host memory"
             }
         ]
     },

resources/seccomp/x86_64-unknown-linux-musl.json

@@ -457,6 +457,18 @@
                     }
                 ]
             },
+            {
+                "syscall": "ioctl",
+                "args": [
+                    {
+                        "index": 1,
+                        "type": "dword",
+                        "op": "eq",
+                        "val": 1075883590,
+                        "comment": "KVM_SET_USER_MEMORY_REGION, used to (un)plug memory for the virtio-mem device"
+                    }
+                ]
+            },
             {
                 "syscall": "sched_yield",
                 "comment": "Used by the rust standard library in std::sync::mpmc. Firecracker uses mpsc channels from this module for inter-thread communication"
@@ -472,6 +484,10 @@
             {
                 "syscall": "restart_syscall",
                 "comment": "automatically issued by the kernel when specific timing-related syscalls (e.g. nanosleep) get interrupted by SIGSTOP"
+            },
+            {
+                "syscall": "mprotect",
+                "comment": "Used by memory hotplug to protect access to underlying host memory"
             }
         ]
     },

src/firecracker/examples/uffd/on_demand_handler.rs

@@ -87,7 +87,7 @@ fn main() {
                         }
                     }
                     userfaultfd::Event::Remove { start, end } => {
-                        uffd_handler.mark_range_removed(start as u64, end as u64)
+                        uffd_handler.unregister_range(start, end)
                     }
                     _ => panic!("Unexpected event on userfaultfd"),
                 }

src/firecracker/examples/uffd/uffd_utils.rs

@@ -9,7 +9,7 @@
     dead_code
 )]
 
-use std::collections::{HashMap, HashSet};
+use std::collections::HashMap;
 use std::ffi::c_void;
 use std::fs::File;
 use std::os::unix::io::{AsRawFd, FromRawFd, IntoRawFd};
@@ -54,7 +54,6 @@ pub struct UffdHandler {
     pub page_size: usize,
     backing_buffer: *const u8,
     uffd: Uffd,
-    removed_pages: HashSet<u64>,
 }
 
 impl UffdHandler {
@@ -125,32 +124,30 @@ impl UffdHandler {
             page_size,
             backing_buffer,
             uffd,
-            removed_pages: HashSet::new(),
         }
     }
 
     pub fn read_event(&mut self) -> Result<Option<Event>, Error> {
         self.uffd.read_event()
     }
 
-    pub fn mark_range_removed(&mut self, start: u64, end: u64) {
-        let pfn_start = start / self.page_size as u64;
-        let pfn_end = end / self.page_size as u64;
-
-        for pfn in pfn_start..pfn_end {
-            self.removed_pages.insert(pfn);
-        }
+    pub fn unregister_range(&mut self, start: *mut c_void, end: *mut c_void) {
+        assert!(
+            (start as usize).is_multiple_of(self.page_size)
+                && (end as usize).is_multiple_of(self.page_size)
+                && end > start
+        );
+        // SAFETY: start and end are valid and provided by UFFD
+        let len = unsafe { end.offset_from_unsigned(start) };
+        self.uffd
+            .unregister(start, len)
+            .expect("range should be valid");
     }
 
     pub fn serve_pf(&mut self, addr: *mut u8, len: usize) -> bool {
         // Find the start of the page that the current faulting address belongs to.
         let dst = (addr as usize & !(self.page_size - 1)) as *mut libc::c_void;
         let fault_page_addr = dst as u64;
-        let fault_pfn = fault_page_addr / self.page_size as u64;
-
-        if self.removed_pages.contains(&fault_pfn) {
-            return self.zero_out(fault_page_addr);
-        }
 
         for region in self.mem_regions.iter() {
             if region.contains(fault_page_addr) {
@@ -193,14 +190,6 @@ impl UffdHandler {
 
         true
     }
-
-    fn zero_out(&mut self, addr: u64) -> bool {
-        match unsafe { self.uffd.zeropage(addr as *mut _, self.page_size, true) } {
-            Ok(_) => true,
-            Err(Error::ZeropageFailed(error)) if error as i32 == libc::EAGAIN => false,
-            r => panic!("Unexpected zeropage result: {:?}", r),
-        }
-    }
 }
 
 #[derive(Debug)]

src/firecracker/src/api_server/parsed_request.rs

@@ -28,6 +28,9 @@ use super::request::pmem::parse_put_pmem;
 use super::request::snapshot::{parse_patch_vm_state, parse_put_snapshot};
 use super::request::version::parse_get_version;
 use super::request::vsock::parse_put_vsock;
+use crate::api_server::request::hotplug::memory::{
+    parse_get_memory_hotplug, parse_patch_memory_hotplug, parse_put_memory_hotplug,
+};
 use crate::api_server::request::serial::parse_put_serial;
 
 #[derive(Debug)]
@@ -85,6 +88,9 @@ impl TryFrom<&Request> for ParsedRequest {
             }
             (Method::Get, "machine-config", None) => parse_get_machine_config(),
             (Method::Get, "mmds", None) => parse_get_mmds(),
+            (Method::Get, "hotplug", None) if path_tokens.next() == Some("memory") => {
+                parse_get_memory_hotplug()
+            }
             (Method::Get, _, Some(_)) => method_to_error(Method::Get),
             (Method::Put, "actions", Some(body)) => parse_put_actions(body),
             (Method::Put, "balloon", Some(body)) => parse_put_balloon(body),
@@ -103,6 +109,9 @@ impl TryFrom<&Request> for ParsedRequest {
             (Method::Put, "snapshot", Some(body)) => parse_put_snapshot(body, path_tokens.next()),
             (Method::Put, "vsock", Some(body)) => parse_put_vsock(body),
             (Method::Put, "entropy", Some(body)) => parse_put_entropy(body),
+            (Method::Put, "hotplug", Some(body)) if path_tokens.next() == Some("memory") => {
+                parse_put_memory_hotplug(body)
+            }
             (Method::Put, _, None) => method_to_error(Method::Put),
             (Method::Patch, "balloon", Some(body)) => parse_patch_balloon(body, path_tokens.next()),
             (Method::Patch, "drives", Some(body)) => parse_patch_drive(body, path_tokens.next()),
@@ -112,6 +121,9 @@ impl TryFrom<&Request> for ParsedRequest {
                 parse_patch_net(body, path_tokens.next())
             }
             (Method::Patch, "vm", Some(body)) => parse_patch_vm_state(body),
+            (Method::Patch, "hotplug", Some(body)) if path_tokens.next() == Some("memory") => {
+                parse_patch_memory_hotplug(body)
+            }
             (Method::Patch, _, None) => method_to_error(Method::Patch),
             (method, unknown_uri, _) => Err(RequestError::InvalidPathMethod(
                 unknown_uri.to_string(),
@@ -175,6 +187,7 @@ impl ParsedRequest {
                     Self::success_response_with_data(balloon_config)
                 }
                 VmmData::BalloonStats(stats) => Self::success_response_with_data(stats),
+                VmmData::VirtioMemStatus(data) => Self::success_response_with_data(data),
                 VmmData::InstanceInformation(info) => Self::success_response_with_data(info),
                 VmmData::VmmVersion(version) => Self::success_response_with_data(
                     &serde_json::json!({ "firecracker_version": version.as_str() }),
@@ -559,6 +572,9 @@ pub mod tests {
                 VmmData::BalloonStats(stats) => {
                     http_response(&serde_json::to_string(stats).unwrap(), 200)
                 }
+                VmmData::VirtioMemStatus(data) => {
+                    http_response(&serde_json::to_string(data).unwrap(), 200)
+                }
                 VmmData::Empty => http_response("", 204),
                 VmmData::FullVmConfig(cfg) => {
                     http_response(&serde_json::to_string(cfg).unwrap(), 200)