diff --git a/vm-virtio/src/pmem.rs b/vm-virtio/src/pmem.rs index da7afcba1..87a97e47a 100644 --- a/vm-virtio/src/pmem.rs +++ b/vm-virtio/src/pmem.rs @@ -26,7 +26,7 @@ use std::sync::Arc; use std::thread; use vm_memory::{ Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, - GuestMemoryError, GuestMemoryMmap, GuestUsize, + GuestMemoryError, GuestMemoryMmap, MmapRegion, }; use vm_migration::{Migratable, MigratableError, Pausable, Snapshottable, Transportable}; use vmm_sys_util::eventfd::EventFd; @@ -322,13 +322,22 @@ pub struct Pmem { interrupt_cb: Option>, epoll_threads: Option>>>, paused: Arc, + + // Hold ownership of the memory that is allocated for the device + // which will be automatically dropped when the device is dropped + _region: MmapRegion, } impl Pmem { - pub fn new(disk: File, addr: GuestAddress, size: GuestUsize, iommu: bool) -> io::Result { + pub fn new( + disk: File, + addr: GuestAddress, + _region: MmapRegion, + iommu: bool, + ) -> io::Result { let config = VirtioPmemConfig { start: addr.raw_value().to_le(), - size: size.to_le(), + size: (_region.size() as u64).to_le(), }; let mut avail_features = 1u64 << VIRTIO_F_VERSION_1; @@ -348,6 +357,7 @@ impl Pmem { interrupt_cb: None, epoll_threads: None, paused: Arc::new(AtomicBool::new(false)), + _region, }) } } diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs index 3dde9e935..27330f9cc 100644 --- a/vmm/src/device_manager.rs +++ b/vmm/src/device_manager.rs @@ -1467,8 +1467,6 @@ impl DeviceManager { .map_err(DeviceManagerError::NewMmapRegion)?; let addr: u64 = mmap_region.as_ptr() as u64; - self._mmap_regions.push(mmap_region); - self.memory_manager .lock() .unwrap() @@ -1482,7 +1480,7 @@ impl DeviceManager { .map_err(DeviceManagerError::MemoryManager)?; let virtio_pmem_device = Arc::new(Mutex::new( - vm_virtio::Pmem::new(file, pmem_guest_addr, size as GuestUsize, pmem_cfg.iommu) + vm_virtio::Pmem::new(file, pmem_guest_addr, mmap_region, pmem_cfg.iommu) .map_err(DeviceManagerError::CreateVirtioPmem)?, ));