feat(snapshot): allow pmem backing file path override on restore

CHANGELOG.md

@@ -10,6 +10,14 @@ and this project adheres to
 
 ### Added
 
+- [#5896](https://github.com/firecracker-microvm/firecracker/pull/5896): Add
+  support for overriding virtio-pmem device backing file paths on snapshot
+  restore via the new `pmem_overrides` field on the `PUT /snapshot/load` API.
+  This mirrors the existing network and vsock override mechanisms and is useful
+  when the host file path baked into the snapshot is no longer valid (for
+  example, when restoring on a different host or under a different jailer
+  chroot).
+
 ### Changed
 
 ### Deprecated

docs/pmem.md

@@ -260,6 +260,38 @@ same backing file as it was configured in the first place. This means all
 `virtio-pmem` backing files should be present in the same locations during
 restore as they were during initial `virtio-pmem` configuration.
 
+### Overriding pmem backing file paths on restore
+
+When the host file path baked into the snapshot is no longer valid - for
+example, when restoring on a different host or under a different jailer chroot -
+the `pmem_overrides` parameter of the snapshot load API can be used to point
+each `virtio-pmem` device at the same backing file in its new location.
+
+This is only intended for relocating the original backing file; it is not a
+mechanism for swapping in a different backing file on restore. The overriding
+file must therefore be the same backing file (and hence the same size) as the
+one used when the snapshot was taken.
+
+```bash
+curl --unix-socket /tmp/firecracker.socket -i \
+    -X PUT 'http://localhost/snapshot/load' \
+    -H  'Accept: application/json' \
+    -H  'Content-Type: application/json' \
+    -d '{
+            "snapshot_path": "./snapshot_file",
+            "mem_backend": {
+                "backend_path": "./mem_file",
+                "backend_type": "File"
+            },
+            "pmem_overrides": [
+                {
+                    "id": "pmem0",
+                    "path_on_host": "/new/path/to/pmem0.img"
+                }
+            ]
+    }'
+```
+
 ## Performance
 
 Even though `virtio-pmem` allows for the direct access of host pages from the

src/firecracker/src/api_server/request/snapshot.rs

@@ -112,6 +112,7 @@ fn parse_put_snapshot_load(body: &Body) -> Result<ParsedRequest, RequestError> {
         network_overrides: snapshot_config.network_overrides,
         vsock_override: snapshot_config.vsock_override,
         clock_realtime: snapshot_config.clock_realtime,
+        pmem_overrides: snapshot_config.pmem_overrides,
     };
 
     // Construct the `ParsedRequest` object.
@@ -127,7 +128,9 @@ fn parse_put_snapshot_load(body: &Body) -> Result<ParsedRequest, RequestError> {
 
 #[cfg(test)]
 mod tests {
-    use vmm::vmm_config::snapshot::{MemBackendConfig, MemBackendType, NetworkOverride};
+    use vmm::vmm_config::snapshot::{
+        MemBackendConfig, MemBackendType, NetworkOverride, PmemOverride,
+    };
 
     use super::*;
     use crate::api_server::parsed_request::tests::{depr_action_from_req, vmm_action_from_request};
@@ -191,6 +194,7 @@ mod tests {
             network_overrides: vec![],
             vsock_override: None,
             clock_realtime: false,
+            pmem_overrides: vec![],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -223,6 +227,7 @@ mod tests {
             network_overrides: vec![],
             vsock_override: None,
             clock_realtime: false,
+            pmem_overrides: vec![],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -255,6 +260,7 @@ mod tests {
             network_overrides: vec![],
             vsock_override: None,
             clock_realtime: false,
+            pmem_overrides: vec![],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -296,6 +302,49 @@ mod tests {
             }],
             vsock_override: None,
             clock_realtime: false,
+            pmem_overrides: vec![],
+        };
+        let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
+        assert!(
+            parsed_request
+                .parsing_info()
+                .take_deprecation_message()
+                .is_none()
+        );
+        assert_eq!(
+            vmm_action_from_request(parsed_request),
+            VmmAction::LoadSnapshot(expected_config)
+        );
+
+        let body = r#"{
+            "snapshot_path": "foo",
+            "mem_backend": {
+                "backend_path": "bar",
+                "backend_type": "File"
+            },
+            "resume_vm": true,
+            "pmem_overrides": [
+                {
+                    "id": "pmem0",
+                    "path_on_host": "/new/path/pmem0.img"
+                }
+            ]
+        }"#;
+        let expected_config = LoadSnapshotParams {
+            snapshot_path: PathBuf::from("foo"),
+            mem_backend: MemBackendConfig {
+                backend_path: PathBuf::from("bar"),
+                backend_type: MemBackendType::File,
+            },
+            track_dirty_pages: false,
+            resume_vm: true,
+            network_overrides: vec![],
+            vsock_override: None,
+            clock_realtime: false,
+            pmem_overrides: vec![PmemOverride {
+                id: String::from("pmem0"),
+                path_on_host: String::from("/new/path/pmem0.img"),
+            }],
         };
         let mut parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert!(
@@ -325,6 +374,7 @@ mod tests {
             network_overrides: vec![],
             vsock_override: None,
             clock_realtime: false,
+            pmem_overrides: vec![],
         };
         let parsed_request = parse_put_snapshot(&Body::new(body), Some("load")).unwrap();
         assert_eq!(

src/firecracker/swagger/firecracker.yaml

@@ -1655,6 +1655,26 @@ definitions:
         description:
           The new path for the backing Unix Domain Socket.
 
+  PmemOverride:
+    type: object
+    description:
+      Allows for changing the backing host file of a virtio-pmem device
+      during snapshot restore.
+    required:
+      - id
+      - path_on_host
+    properties:
+      id:
+        type: string
+        description:
+          The ID of the pmem device to modify.
+      path_on_host:
+        type: string
+        description:
+          The new host path for the pmem device's backing file. The file must
+          be at least as large as the backing file used when the snapshot was
+          taken.
+
   SnapshotLoadParams:
     type: object
     description:
@@ -1710,6 +1730,11 @@ definitions:
           elapsed since the snapshot was taken. When false (default), kvmclock resumes
           from where it was at snapshot time. This option may be extended to other clock
           sources and CPU architectures in the future."
+      pmem_overrides:
+        type: array
+        description: Pmem device backing file paths to override on snapshot restore.
+        items:
+          $ref: "#/definitions/PmemOverride"
 
 
   TokenBucket:

src/vmm/src/persist.rs

@@ -34,7 +34,9 @@ use crate::utils::u64_to_usize;
 use crate::vmm_config::boot_source::BootSourceConfig;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::{HugePageConfig, MachineConfigError, MachineConfigUpdate};
-use crate::vmm_config::snapshot::{CreateSnapshotParams, LoadSnapshotParams, MemBackendType};
+use crate::vmm_config::snapshot::{
+    CreateSnapshotParams, LoadSnapshotParams, MemBackendType, PmemOverride,
+};
 use crate::vstate::kvm::KvmState;
 use crate::vstate::memory::{
     self, GuestMemoryState, GuestRegionMmap, GuestRegionType, MemoryError,
@@ -412,6 +414,8 @@ pub fn restore_from_snapshot(
             .clone_from(&vsock_override.uds_path);
     }
 
+    apply_pmem_overrides(&mut microvm_state, &params.pmem_overrides)?;
+
     let track_dirty_pages = params.track_dirty_pages;
 
     let vcpu_count = microvm_state
@@ -475,6 +479,38 @@ pub fn restore_from_snapshot(
     .map_err(RestoreFromSnapshotError::Build)
 }
 
+/// Applies the pmem backing file path overrides to the restored microVM state.
+///
+/// Each override identifies a pmem device by its id and replaces its
+/// `path_on_host`. Both MMIO and PCI transports are searched. An override
+/// targeting an id that does not match any pmem device returns an error.
+fn apply_pmem_overrides(
+    microvm_state: &mut MicrovmState,
+    overrides: &[PmemOverride],
+) -> Result<(), SnapshotStateFromFileError> {
+    for entry in overrides {
+        microvm_state
+            .device_states
+            .mmio_state
+            .pmem_devices
+            .iter_mut()
+            .map(|device| &mut device.device_state)
+            .chain(
+                microvm_state
+                    .device_states
+                    .pci_state
+                    .pmem_devices
+                    .iter_mut()
+                    .map(|device| &mut device.device_state),
+            )
+            .find(|state| state.config.id == entry.id)
+            .map(|state| state.config.path_on_host.clone_from(&entry.path_on_host))
+            .ok_or_else(|| SnapshotStateFromFileError::UnknownPmemDevice(entry.id.clone()))?;
+    }
+
+    Ok(())
+}
+
 /// Error type for [`snapshot_state_from_file`]
 #[derive(Debug, thiserror::Error, displaydoc::Display)]
 pub enum SnapshotStateFromFileError {
@@ -486,6 +522,8 @@ pub enum SnapshotStateFromFileError {
     UnknownNetworkDevice,
     /// Unknown Vsock Device.
     UnknownVsockDevice,
+    /// Unknown Pmem device: {0}
+    UnknownPmemDevice(String),
 }
 
 fn snapshot_state_from_file(
@@ -657,14 +695,15 @@ mod tests {
     use crate::builder::tests::insert_vmgenid_device;
     use crate::builder::tests::{
         CustomBlockConfig, default_kernel_cmdline, default_vmm, insert_balloon_device,
-        insert_block_devices, insert_net_device, insert_vsock_device,
+        insert_block_devices, insert_net_device, insert_pmem_devices, insert_vsock_device,
     };
     #[cfg(target_arch = "aarch64")]
     use crate::construct_kvm_mpidrs;
     use crate::devices::virtio::block::CacheType;
     use crate::snapshot::Persist;
     use crate::vmm_config::balloon::BalloonDeviceConfig;
     use crate::vmm_config::net::NetworkInterfaceConfig;
+    use crate::vmm_config::pmem::PmemConfig;
     use crate::vmm_config::vsock::tests::default_config;
     use crate::vstate::memory::{GuestMemoryRegionState, GuestRegionType};
 
@@ -717,6 +756,18 @@ mod tests {
 
         insert_vsock_device(&mut vmm, &mut cmdline, &mut event_manager, vsock_config);
 
+        // Add a pmem device. The returned temp file is dropped immediately; only
+        // the saved device state (which carries the configured id/path) is needed.
+        let pmem_config = PmemConfig {
+            id: String::from("pmem0"),
+            path_on_host: String::new(),
+            root_device: false,
+            read_only: false,
+            rate_limiter: None,
+        };
+        let _pmem_files =
+            insert_pmem_devices(&mut vmm, &mut cmdline, &mut event_manager, vec![pmem_config]);
+
         #[cfg(target_arch = "x86_64")]
         insert_vmgenid_device(&mut vmm);
         #[cfg(target_arch = "x86_64")]
@@ -765,6 +816,97 @@ mod tests {
         )
     }
 
+    fn microvm_state_with_devices() -> MicrovmState {
+        let vmm = default_vmm_with_devices();
+        let device_states = vmm.device_manager.save();
+
+        let vcpu_states = vec![VcpuState::default()];
+        #[cfg(target_arch = "aarch64")]
+        let mpidrs = construct_kvm_mpidrs(&vcpu_states);
+        MicrovmState {
+            device_states,
+            vcpu_states,
+            kvm_state: Default::default(),
+            vm_info: VmInfo {
+                mem_size_mib: 1u64,
+                ..Default::default()
+            },
+            #[cfg(target_arch = "aarch64")]
+            vm_state: vmm.vm.as_kvm().unwrap().save_state(&mpidrs).unwrap(),
+            #[cfg(target_arch = "x86_64")]
+            vm_state: vmm.vm.as_kvm().unwrap().save_state().unwrap(),
+        }
+    }
+
+    fn pmem_path(state: &MicrovmState, id: &str) -> Option<String> {
+        state
+            .device_states
+            .mmio_state
+            .pmem_devices
+            .iter()
+            .map(|device| &device.device_state)
+            .chain(
+                state
+                    .device_states
+                    .pci_state
+                    .pmem_devices
+                    .iter()
+                    .map(|device| &device.device_state),
+            )
+            .find(|s| s.config.id == id)
+            .map(|s| s.config.path_on_host.clone())
+    }
+
+    #[test]
+    fn test_apply_pmem_overrides_happy_path() {
+        let mut microvm_state = microvm_state_with_devices();
+        // Sanity: the device exists and its path differs from what we override to.
+        assert!(pmem_path(&microvm_state, "pmem0").is_some());
+
+        let overrides = vec![PmemOverride {
+            id: String::from("pmem0"),
+            path_on_host: String::from("/new/backing/file"),
+        }];
+
+        apply_pmem_overrides(&mut microvm_state, &overrides).unwrap();
+
+        assert_eq!(
+            pmem_path(&microvm_state, "pmem0").as_deref(),
+            Some("/new/backing/file")
+        );
+    }
+
+    #[test]
+    fn test_apply_pmem_overrides_unknown_device() {
+        let mut microvm_state = microvm_state_with_devices();
+
+        let overrides = vec![PmemOverride {
+            id: String::from("does-not-exist"),
+            path_on_host: String::from("/new/backing/file"),
+        }];
+
+        let err = apply_pmem_overrides(&mut microvm_state, &overrides).unwrap_err();
+        assert!(
+            matches!(
+                &err,
+                SnapshotStateFromFileError::UnknownPmemDevice(id) if id == "does-not-exist"
+            ),
+            "unexpected error: {err:?}"
+        );
+        // The existing device must be left untouched.
+        assert!(pmem_path(&microvm_state, "pmem0").is_some());
+    }
+
+    #[test]
+    fn test_apply_pmem_overrides_empty_is_noop() {
+        let mut microvm_state = microvm_state_with_devices();
+        let original = pmem_path(&microvm_state, "pmem0");
+
+        apply_pmem_overrides(&mut microvm_state, &[]).unwrap();
+
+        assert_eq!(pmem_path(&microvm_state, "pmem0"), original);
+    }
+
     #[test]
     fn test_create_guest_memory() {
         let mem_state = GuestMemoryState {