refactor: remove SharedDeviceType

ShadowCurse · ShadowCurse · commit b7d041f1e000 · 2025-10-22T12:35:28.000+01:00
This wrapper was only used in the `update_from_restored_device`
function which was instantly unwrapping the value and
just adding device to the VmResources. Instead of this
round trip just add devices directly since fields are
public.
This also removes the check for the addition of balloon device,
but this check is not needed anyway, since snapshot should not
contain conflicting configs.

Signed-off-by: Egor Lazarchuk &lt;yegorlz@amazon.co.uk&gt;
diff --git a/src/vmm/src/device_manager/pci_mngr.rs b/src/vmm/src/device_manager/pci_mngr.rs
@@ -10,7 +10,7 @@ use event_manager::{MutEventSubscriber, SubscriberOps};
 use log::{debug, error, warn};
 use serde::{Deserialize, Serialize};
 
-use super::persist::{MmdsState, SharedDeviceType};
+use super::persist::MmdsState;
 use crate::devices::pci::PciSegment;
 use crate::devices::virtio::balloon::Balloon;
 use crate::devices::virtio::balloon::persist::{BalloonConstructorArgs, BalloonState};
@@ -431,8 +431,8 @@ impl<'a> Persist<'a> for PciDevices {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Balloon(device.clone()))
-                .unwrap();
+                .balloon
+                .set_device(device.clone());
 
             pci_devices
                 .restore_pci_device(
@@ -456,8 +456,8 @@ impl<'a> Persist<'a> for PciDevices {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::VirtioBlock(device.clone()))
-                .unwrap();
+                .block
+                .add_virtio_device(device.clone());
 
             pci_devices
                 .restore_pci_device(
@@ -506,8 +506,8 @@ impl<'a> Persist<'a> for PciDevices {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Network(device.clone()))
-                .unwrap();
+                .net_builder
+                .add_device(device.clone());
 
             pci_devices
                 .restore_pci_device(
@@ -539,8 +539,8 @@ impl<'a> Persist<'a> for PciDevices {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Vsock(device.clone()))
-                .unwrap();
+                .vsock
+                .set_device(device.clone());
 
             pci_devices
                 .restore_pci_device(
@@ -562,8 +562,8 @@ impl<'a> Persist<'a> for PciDevices {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Entropy(device.clone()))
-                .unwrap();
+                .entropy
+                .set_device(device.clone());
 
             pci_devices
                 .restore_pci_device(
@@ -590,8 +590,8 @@ impl<'a> Persist<'a> for PciDevices {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Pmem(device.clone()))
-                .unwrap();
+                .pmem
+                .add_device(device.clone());
 
             pci_devices
                 .restore_pci_device(
diff --git a/src/vmm/src/device_manager/persist.rs b/src/vmm/src/device_manager/persist.rs
@@ -44,7 +44,7 @@ use crate::devices::virtio::vsock::persist::{
 };
 use crate::devices::virtio::vsock::{Vsock, VsockError, VsockUnixBackend, VsockUnixBackendError};
 use crate::mmds::data_store::MmdsVersion;
-use crate::resources::{ResourcesError, VmResources};
+use crate::resources::VmResources;
 use crate::snapshot::Persist;
 use crate::vmm_config::mmds::MmdsConfigError;
 use crate::vstate::bus::BusError;
@@ -79,8 +79,6 @@ pub enum DevicePersistError {
     Entropy(#[from] EntropyError),
     /// Pmem: {0}
     Pmem(#[from] PmemError),
-    /// Resource misconfiguration: {0}. Is the snapshot file corrupted?
-    ResourcesError(#[from] ResourcesError),
     /// Could not activate device: {0}
     DeviceActivation(#[from] ActivateError),
 }
@@ -136,18 +134,6 @@ pub struct DeviceStates {
     pub pmem_devices: Vec<VirtioDeviceState<PmemState>>,
 }
 
-/// A type used to extract the concrete `Arc<Mutex<T>>` for each of the device
-/// types when restoring from a snapshot.
-#[derive(Debug)]
-pub enum SharedDeviceType {
-    VirtioBlock(Arc<Mutex<Block>>),
-    Network(Arc<Mutex<Net>>),
-    Balloon(Arc<Mutex<Balloon>>),
-    Vsock(Arc<Mutex<Vsock<VsockUnixBackend>>>),
-    Entropy(Arc<Mutex<Entropy>>),
-    Pmem(Arc<Mutex<Pmem>>),
-}
-
 pub struct MMIODevManagerConstructorArgs<'a> {
     pub mem: &'a GuestMemoryMmap,
     pub vm: &'a Arc<Vm>,
@@ -437,7 +423,8 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Balloon(device.clone()))?;
+                .balloon
+                .set_device(device.clone());
 
             restore_helper(
                 device.clone(),
@@ -459,7 +446,8 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::VirtioBlock(device.clone()))?;
+                .block
+                .add_virtio_device(device.clone());
 
             restore_helper(
                 device.clone(),
@@ -498,7 +486,8 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Network(device.clone()))?;
+                .net_builder
+                .add_device(device.clone());
 
             restore_helper(
                 device.clone(),
@@ -527,7 +516,8 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Vsock(device.clone()))?;
+                .vsock
+                .set_device(device.clone());
 
             restore_helper(
                 device.clone(),
@@ -551,7 +541,8 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Entropy(device.clone()))?;
+                .entropy
+                .set_device(device.clone());
 
             restore_helper(
                 device.clone(),
@@ -576,7 +567,8 @@ impl<'a> Persist<'a> for MMIODeviceManager {
 
             constructor_args
                 .vm_resources
-                .update_from_restored_device(SharedDeviceType::Pmem(device.clone()))?;
+                .pmem
+                .add_device(device.clone());
 
             restore_helper(
                 device.clone(),
diff --git a/src/vmm/src/resources.rs b/src/vmm/src/resources.rs
@@ -9,7 +9,6 @@ use serde::{Deserialize, Serialize};
 use vm_memory::GuestAddress;
 
 use crate::cpu_config::templates::CustomCpuTemplate;
-use crate::device_manager::persist::SharedDeviceType;
 use crate::logger::info;
 use crate::mmds;
 use crate::mmds::data_store::{Mmds, MmdsVersion};
@@ -230,40 +229,6 @@ impl VmResources {
         Ok(mmds.lock().expect("Poisoned lock"))
     }
 
-    /// Updates the resources from a restored device (used for configuring resources when
-    /// restoring from a snapshot).
-    pub fn update_from_restored_device(
-        &mut self,
-        device: SharedDeviceType,
-    ) -> Result<(), ResourcesError> {
-        match device {
-            SharedDeviceType::VirtioBlock(block) => {
-                self.block.add_virtio_device(block);
-            }
-            SharedDeviceType::Network(network) => {
-                self.net_builder.add_device(network);
-            }
-            SharedDeviceType::Balloon(balloon) => {
-                self.balloon.set_device(balloon);
-
-                if self.machine_config.huge_pages != HugePageConfig::None {
-                    return Err(ResourcesError::BalloonDevice(BalloonConfigError::HugePages));
-                }
-            }
-            SharedDeviceType::Vsock(vsock) => {
-                self.vsock.set_device(vsock);
-            }
-            SharedDeviceType::Entropy(entropy) => {
-                self.entropy.set_device(entropy);
-            }
-            SharedDeviceType::Pmem(pmem) => {
-                self.pmem.add_device(pmem);
-            }
-        }
-
-        Ok(())
-    }
-
     /// Add a custom CPU template to the VM resources
     /// to configure vCPUs.
     pub fn set_custom_cpu_template(&mut self, cpu_template: CustomCpuTemplate) {
@@ -563,7 +528,6 @@ mod tests {
     use crate::HTTP_MAX_PAYLOAD_SIZE;
     use crate::cpu_config::templates::test_utils::TEST_TEMPLATE_JSON;
     use crate::cpu_config::templates::{CpuTemplateType, StaticCpuTemplate};
-    use crate::devices::virtio::balloon::Balloon;
     use crate::devices::virtio::block::virtio::VirtioBlockError;
     use crate::devices::virtio::block::{BlockError, CacheType};
     use crate::devices::virtio::vsock::VSOCK_DEV_ID;
@@ -1539,31 +1503,6 @@ mod tests {
             .unwrap_err();
     }
 
-    #[test]
-    fn test_negative_restore_balloon_device_with_huge_pages() {
-        let mut vm_resources = default_vm_resources();
-        vm_resources.balloon = BalloonBuilder::new();
-        vm_resources
-            .update_machine_config(&MachineConfigUpdate {
-                huge_pages: Some(HugePageConfig::Hugetlbfs2M),
-                ..Default::default()
-            })
-            .unwrap();
-        let err = vm_resources
-            .update_from_restored_device(SharedDeviceType::Balloon(Arc::new(Mutex::new(
-                Balloon::new(128, false, 0).unwrap(),
-            ))))
-            .unwrap_err();
-        assert!(
-            matches!(
-                err,
-                ResourcesError::BalloonDevice(BalloonConfigError::HugePages)
-            ),
-            "{:?}",
-            err
-        );
-    }
-
     #[test]
     fn test_set_entropy_device() {
         let mut vm_resources = default_vm_resources();
