From e1c2b179e7d67fa66f704c9c37dceb70465e213c Mon Sep 17 00:00:00 2001 From: Dylan Reid Date: Mon, 4 May 2026 18:29:03 -0700 Subject: [PATCH] vmm: device_manager: prune balloon and virtio_mem refs on eject self.balloon and self.virtio_mem_devices are not updated when eject_device removes the underlying device, leaving stale Arc> entries that resize_balloon / balloon_size and the virtio-mem DMA-handler iteration would dereference if reached after eject. Clear self.balloon and retain-out the matching virtio-mem entry in the PciDeviceHandle::Virtio eject arm, identifying the ejected device by Arc pointer-equality against the already-resolved Arc>. This change is defensive: DeviceManager::remove_device currently rejects VirtioDeviceType::Balloon and VirtioDeviceType::Mem with RemovalNotAllowed before pci_devices_down is set, so the guest never sees an eject notification and eject_device is never reached for either type today. If the allowlist is extended later, this cleanup keeps the post-eject state consistent. Signed-off-by: Dylan Reid --- vmm/src/device_manager.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vmm/src/device_manager.rs b/vmm/src/device_manager.rs index 83a3e64a8..9b8c20071 100644 --- a/vmm/src/device_manager.rs +++ b/vmm/src/device_manager.rs @@ -4923,6 +4923,19 @@ impl DeviceManager { self.virtio_devices .retain(|handler| !Arc::ptr_eq(&handler.virtio_device, &virtio_device)); + + // Drop any extra arcs that are kept for memory devices. + if let Some(balloon) = &self.balloon { + let balloon_dyn = + Arc::clone(balloon) as Arc>; + if Arc::ptr_eq(&balloon_dyn, &virtio_device) { + self.balloon = None; + } + } + self.virtio_mem_devices.retain(|mem| { + let mem_dyn = Arc::clone(mem) as Arc>; + !Arc::ptr_eq(&mem_dyn, &virtio_device) + }); } event!(