Pmem: Fix resource leaks

src/vmm/src/builder.rs

@@ -51,6 +51,7 @@ use crate::utils::mib_to_bytes;
 use crate::vmm_config::instance_info::InstanceInfo;
 use crate::vmm_config::machine_config::MachineConfigError;
 use crate::vmm_config::memory_hotplug::MemoryHotplugConfig;
+use crate::vmm_config::pmem::PmemConfig;
 use crate::vstate::kvm::{Kvm, KvmError};
 use crate::vstate::memory::GuestRegionMmap;
 #[cfg(target_arch = "aarch64")]
@@ -246,7 +247,7 @@ pub fn build_microvm_for_boot(
         &mut device_manager,
         &vm,
         &mut boot_cmdline,
-        vm_resources.pmem.devices.iter(),
+        &vm_resources.pmem.configs,
         event_manager,
     )?;
 
@@ -726,36 +727,26 @@ fn attach_net_devices<'a, I: Iterator<Item = &'a Arc<Mutex<Net>>> + Debug>(
     Ok(())
 }
 
-fn attach_pmem_devices<'a, I: Iterator<Item = &'a Arc<Mutex<Pmem>>> + Debug>(
+fn attach_pmem_devices(
     device_manager: &mut DeviceManager,
     vm: &Arc<Vm>,
     cmdline: &mut LoaderKernelCmdline,
-    pmem_devices: I,
+    configs: &[PmemConfig],
     event_manager: &mut EventManager,
 ) -> Result<(), StartMicrovmError> {
-    for (i, device) in pmem_devices.enumerate() {
-        let id = {
-            let mut locked_dev = device.lock().expect("Poisoned lock");
-            if locked_dev.config.root_device {
-                cmdline.insert_str(format!("root=/dev/pmem{i}"))?;
-                match locked_dev.config.read_only {
-                    true => cmdline.insert_str("ro")?,
-                    false => cmdline.insert_str("rw")?,
-                }
+    for (i, config) in configs.iter().enumerate() {
+        if config.root_device {
+            cmdline.insert_str(format!("root=/dev/pmem{i}"))?;
+            match config.read_only {
+                true => cmdline.insert_str("ro")?,
+                false => cmdline.insert_str("rw")?,
             }
-            locked_dev.alloc_region(vm.as_ref());
-            locked_dev.set_mem_region(vm.as_ref())?;
-            locked_dev.config.id.to_string()
-        };
+        }
+        let id = config.id.clone();
+        let pmem = Pmem::new(vm.clone(), config.clone())?;
+        let device = Arc::new(Mutex::new(pmem));
 
-        device_manager.attach_virtio_device(
-            vm,
-            id,
-            device.clone(),
-            cmdline,
-            event_manager,
-            false,
-        )?;
+        device_manager.attach_virtio_device(vm, id, device, cmdline, event_manager, false)?;
     }
     Ok(())
 }
@@ -1049,7 +1040,7 @@ pub(crate) mod tests {
             &mut vmm.device_manager,
             &vmm.vm,
             cmdline,
-            builder.devices.iter(),
+            &builder.configs,
             event_manager,
         )
         .unwrap();

src/vmm/src/device_manager/pci_mngr.rs

@@ -576,15 +576,16 @@ impl<'a> Persist<'a> for PciDevices {
             let device = Arc::new(Mutex::new(Pmem::restore(
                 PmemConstructorArgs {
                     mem,
-                    vm: constructor_args.vm.as_ref(),
+                    vm: constructor_args.vm.clone(),
                 },
                 &pmem_state.device_state,
             )?));
 
             constructor_args
                 .vm_resources
                 .pmem
-                .add_device(device.clone());
+                .configs
+                .push(pmem_state.device_state.config.clone());
 
             pci_devices.restore_pci_device(
                 constructor_args.vm,

src/vmm/src/device_manager/persist.rs

@@ -562,15 +562,16 @@ impl<'a> Persist<'a> for MMIODeviceManager {
             let device = Arc::new(Mutex::new(Pmem::restore(
                 PmemConstructorArgs {
                     mem,
-                    vm: vm.as_ref(),
+                    vm: vm.clone(),
                 },
                 &pmem_state.device_state,
             )?));
 
             constructor_args
                 .vm_resources
                 .pmem
-                .add_device(device.clone());
+                .configs
+                .push(pmem_state.device_state.config.clone());
 
             restore_helper(
                 device,