[virtio-mem] Add hotplug/hotunplug support with statically allocated memory region

docs/device-api.md

@@ -104,6 +104,7 @@ specification:
 | `MemoryHotplugConfig`     | total_size_mib     |    O     |       O        |      O       |        O         |     O      |      O       |     O      |   **R**    |
 |                           | slot_size_mib      |    O     |       O        |      O       |        O         |     O      |      O       |     O      |   **R**    |
 |                           | block_size_mi      |    O     |       O        |      O       |        O         |     O      |      O       |     O      |   **R**    |
+| `MemoryHotplugSizeUpdate` | requested_size_mib |    O     |       O        |      O       |        O         |     O      |      O       |     O      |   **R**    |
 
 \* `Drive`'s `drive_id`, `is_root_device` and `partuuid` can be configured by
 either virtio-block or vhost-user-block devices.

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,25 @@ 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) {
+        // 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 +185,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,7 +28,7 @@ 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_put_memory_hotplug,
+    parse_get_memory_hotplug, parse_patch_memory_hotplug, parse_put_memory_hotplug,
 };
 use crate::api_server::request::serial::parse_put_serial;
 
@@ -119,6 +119,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(),

src/firecracker/src/api_server/request/hotplug/memory.rs

@@ -4,7 +4,7 @@
 use micro_http::Body;
 use vmm::logger::{IncMetric, METRICS};
 use vmm::rpc_interface::VmmAction;
-use vmm::vmm_config::memory_hotplug::MemoryHotplugConfig;
+use vmm::vmm_config::memory_hotplug::{MemoryHotplugConfig, MemoryHotplugSizeUpdate};
 
 use crate::api_server::parsed_request::{ParsedRequest, RequestError};
 
@@ -23,11 +23,23 @@ pub(crate) fn parse_get_memory_hotplug() -> Result<ParsedRequest, RequestError>
     Ok(ParsedRequest::new_sync(VmmAction::GetMemoryHotplugStatus))
 }
 
+pub(crate) fn parse_patch_memory_hotplug(body: &Body) -> Result<ParsedRequest, RequestError> {
+    METRICS.patch_api_requests.hotplug_memory_count.inc();
+    let config =
+        serde_json::from_slice::<MemoryHotplugSizeUpdate>(body.raw()).inspect_err(|_| {
+            METRICS.patch_api_requests.hotplug_memory_fails.inc();
+        })?;
+    Ok(ParsedRequest::new_sync(VmmAction::UpdateMemoryHotplugSize(
+        config,
+    )))
+}
+
 #[cfg(test)]
 mod tests {
     use vmm::devices::virtio::mem::{
         VIRTIO_MEM_DEFAULT_BLOCK_SIZE_MIB, VIRTIO_MEM_DEFAULT_SLOT_SIZE_MIB,
     };
+    use vmm::vmm_config::memory_hotplug::MemoryHotplugSizeUpdate;
 
     use super::*;
     use crate::api_server::parsed_request::tests::vmm_action_from_request;
@@ -80,4 +92,27 @@ mod tests {
             VmmAction::GetMemoryHotplugStatus
         );
     }
+
+    #[test]
+    fn test_parse_patch_memory_hotplug_request() {
+        parse_patch_memory_hotplug(&Body::new("invalid_payload")).unwrap_err();
+
+        // PATCH with invalid fields.
+        let body = r#"{
+            "requested_size_mib": "bar"
+        }"#;
+        parse_patch_memory_hotplug(&Body::new(body)).unwrap_err();
+
+        // PATCH with valid input fields.
+        let body = r#"{
+            "requested_size_mib": 2048
+        }"#;
+        let expected_config = MemoryHotplugSizeUpdate {
+            requested_size_mib: 2048,
+        };
+        assert_eq!(
+            vmm_action_from_request(parse_patch_memory_hotplug(&Body::new(body)).unwrap()),
+            VmmAction::UpdateMemoryHotplugSize(expected_config)
+        );
+    }
 }

src/firecracker/swagger/firecracker.yaml

@@ -549,6 +549,26 @@ paths:
           description: Internal server error
           schema:
             $ref: "#/definitions/Error"
+    patch:
+      summary: Updates the size of the hotpluggable memory region
+      operationId: patchMemoryHotplug
+      description:
+        Updates the size of the hotpluggable memory region. The guest will plug and unplug memory to
+        hit the requested memory.
+      parameters:
+        - name: body
+          in: body
+          description: Hotpluggable memory size update
+          required: true
+          schema:
+            $ref: "#/definitions/MemoryHotplugSizeUpdate"
+      responses:
+        204:
+          description: Hotpluggable memory configured
+        default:
+          description: Internal server error
+          schema:
+            $ref: "#/definitions/Error"
     get:
       summary: Retrieves the status of the hotpluggable memory
       operationId: getMemoryHotplug
@@ -1422,6 +1442,15 @@ definitions:
         description: (Logical) Block size for the hotpluggable memory in MiB. This will determine the logical
           granularity of hot-plug memory for the guest. Refer to the device documentation on how to tune this value.
 
+  MemoryHotplugSizeUpdate:
+    type: object
+    description:
+      An update to the size of the hotpluggable memory region.
+    properties:
+      requested_size_mib:
+        type: integer
+        description: New target region size.
+
   MemoryHotplugStatus:
     type: object
     description:

src/vmm/Cargo.toml

@@ -23,6 +23,7 @@ aws-lc-rs = { version = "1.14.0", features = ["bindgen"] }
 base64 = "0.22.1"
 bincode = { version = "2.0.1", features = ["serde"] }
 bitflags = "2.9.4"
+bitvec = { version = "1.0.1", features = ["atomic", "serde"] }
 byteorder = "1.5.0"
 crc64 = "2.0.0"
 derive_more = { version = "2.0.1", default-features = false, features = [

src/vmm/src/builder.rs

@@ -32,7 +32,7 @@ use crate::device_manager::{
 use crate::devices::acpi::vmgenid::VmGenIdError;
 use crate::devices::virtio::balloon::Balloon;
 use crate::devices::virtio::block::device::Block;
-use crate::devices::virtio::mem::VirtioMem;
+use crate::devices::virtio::mem::{VIRTIO_MEM_GUEST_ADDRESS, VirtioMem};
 use crate::devices::virtio::net::Net;
 use crate::devices::virtio::rng::Entropy;
 use crate::devices::virtio::vsock::{Vsock, VsockUnixBackend};
@@ -44,6 +44,7 @@ use crate::persist::{MicrovmState, MicrovmStateError};
 use crate::resources::VmResources;
 use crate::seccomp::BpfThreadMap;
 use crate::snapshot::Persist;
+use crate::utils::mib_to_bytes;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::MachineConfigError;
 use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
@@ -172,6 +173,18 @@ pub fn build_microvm_for_boot(
     let (mut vcpus, vcpus_exit_evt) = vm.create_vcpus(vm_resources.machine_config.vcpu_count)?;
     vm.register_dram_memory_regions(guest_memory)?;
 
+    // Allocate memory as soon as possible to make hotpluggable memory available to all consumers,
+    // before they clone the GuestMemoryMmap object
+    if let Some(memory_hotplug) = &vm_resources.memory_hotplug {
+        let hotplug_memory_region = vm_resources
+            .allocate_memory_region(
+                VIRTIO_MEM_GUEST_ADDRESS,
+                mib_to_bytes(memory_hotplug.total_size_mib),
+            )
+            .map_err(StartMicrovmError::GuestMemory)?;
+        vm.register_hotpluggable_memory_region(hotplug_memory_region)?;
+    }
+
     let mut device_manager = DeviceManager::new(
         event_manager,
         &vcpus_exit_evt,