snapshot: add container_snapshot_path to load snapshot request

src/api_server/src/request/snapshot.rs

@@ -94,6 +94,7 @@ fn parse_put_snapshot_load(body: &Body) -> Result<ParsedRequest, Error> {
         mem_backend,
         enable_diff_snapshots: snapshot_config.enable_diff_snapshots,
         resume_vm: snapshot_config.resume_vm,
+        container_snapshot_path: snapshot_config.container_snapshot_path,
     };
 
     // Construct the `ParsedRequest` object.

src/api_server/swagger/firecracker.yaml

@@ -1188,6 +1188,7 @@ definitions:
       the two `mem_*` fields must be present in the body of the request.
     required:
       - snapshot_path
+      - container_snapshot_path
     properties:
       enable_diff_snapshots:
         type: boolean
@@ -1212,6 +1213,10 @@ definitions:
         type: boolean
         description:
           When set to true, the vm is also resumed if the snapshot load is successful.
+      container_snapshot_path:
+        type: string
+        description:
+          Path to the disk device backing the disk state at the time of the snapshot creation.
 
   TokenBucket:
     type: object

src/vmm/src/devices/virtio/block/persist.rs

@@ -95,7 +95,7 @@ pub struct BlockState {
     )]
     cache_type: CacheTypeState,
     root_device: bool,
-    disk_path: String,
+    pub disk_path: String,
     virtio_state: VirtioDeviceState,
     rate_limiter_state: RateLimiterState,
     #[version(start = 3)]

src/vmm/src/persist.rs

@@ -521,7 +521,25 @@ pub fn restore_from_snapshot(
     version_map: VersionMap,
     vm_resources: &mut VmResources,
 ) -> Result<Arc<Mutex<Vmm>>, RestoreFromSnapshotError> {
-    let microvm_state = snapshot_state_from_file(&params.snapshot_path, version_map)?;
+    let mut microvm_state = snapshot_state_from_file(&params.snapshot_path, version_map)?;
+
+    let container_snapshot_path = &params.container_snapshot_path;
+    // We assume that each microVM is backed by exactly one container image
+    // snapshot device (i.e., that no more than one container is run on each microVM).
+    assert_eq!(microvm_state.device_states.block_devices.len(), 2);
+    for i in 0..2 {
+        // We assume that one of the block devices is the rootfs, the other being the
+        // container image snapshot.
+        if microvm_state.device_states.block_devices[i]
+            .device_state
+            .disk_path
+            .contains("snap")
+        {
+            microvm_state.device_states.block_devices[i]
+                .device_state
+                .disk_path = container_snapshot_path.clone();
+        }
+    }
 
     // Some sanity checks before building the microvm.
     snapshot_state_sanity_check(&microvm_state)?;

src/vmm/src/vmm_config/snapshot.rs

@@ -63,6 +63,8 @@ pub struct LoadSnapshotParams {
     /// When set to true, the vm is also resumed if the snapshot load
     /// is successful.
     pub resume_vm: bool,
+    /// Path to the disk device backing the container snapshot.
+    pub container_snapshot_path: String,
 }
 
 /// Stores the configuration for loading a snapshot that is provided by the user.
@@ -85,6 +87,8 @@ pub struct LoadSnapshotConfig {
     /// Whether or not to resume the vm post snapshot load.
     #[serde(default)]
     pub resume_vm: bool,
+    /// Path to the disk device backing the container snapshot.
+    pub container_snapshot_path: String,
 }
 
 /// Stores the configuration used for managing snapshot memory.