From fde89896d262f98174a9c73129a573758b71e771 Mon Sep 17 00:00:00 2001
From: Egor Lazarchuk <yegorlz@amazon.co.uk>
Date: Thu, 30 Nov 2023 16:32:28 +0000
Subject: [PATCH] refactor(block): readding skip of vhost-user-block from
 snapshots

Adding back the skip of vhost-user-block devices.

Signed-off-by: Egor Lazarchuk <yegorlz@amazon.co.uk>
---
 src/vmm/src/device_manager/persist.rs | 39 ++++++++++++++++-----------
 src/vmm/src/persist.rs                |  2 +-
 src/vmm/tests/integration_tests.rs    |  2 +-
 3 files changed, 25 insertions(+), 18 deletions(-)

diff --git a/src/vmm/src/device_manager/persist.rs b/src/vmm/src/device_manager/persist.rs
index 7be444d0ab97..2aef7f9060fe 100644
--- a/src/vmm/src/device_manager/persist.rs
+++ b/src/vmm/src/device_manager/persist.rs
@@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex};
 
 use event_manager::{MutEventSubscriber, SubscriberOps};
 use kvm_ioctls::VmFd;
-use log::error;
+use log::{error, warn};
 use snapshot::Persist;
 use versionize::{VersionMap, Versionize, VersionizeError, VersionizeResult};
 use versionize_derive::Versionize;
@@ -175,8 +175,8 @@ pub struct DeviceStates {
     #[cfg(target_arch = "aarch64")]
     // State of legacy devices in MMIO space.
     pub legacy_devices: Vec<ConnectedLegacyState>,
-    /// Virtio block device states.
-    pub virtio_block_devices: Vec<ConnectedBlockState>,
+    /// Block device states.
+    pub block_devices: Vec<ConnectedBlockState>,
     /// Net device states.
     pub net_devices: Vec<ConnectedNetState>,
     /// Vsock device state.
@@ -274,13 +274,20 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                 // Both virtio-block and vhost-user-block share same device type.
                 TYPE_BLOCK => {
                     let block = locked_device.as_mut_any().downcast_mut::<Block>().unwrap();
-                    block.prepare_save();
-                    states.virtio_block_devices.push(ConnectedBlockState {
-                        device_id: devid.clone(),
-                        device_state: block.save(),
-                        transport_state,
-                        device_info: device_info.clone(),
-                    })
+                    if block.is_vhost_user() {
+                        warn!(
+                            "Skipping vhost-user-block device. VhostUserBlock does not support \
+                             snapshotting yet"
+                        );
+                    } else {
+                        block.prepare_save();
+                        states.block_devices.push(ConnectedBlockState {
+                            device_id: devid.clone(),
+                            device_state: block.save(),
+                            transport_state,
+                            device_info: device_info.clone(),
+                        })
+                    }
                 }
                 TYPE_NET => {
                     let net = locked_device.as_any().downcast_ref::<Net>().unwrap();
@@ -470,10 +477,10 @@ impl<'a> Persist<'a> for MMIODeviceManager {
             )?;
         }
 
-        for virtio_block_state in &state.virtio_block_devices {
+        for block_state in &state.block_devices {
             let device = Arc::new(Mutex::new(Block::restore(
                 BlockConstructorArgs { mem: mem.clone() },
-                &virtio_block_state.device_state,
+                &block_state.device_state,
             )?));
 
             (constructor_args.for_each_restored_device)(
@@ -485,9 +492,9 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                 device.clone(),
                 false,
                 device,
-                &virtio_block_state.device_id,
-                &virtio_block_state.transport_state,
-                &virtio_block_state.device_info,
+                &block_state.device_id,
+                &block_state.transport_state,
+                &block_state.device_info,
                 constructor_args.event_manager,
             )?;
         }
@@ -640,7 +647,7 @@ mod tests {
     impl PartialEq for DeviceStates {
         fn eq(&self, other: &DeviceStates) -> bool {
             self.balloon_device == other.balloon_device
-                && self.virtio_block_devices == other.virtio_block_devices
+                && self.block_devices == other.block_devices
                 && self.net_devices == other.net_devices
                 && self.vsock_device == other.vsock_device
         }
diff --git a/src/vmm/src/persist.rs b/src/vmm/src/persist.rs
index cd799f30b707..6635131245b3 100644
--- a/src/vmm/src/persist.rs
+++ b/src/vmm/src/persist.rs
@@ -650,7 +650,7 @@ mod tests {
 
         // Only checking that all devices are saved, actual device state
         // is tested by that device's tests.
-        assert_eq!(states.virtio_block_devices.len(), 1);
+        assert_eq!(states.block_devices.len(), 1);
         assert_eq!(states.net_devices.len(), 1);
         assert!(states.vsock_device.is_some());
         assert!(states.balloon_device.is_some());
diff --git a/src/vmm/tests/integration_tests.rs b/src/vmm/tests/integration_tests.rs
index f06836d43bef..8db755b66d87 100644
--- a/src/vmm/tests/integration_tests.rs
+++ b/src/vmm/tests/integration_tests.rs
@@ -226,7 +226,7 @@ fn verify_create_snapshot(is_diff: bool) -> (TempFile, TempFile) {
     assert_eq!(
         restored_microvm_state
             .device_states
-            .virtio_block_devices
+            .block_devices
             .len(),
         0
     );