feat: added vhost-user to device manager save/restore paths

ShadowCurse · ShadowCurse · commit 6f9c738d3727 · 2023-11-08T10:23:36.000Z
Even though we don't support creation of snapshots with
vhost-user-block device we add it to the restoration path
for future use.

Signed-off-by: Egor Lazarchuk &lt;yegorlz@amazon.co.uk&gt;
diff --git a/src/vmm/src/device_manager/mmio.rs b/src/vmm/src/device_manager/mmio.rs
@@ -396,14 +396,18 @@ impl MMIODeviceManager {
                         }
                     }
                     TYPE_BLOCK => {
-                        let block = virtio.as_mut_any().downcast_mut::<VirtioBlock>().unwrap();
-                        // If device is activated, kick the block queue(s) to make up for any
-                        // pending or in-flight epoll events we may have not captured in snapshot.
-                        // No need to kick Ratelimiters because they are restored 'unblocked' so
-                        // any inflight `timer_fd` events can be safely discarded.
-                        if block.is_activated() {
-                            info!("kick block {}.", id);
-                            block.process_virtio_queues();
+                        // We only care about kicking virtio block.
+                        // If we need to kick vhost-user-block we can do nothing.
+                        if let Some(block) = virtio.as_mut_any().downcast_mut::<VirtioBlock>() {
+                            // If device is activated, kick the block queue(s) to make up for any
+                            // pending or in-flight epoll events we may have not captured in
+                            // snapshot. No need to kick Ratelimiters
+                            // because they are restored 'unblocked' so
+                            // any inflight `timer_fd` events can be safely discarded.
+                            if block.is_activated() {
+                                info!("kick block {}.", id);
+                                block.process_virtio_queues();
+                            }
                         }
                     }
                     TYPE_NET => {
diff --git 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;
@@ -30,12 +30,17 @@ use crate::devices::virtio::rng::persist::{
     EntropyConstructorArgs, EntropyPersistError as EntropyError, EntropyState,
 };
 use crate::devices::virtio::rng::Entropy;
+use crate::devices::virtio::vhost_user_block::device::VhostUserBlock;
+use crate::devices::virtio::vhost_user_block::persist::{
+    VhostUserBlockConstructorArgs, VhostUserBlockState,
+};
 use crate::devices::virtio::vsock::persist::{
     VsockConstructorArgs, VsockState, VsockUdsConstructorArgs,
 };
 use crate::devices::virtio::vsock::{Vsock, VsockError, VsockUnixBackend, VsockUnixBackendError};
 use crate::devices::virtio::{
-    MmioTransport, VirtioDevice, TYPE_BALLOON, TYPE_BLOCK, TYPE_NET, TYPE_RNG, TYPE_VSOCK,
+    MmioTransport, VhostUserBlockError, VirtioDevice, TYPE_BALLOON, TYPE_BLOCK, TYPE_NET, TYPE_RNG,
+    TYPE_VSOCK,
 };
 #[cfg(target_arch = "aarch64")]
 use crate::logger;
@@ -49,7 +54,8 @@ use crate::EventManager;
 #[derive(Debug, derive_more::From)]
 pub enum DevicePersistError {
     Balloon(BalloonError),
-    Block(VirtioBlockError),
+    VirtioBlock(VirtioBlockError),
+    VhostUserBlock(VhostUserBlockError),
     DeviceManager(super::mmio::MmioError),
     MmioTransport,
     #[cfg(target_arch = "aarch64")]
@@ -75,10 +81,10 @@ pub struct ConnectedBalloonState {
     pub device_info: MMIODeviceInfo,
 }
 
-/// Holds the state of a block device connected to the MMIO space.
+/// Holds the state of a virtio block device connected to the MMIO space.
 // NOTICE: Any changes to this structure require a snapshot version bump.
 #[derive(Debug, Clone, Versionize)]
-pub struct ConnectedBlockState {
+pub struct ConnectedVirtioBlockState {
     /// Device identifier.
     pub device_id: String,
     /// Device state.
@@ -89,6 +95,20 @@ pub struct ConnectedBlockState {
     pub device_info: MMIODeviceInfo,
 }
 
+/// Holds the state of a vhost-user block device connected to the MMIO space.
+// NOTICE: Any changes to this structure require a snapshot version bump.
+#[derive(Debug, Clone, Versionize)]
+pub struct ConnectedVhostUserBlockState {
+    /// Device identifier.
+    pub device_id: String,
+    /// Device state.
+    pub device_state: VhostUserBlockState,
+    /// Mmio transport state.
+    pub transport_state: MmioTransportState,
+    /// VmmResources.
+    pub device_info: MMIODeviceInfo,
+}
+
 /// Holds the state of a net device connected to the MMIO space.
 // NOTICE: Any changes to this structure require a snapshot version bump.
 #[derive(Debug, Clone, Versionize)]
@@ -174,8 +194,11 @@ pub struct DeviceStates {
     #[cfg(target_arch = "aarch64")]
     // State of legacy devices in MMIO space.
     pub legacy_devices: Vec<ConnectedLegacyState>,
-    /// Block device states.
-    pub block_devices: Vec<ConnectedBlockState>,
+    /// Virtio block device states.
+    pub virtio_block_devices: Vec<ConnectedVirtioBlockState>,
+    /// Vhost-user block device states.
+    #[version(start = 5, de_fn = "de_vhost_user_block")]
+    pub vhost_user_block_devices: Vec<ConnectedVhostUserBlockState>,
     /// Net device states.
     pub net_devices: Vec<ConnectedNetState>,
     /// Vsock device state.
@@ -195,13 +218,23 @@ pub struct DeviceStates {
 /// from a snapshot.
 #[derive(Debug)]
 pub enum SharedDeviceType {
-    Block(Arc<Mutex<VirtioBlock>>),
+    VirtioBlock(Arc<Mutex<VirtioBlock>>),
+    VhostUserBlock(Arc<Mutex<VhostUserBlock>>),
     Network(Arc<Mutex<Net>>),
     Balloon(Arc<Mutex<Balloon>>),
     Vsock(Arc<Mutex<Vsock<VsockUnixBackend>>>),
     Entropy(Arc<Mutex<Entropy>>),
 }
 
+impl DeviceStates {
+    fn de_vhost_user_block(&mut self, source_version: u16) -> VersionizeResult<()> {
+        if source_version < 5 {
+            self.vhost_user_block_devices = vec![];
+        }
+
+        Ok(())
+    }
+}
 pub struct MMIODevManagerConstructorArgs<'a> {
     pub mem: GuestMemoryMmap,
     pub vm: &'a VmFd,
@@ -229,16 +262,7 @@ impl<'a> Persist<'a> for MMIODeviceManager {
     type Error = DevicePersistError;
 
     fn save(&self) -> Self::State {
-        let mut states = DeviceStates {
-            balloon_device: None,
-            block_devices: Vec::new(),
-            net_devices: Vec::new(),
-            vsock_device: None,
-            #[cfg(target_arch = "aarch64")]
-            legacy_devices: Vec::new(),
-            mmds_version: None,
-            entropy_device: None,
-        };
+        let mut states = DeviceStates::default();
         let _: Result<(), ()> = self.for_each_device(|devtype, devid, device_info, bus_dev| {
             if *devtype == crate::arch::DeviceType::BootTimer {
                 // No need to save BootTimer state.
@@ -279,18 +303,22 @@ impl<'a> Persist<'a> for MMIODeviceManager {
                         device_info: device_info.clone(),
                     });
                 }
+                // Both virtio-block and vhost-user-block share same device type.
                 TYPE_BLOCK => {
-                    let block = locked_device
-                        .as_mut_any()
-                        .downcast_mut::<VirtioBlock>()
-                        .unwrap();
-                    block.prepare_save();
-                    states.block_devices.push(ConnectedBlockState {
-                        device_id: devid.clone(),
-                        device_state: block.save(),
-                        transport_state,
-                        device_info: device_info.clone(),
-                    });
+                    if let Some(block) = locked_device.as_mut_any().downcast_mut::<VirtioBlock>() {
+                        block.prepare_save();
+                        states.virtio_block_devices.push(ConnectedVirtioBlockState {
+                            device_id: devid.clone(),
+                            device_state: block.save(),
+                            transport_state,
+                            device_info: device_info.clone(),
+                        })
+                    } else {
+                        warn!(
+                            "Skipping vhost-user-block device. VhostUserBlock does not support \
+                             snapshotting yet"
+                        );
+                    }
                 }
                 TYPE_NET => {
                     let net = locked_device.as_any().downcast_ref::<Net>().unwrap();
@@ -477,23 +505,44 @@ impl<'a> Persist<'a> for MMIODeviceManager {
             )?;
         }
 
-        for block_state in &state.block_devices {
+        for virtio_block_state in &state.virtio_block_devices {
             let device = Arc::new(Mutex::new(VirtioBlock::restore(
                 VirtioBlockConstructorArgs { mem: mem.clone() },
-                &block_state.device_state,
+                &virtio_block_state.device_state,
+            )?));
+
+            (constructor_args.for_each_restored_device)(
+                constructor_args.vm_resources,
+                SharedDeviceType::VirtioBlock(device.clone()),
+            );
+
+            restore_helper(
+                device.clone(),
+                device,
+                &virtio_block_state.device_id,
+                &virtio_block_state.transport_state,
+                &virtio_block_state.device_info,
+                constructor_args.event_manager,
+            )?;
+        }
+
+        for vhost_user_block_state in &state.vhost_user_block_devices {
+            let device = Arc::new(Mutex::new(VhostUserBlock::restore(
+                VhostUserBlockConstructorArgs { mem: mem.clone() },
+                &vhost_user_block_state.device_state,
             )?));
 
             (constructor_args.for_each_restored_device)(
                 constructor_args.vm_resources,
-                SharedDeviceType::Block(device.clone()),
+                SharedDeviceType::VhostUserBlock(device.clone()),
             );
 
             restore_helper(
                 device.clone(),
                 device,
-                &block_state.device_id,
-                &block_state.transport_state,
-                &block_state.device_info,
+                &vhost_user_block_state.device_id,
+                &vhost_user_block_state.transport_state,
+                &vhost_user_block_state.device_info,
                 constructor_args.event_manager,
             )?;
         }
@@ -619,8 +668,8 @@ mod tests {
         }
     }
 
-    impl PartialEq for ConnectedBlockState {
-        fn eq(&self, other: &ConnectedBlockState) -> bool {
+    impl PartialEq for ConnectedVirtioBlockState {
+        fn eq(&self, other: &ConnectedVirtioBlockState) -> bool {
             // Actual device state equality is checked by the device's tests.
             self.transport_state == other.transport_state && self.device_info == other.device_info
         }
@@ -643,7 +692,7 @@ mod tests {
     impl PartialEq for DeviceStates {
         fn eq(&self, other: &DeviceStates) -> bool {
             self.balloon_device == other.balloon_device
-                && self.block_devices == other.block_devices
+                && self.virtio_block_devices == other.virtio_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
@@ -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.block_devices.len(), 1);
+        assert_eq!(states.virtio_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/src/resources.rs b/src/vmm/src/resources.rs
@@ -204,8 +204,12 @@ impl VmResources {
     /// restoring from a snapshot).
     pub fn update_from_restored_device(&mut self, device: SharedDeviceType) {
         match device {
-            SharedDeviceType::Block(block) => {
-                self.block.add_device(block);
+            SharedDeviceType::VirtioBlock(block) => {
+                self.block.add_virtio_device(block);
+            }
+
+            SharedDeviceType::VhostUserBlock(block) => {
+                self.block.add_vhost_user_device(block);
             }
 
             SharedDeviceType::Network(network) => {
diff --git a/src/vmm/src/version_map.rs b/src/vmm/src/version_map.rs
@@ -86,6 +86,7 @@ lazy_static! {
         version_map.new_version();
 
         version_map.set_type_version(VirtioDeviceState::type_id(), 2);
+        version_map.set_type_version(DeviceStates::type_id(), 5);
 
         version_map
     };
diff --git a/src/vmm/src/vmm_config/drive.rs b/src/vmm/src/vmm_config/drive.rs
@@ -125,7 +125,7 @@ impl BlockBuilder {
     }
 
     /// Inserts an existing block device.
-    pub fn add_device(&mut self, block_device: Arc<Mutex<VirtioBlock>>) {
+    pub fn add_virtio_device(&mut self, block_device: Arc<Mutex<VirtioBlock>>) {
         if block_device.lock().expect("Poisoned lock").is_root_device() {
             self.devices
                 .push_front(BlockDeviceType::VirtioBlock(block_device));
@@ -135,6 +135,17 @@ impl BlockBuilder {
         }
     }
 
+    /// Inserts an existing block device.
+    pub fn add_vhost_user_device(&mut self, block_device: Arc<Mutex<VhostUserBlock>>) {
+        if block_device.lock().expect("Poisoned lock").root_device {
+            self.devices
+                .push_front(BlockDeviceType::VhostUserBlock(block_device));
+        } else {
+            self.devices
+                .push_back(BlockDeviceType::VhostUserBlock(block_device));
+        }
+    }
+
     /// Inserts a `Block` in the block devices list using the specified configuration.
     /// If a block with the same id already exists, it will overwrite it.
     /// Inserting a secondary root block device will fail.
@@ -633,7 +644,7 @@ mod tests {
 
         let block = VirtioBlock::new(config).unwrap();
 
-        block_devs.add_device(Arc::new(Mutex::new(block)));
+        block_devs.add_virtio_device(Arc::new(Mutex::new(block)));
         assert_eq!(block_devs.devices.len(), 1);
         match block_devs.devices.pop_back().unwrap() {
             BlockDeviceType::VirtioBlock(ref b) => assert_eq!(b.lock().unwrap().id(), block_id),
diff --git a/src/vmm/tests/integration_tests.rs b/src/vmm/tests/integration_tests.rs
@@ -223,7 +223,13 @@ fn verify_create_snapshot(is_diff: bool) -> (TempFile, TempFile) {
 
     // Verify deserialized data.
     // The default vmm has no devices and one vCPU.
-    assert_eq!(restored_microvm_state.device_states.block_devices.len(), 0);
+    assert_eq!(
+        restored_microvm_state
+            .device_states
+            .virtio_block_devices
+            .len(),
+        0
+    );
     assert_eq!(restored_microvm_state.device_states.net_devices.len(), 0);
     assert!(restored_microvm_state.device_states.vsock_device.is_none());
     assert_eq!(restored_microvm_state.vcpu_states.len(), 1);
