From e5faf365d159e6d2202f04ded9d22ec0b59a2f69 Mon Sep 17 00:00:00 2001
From: Georgiy Lebedev <lebedev.gk@phystech.edu>
Date: Thu, 24 Aug 2023 16:52:31 +0600
Subject: [PATCH] snapshot: add `container_snapshot_path` to load snapshot
 request

When a snapshot of a VM created by firecracker-containerd is restored,
due to the non-deterministic container snapshot path (it depends
on the containerd snapshotter implementation), the container snapshot
path at the time of the snapshot creation is different from the
container snapshot path at the time of the snapshot loading.

Firecracker does not support renaming resources at snapshot-restore, so
as a workaround we manually substitute the VM state with the path of the
block device backing the container snapshot to the path of the new
container snapshot path received from the LoadSnapshot request.

Closes #4014

Signed-off-by: Georgiy Lebedev <lebedev.gk@phystech.edu>
---
 src/api_server/src/request/snapshot.rs      |  1 +
 src/api_server/swagger/firecracker.yaml     |  5 +++++
 src/vmm/src/devices/virtio/block/persist.rs |  2 +-
 src/vmm/src/persist.rs                      | 20 +++++++++++++++++++-
 src/vmm/src/vmm_config/snapshot.rs          |  4 ++++
 5 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/api_server/src/request/snapshot.rs b/src/api_server/src/request/snapshot.rs
index 221b975178d..9e825061484 100644
--- a/src/api_server/src/request/snapshot.rs
+++ b/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.
diff --git a/src/api_server/swagger/firecracker.yaml b/src/api_server/swagger/firecracker.yaml
index 7a4e2151d91..392556585d4 100644
--- a/src/api_server/swagger/firecracker.yaml
+++ b/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
diff --git a/src/vmm/src/devices/virtio/block/persist.rs b/src/vmm/src/devices/virtio/block/persist.rs
index bb674379a01..533b178a214 100644
--- a/src/vmm/src/devices/virtio/block/persist.rs
+++ b/src/vmm/src/devices/virtio/block/persist.rs
@@ -86,7 +86,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)]
diff --git a/src/vmm/src/persist.rs b/src/vmm/src/persist.rs
index abf649e3a75..6f82b8e4e17 100644
--- a/src/vmm/src/persist.rs
+++ b/src/vmm/src/persist.rs
@@ -523,7 +523,25 @@ pub fn restore_from_snapshot(
     version_map: VersionMap,
     vm_resources: &mut VmResources,
 ) -> std::result::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)?;
diff --git a/src/vmm/src/vmm_config/snapshot.rs b/src/vmm/src/vmm_config/snapshot.rs
index d7d73345f40..078724138d4 100644
--- a/src/vmm/src/vmm_config/snapshot.rs
+++ b/src/vmm/src/vmm_config/snapshot.rs
@@ -61,6 +61,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.
@@ -83,6 +85,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.