From 0dc6a29c7e57797d4cffd05f80caf2a543e8c4a8 Mon Sep 17 00:00:00 2001
From: Georgiy Lebedev <lebedev.gk@phystech.edu>
Date: Mon, 7 Aug 2023 14:36: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 3563c85f9b9..cffa2e7cb0d 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 4cfb6126561..52c2d65d48c 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 9c9253bdc03..fbe36f2c2ef 100644
--- a/src/vmm/src/devices/virtio/block/persist.rs
+++ b/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)]
diff --git a/src/vmm/src/persist.rs b/src/vmm/src/persist.rs
index 3624ffd1fe1..8197b164697 100644
--- a/src/vmm/src/persist.rs
+++ b/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)?;
diff --git a/src/vmm/src/vmm_config/snapshot.rs b/src/vmm/src/vmm_config/snapshot.rs
index 28b111b9f7d..8a521d1fb45 100644
--- a/src/vmm/src/vmm_config/snapshot.rs
+++ b/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.