Block refactor

src/vmm/src/builder.rs

@@ -43,6 +43,7 @@ use crate::devices::legacy::serial::SerialOut;
 use crate::devices::legacy::RTCDevice;
 use crate::devices::legacy::{EventFdTrigger, SerialEventsWrapper, SerialWrapper};
 use crate::devices::virtio::balloon::Balloon;
+use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::mmio::MmioTransport;
 use crate::devices::virtio::net::Net;
@@ -54,7 +55,6 @@ use crate::persist::{MicrovmState, MicrovmStateError};
 use crate::resources::VmResources;
 use crate::snapshot::Persist;
 use crate::vmm_config::boot_source::BootConfig;
-use crate::vmm_config::drive::BlockDeviceType;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::{VmConfig, VmConfigError};
 use crate::vstate::memory::{GuestAddress, GuestMemory, GuestMemoryExtension, GuestMemoryMmap};
@@ -840,51 +840,38 @@ fn attach_entropy_device(
     )
 }
 
-fn attach_block_devices<'a, I: Iterator<Item = &'a BlockDeviceType> + Debug>(
+fn attach_block_devices<'a, I: Iterator<Item = &'a Arc<Mutex<Block>>> + Debug>(
     vmm: &mut Vmm,
     cmdline: &mut LoaderKernelCmdline,
     blocks: I,
     event_manager: &mut EventManager,
 ) -> Result<(), StartMicrovmError> {
-    macro_rules! attach_block {
-        ($block:ident, $is_vhost_user:ident) => {
-            let id = {
-                let locked = $block.lock().expect("Poisoned lock");
-                if locked.root_device {
-                    match locked.partuuid {
-                        Some(ref partuuid) => {
-                            cmdline.insert_str(format!("root=PARTUUID={}", partuuid))?
-                        }
-                        None => cmdline.insert_str("root=/dev/vda")?,
-                    }
-                    match locked.read_only {
-                        true => cmdline.insert_str("ro")?,
-                        false => cmdline.insert_str("rw")?,
+    for block in blocks {
+        let (id, is_vhost_user) = {
+            let locked = block.lock().expect("Poisoned lock");
+            if locked.root_device() {
+                match locked.partuuid() {
+                    Some(ref partuuid) => {
+                        cmdline.insert_str(format!("root=PARTUUID={}", partuuid))?
                     }
+                    None => cmdline.insert_str("root=/dev/vda")?,
+                }
+                match locked.read_only() {
+                    true => cmdline.insert_str("ro")?,
+                    false => cmdline.insert_str("rw")?,
                 }
-                locked.id.clone()
-            };
-            // The device mutex mustn't be locked here otherwise it will deadlock.
-            attach_virtio_device(
-                event_manager,
-                vmm,
-                id,
-                $block.clone(),
-                cmdline,
-                $is_vhost_user,
-            )?;
-        };
-    }
-
-    for block in blocks {
-        match block {
-            BlockDeviceType::VirtioBlock(block) => {
-                attach_block!(block, false);
-            }
-            BlockDeviceType::VhostUserBlock(block) => {
-                attach_block!(block, true);
             }
+            (locked.id().to_string(), locked.is_vhost_user())
         };
+        // The device mutex mustn't be locked here otherwise it will deadlock.
+        attach_virtio_device(
+            event_manager,
+            vmm,
+            id,
+            block.clone(),
+            cmdline,
+            is_vhost_user,
+        )?;
     }
     Ok(())
 }
@@ -948,7 +935,7 @@ pub mod tests {
 
     use super::*;
     use crate::arch::DeviceType;
-    use crate::devices::virtio::block_common::CacheType;
+    use crate::devices::virtio::block::CacheType;
     use crate::devices::virtio::rng::device::ENTROPY_DEV_ID;
     use crate::devices::virtio::vsock::{TYPE_VSOCK, VSOCK_DEV_ID};
     use crate::devices::virtio::{TYPE_BALLOON, TYPE_BLOCK, TYPE_RNG};

src/vmm/src/device_manager/mmio.rs

@@ -23,11 +23,11 @@
 use crate::devices::legacy::RTCDevice;
 use crate::devices::pseudo::BootTimer;
 use crate::devices::virtio::balloon::Balloon;
+use crate::devices::virtio::block::device::Block;
 use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::mmio::MmioTransport;
 use crate::devices::virtio::net::Net;
 use crate::devices::virtio::rng::Entropy;
-use crate::devices::virtio::virtio_block::VirtioBlock;
 use crate::devices::virtio::vsock::TYPE_VSOCK;
 use crate::devices::virtio::{TYPE_BALLOON, TYPE_BLOCK, TYPE_NET, TYPE_RNG};
 use crate::devices::BusDevice;
@@ -400,7 +400,7 @@
                     TYPE_BLOCK => {
                         // 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 let Some(block) = virtio.as_mut_any().downcast_mut::<Block>() {
                             // 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

src/vmm/src/device_manager/persist.rs

@@ -17,6 +17,9 @@
 use crate::arch::DeviceType;
 use crate::devices::virtio::balloon::persist::{BalloonConstructorArgs, BalloonState};
 use crate::devices::virtio::balloon::{Balloon, BalloonError};
+use crate::devices::virtio::block::device::Block;
+use crate::devices::virtio::block::persist::{BlockConstructorArgs, BlockState};
+use crate::devices::virtio::block::BlockError;
 use crate::devices::virtio::device::VirtioDevice;
 use crate::devices::virtio::mmio::MmioTransport;
 use crate::devices::virtio::net::persist::{
@@ -28,13 +31,6 @@
     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::vhost_user_block::VhostUserBlockError;
-use crate::devices::virtio::virtio_block::persist::{VirtioBlockConstructorArgs, VirtioBlockState};
-use crate::devices::virtio::virtio_block::{VirtioBlock, VirtioBlockError};
 use crate::devices::virtio::vsock::persist::{
     VsockConstructorArgs, VsockState, VsockUdsConstructorArgs,
 };
@@ -54,10 +50,8 @@
 pub enum DevicePersistError {
     /// Balloon: {0}
     Balloon(#[from] BalloonError),
-    /// VirtioBlock: {0}
-    VirtioBlock(#[from] VirtioBlockError),
-    /// VhostUserBlock: {0}
-    VhostUserBlock(#[from] VhostUserBlockError),
+    /// Block: {0}
+    Block(#[from] BlockError),
     /// Device manager: {0}
     DeviceManager(#[from] super::mmio::MmioError),
     /// Mmio transport
@@ -92,24 +86,11 @@
 
 /// Holds the state of a virtio block device connected to the MMIO space.
 #[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ConnectedVirtioBlockState {
-    /// Device identifier.
-    pub device_id: String,
-    /// Device state.
-    pub device_state: VirtioBlockState,
-    /// Mmio transport state.
-    pub transport_state: MmioTransportState,
-    /// VmmResources.
-    pub device_info: MMIODeviceInfo,
-}
-
-/// Holds the state of a vhost-user block device connected to the MMIO space.
-#[derive(Debug, Clone, Serialize, Deserialize)]
-pub struct ConnectedVhostUserBlockState {
+pub struct ConnectedBlockState {
     /// Device identifier.
     pub device_id: String,
     /// Device state.
-    pub device_state: VhostUserBlockState,
+    pub device_state: BlockState,
     /// Mmio transport state.
     pub transport_state: MmioTransportState,
     /// VmmResources.
@@ -196,10 +177,8 @@
     #[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<ConnectedVirtioBlockState>,
-    /// Vhost-user block device states.
-    pub vhost_user_block_devices: Vec<ConnectedVhostUserBlockState>,
+    /// Block device states.
+    pub block_devices: Vec<ConnectedBlockState>,
     /// Net device states.
     pub net_devices: Vec<ConnectedNetState>,
     /// Vsock device state.
@@ -216,8 +195,7 @@
 /// types when restoring from a snapshot.
 #[derive(Debug)]
 pub enum SharedDeviceType {
-    VirtioBlock(Arc<Mutex<VirtioBlock>>),
-    VhostUserBlock(Arc<Mutex<VhostUserBlock>>),
+    VirtioBlock(Arc<Mutex<Block>>),
     Network(Arc<Mutex<Net>>),
     Balloon(Arc<Mutex<Balloon>>),
     Vsock(Arc<Mutex<Vsock<VsockUnixBackend>>>),
@@ -294,19 +272,20 @@
                 }
                 // Both virtio-block and vhost-user-block share same device type.
                 TYPE_BLOCK => {
-                    if let Some(block) = locked_device.as_mut_any().downcast_mut::<VirtioBlock>() {
+                    let block = locked_device.as_mut_any().downcast_mut::<Block>().unwrap();
+                    if block.is_vhost_user() {
+                        warn!(
+                            "Skipping vhost-user-block device. VhostUserBlock does not support \
+                             snapshotting yet"
+                        );
+                    } else {
                         block.prepare_save();
-                        states.virtio_block_devices.push(ConnectedVirtioBlockState {
+                        states.block_devices.push(ConnectedBlockState {
                             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 => {
@@ -497,10 +476,10 @@
             )?;
         }
 
-        for virtio_block_state in &state.virtio_block_devices {
-            let device = Arc::new(Mutex::new(VirtioBlock::restore(
-                VirtioBlockConstructorArgs { mem: mem.clone() },
-                &virtio_block_state.device_state,
+        for block_state in &state.block_devices {
+            let device = Arc::new(Mutex::new(Block::restore(
+                BlockConstructorArgs { mem: mem.clone() },
+                &block_state.device_state,
             )?));
 
             (constructor_args.for_each_restored_device)(
@@ -512,31 +491,9 @@
                 device.clone(),
                 false,
                 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::VhostUserBlock(device.clone()),
-            );
-
-            restore_helper(
-                device.clone(),
-                true,
-                device,
-                &vhost_user_block_state.device_id,
-                &vhost_user_block_state.transport_state,
-                &vhost_user_block_state.device_info,
+                &block_state.device_id,
+                &block_state.transport_state,
+                &block_state.device_info,
                 constructor_args.event_manager,
             )?;
         }
@@ -650,7 +607,7 @@
 
     use super::*;
     use crate::builder::tests::*;
-    use crate::devices::virtio::block_common::CacheType;
+    use crate::devices::virtio::block::CacheType;
     use crate::resources::VmmConfig;
     use crate::snapshot::Snapshot;
     use crate::vmm_config::balloon::BalloonDeviceConfig;
@@ -665,8 +622,8 @@
         }
     }
 
-    impl PartialEq for ConnectedVirtioBlockState {
-        fn eq(&self, other: &ConnectedVirtioBlockState) -> bool {
+    impl PartialEq for ConnectedBlockState {
+        fn eq(&self, other: &ConnectedBlockState) -> bool {
             // Actual device state equality is checked by the device's tests.
             self.transport_state == other.transport_state && self.device_info == other.device_info
         }
@@ -689,7 +646,7 @@
     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
         }