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<Mutex<...>> 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<Mutex<dyn VirtioDevice>>.

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 <dgreid@fb.com>

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<Mutex<dyn virtio_devices::VirtioDevice>>;
+                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<Mutex<dyn virtio_devices::VirtioDevice>>;
+                !Arc::ptr_eq(&mem_dyn, &virtio_device)
+            });
         }
 
         event!(