From 1b1f2a97cd3e1c76bccb1e9016600ed924b2606f Mon Sep 17 00:00:00 2001 From: John Starks Date: Wed, 25 Feb 2026 20:27:24 +0000 Subject: [PATCH 1/6] task_control: add poll_stop() method Add TaskControl::poll_stop() as the poll-based variant of stop(). Returns Poll::Ready(true) when the task was running and has stopped, Poll::Ready(false) if already stopped/not inserted/completed, and Poll::Pending while waiting for the task to stop. Rewrite stop() as poll_fn(|cx| self.poll_stop(cx)).await. This is a prerequisite for virtio poll_disable, where devices need a non-async way to drive worker shutdown from poll_disable(). --- support/task_control/src/lib.rs | 85 ++++++++++++++++++++++++--------- 1 file changed, 63 insertions(+), 22 deletions(-) diff --git a/support/task_control/src/lib.rs b/support/task_control/src/lib.rs index 18b85eb992..8fc78c773c 100644 --- a/support/task_control/src/lib.rs +++ b/support/task_control/src/lib.rs @@ -517,43 +517,51 @@ impl, S: 'static + Send> TaskControl { .await } - /// Stops the task, waiting for it to be cancelled. + /// Poll variant of [`stop`](Self::stop). Signals the task to stop and polls + /// for completion. /// - /// Returns true if the task was previously running. Returns false if the - /// task was not running, not inserted, or had already completed. - pub async fn stop(&mut self) -> bool { + /// Returns `Poll::Ready(true)` if the task was previously running and has + /// now stopped. Returns `Poll::Ready(false)` if the task was not running, + /// not inserted, or had already completed. Returns `Poll::Pending` if the + /// task has not yet stopped. + pub fn poll_stop(&mut self, cx: &mut Context<'_>) -> Poll { match &mut self.inner { Inner::WithState { activity, shared, .. } => match activity { Activity::Running => { - let task_and_state = poll_fn(|cx| { - let mut shared = shared.lock(); - shared.stop = true; - if shared.task_and_state.is_none() || !shared.calls.is_empty() { - shared.outer_waker = Some(cx.waker().clone()); - let waker = shared.inner_waker.take(); - drop(shared); - if let Some(waker) = waker { - waker.wake(); - } - return Poll::Pending; + let mut shared = shared.lock(); + shared.stop = true; + if shared.task_and_state.is_none() || !shared.calls.is_empty() { + shared.outer_waker = Some(cx.waker().clone()); + let waker = shared.inner_waker.take(); + drop(shared); + if let Some(waker) = waker { + waker.wake(); } - Poll::Ready(shared.task_and_state.take().unwrap()) - }) - .await; - + return Poll::Pending; + } + let task_and_state = shared.task_and_state.take().unwrap(); + drop(shared); let done = task_and_state.done; *activity = Activity::Stopped(task_and_state); - !done + Poll::Ready(!done) } - _ => false, + _ => Poll::Ready(false), }, - Inner::NoState(_) => false, + Inner::NoState(_) => Poll::Ready(false), Inner::Invalid => unreachable!(), } } + /// Stops the task, waiting for it to be cancelled. + /// + /// Returns true if the task was previously running. Returns false if the + /// task was not running, not inserted, or had already completed. + pub async fn stop(&mut self) -> bool { + poll_fn(|cx| self.poll_stop(cx)).await + } + /// Removes the task state. /// /// Panics if the task is not stopped. @@ -655,4 +663,37 @@ mod tests { assert!(t.stop().await); assert_eq!(t.task_mut().0, 8); } + + #[async_test] + async fn test_poll_stop(driver: DefaultDriver) { + let mut t = TaskControl::new(Foo(5)); + + // poll_stop on a task without state returns Ready(false). + assert_eq!( + std::future::poll_fn(|cx| Poll::Ready(t.poll_stop(cx))).await, + Poll::Ready(false) + ); + + t.insert(&driver, "test", false); + + // poll_stop on a stopped (not started) task returns Ready(false). + assert_eq!( + std::future::poll_fn(|cx| Poll::Ready(t.poll_stop(cx))).await, + Poll::Ready(false) + ); + + assert!(t.start()); + yield_once().await; + + // poll_stop drives the task to stop, equivalent to stop().await. + let result = std::future::poll_fn(|cx| t.poll_stop(cx)).await; + assert!(result); // was running + assert_eq!(t.task().0, 6); + + // poll_stop after already stopped returns Ready(false). + assert_eq!( + std::future::poll_fn(|cx| Poll::Ready(t.poll_stop(cx))).await, + Poll::Ready(false) + ); + } } From 45bdaec269ab597960ce7d127b2fc5e735750560 Mon Sep 17 00:00:00 2001 From: John Starks Date: Wed, 25 Feb 2026 20:34:42 +0000 Subject: [PATCH 2/6] virtio: change VirtioDevice::disable() to poll_disable() Change the VirtioDevice trait method from sync fn disable(&mut self) to fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()>. Device implementations now use TaskControl::poll_stop() to stop workers instead of spawning detached fire-and-forget tasks. This eliminates the spawn+detach pattern that could outlive the device. Updated implementations: - virtio-net: sends CoordinatorMessage::Disable, returns Ready (unchanged semantics -- coordinator handles cleanup internally) - virtiofs: poll_stop on all queue workers - virtio-pmem: poll_stop on single worker - virtio-9p: poll_stop on single worker - TestDevice: poll_stop on all workers Transport call sites (PCI and MMIO) use a noop waker to initiate the disable, maintaining the same fire-and-forget behavior for now. Proper async completion via PollDevice will be added in a follow-up. Removed now-unused dependencies: pal_async from virtio_p9, virtio_pmem, and virtiofs; futures from virtiofs. --- Cargo.lock | 4 ---- vm/devices/virtio/virtio/src/common.rs | 10 ++++++++- vm/devices/virtio/virtio/src/tests.rs | 18 +++++----------- .../virtio/virtio/src/transport/mmio.rs | 8 +++++-- vm/devices/virtio/virtio/src/transport/pci.rs | 8 +++++-- vm/devices/virtio/virtio_net/src/lib.rs | 4 +++- vm/devices/virtio/virtio_p9/Cargo.toml | 1 - vm/devices/virtio/virtio_p9/src/lib.rs | 20 ++++++++---------- vm/devices/virtio/virtio_pmem/Cargo.toml | 2 -- vm/devices/virtio/virtio_pmem/src/lib.rs | 16 +++++++------- vm/devices/virtio/virtiofs/Cargo.toml | 2 -- vm/devices/virtio/virtiofs/src/virtio.rs | 21 ++++++++----------- 12 files changed, 54 insertions(+), 60 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0a31c69ff0..2ab00bc703 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9005,7 +9005,6 @@ dependencies = [ "event-listener", "guestmem", "inspect", - "pal_async", "plan9", "task_control", "tracing", @@ -9025,7 +9024,6 @@ dependencies = [ "fs-err", "guestmem", "inspect", - "pal_async", "sparse_mmap", "task_control", "tracing", @@ -9052,14 +9050,12 @@ dependencies = [ "async-trait", "event-listener", "fuse", - "futures", "guestmem", "inspect", "lx", "lxutil", "ntapi", "pal", - "pal_async", "parking_lot", "task_control", "tracing", diff --git a/vm/devices/virtio/virtio/src/common.rs b/vm/devices/virtio/virtio/src/common.rs index d309e0514f..72784d2c9f 100644 --- a/vm/devices/virtio/virtio/src/common.rs +++ b/vm/devices/virtio/virtio/src/common.rs @@ -442,7 +442,15 @@ pub trait VirtioDevice: inspect::InspectMut + Send { fn read_registers_u32(&self, offset: u16) -> u32; fn write_registers_u32(&mut self, offset: u16, val: u32); fn enable(&mut self, resources: Resources); - fn disable(&mut self); + /// Poll the device to complete a disable/reset operation. + /// + /// This is called when the guest writes status=0 (device reset). The device + /// should stop workers and drain any in-flight IO. Returns `Poll::Ready(())` + /// when the disable is complete, or `Poll::Pending` if more work is needed. + /// + /// Devices that don't need async cleanup can return `Poll::Ready(())` + /// immediately. + fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()>; } pub struct QueueResources { diff --git a/vm/devices/virtio/virtio/src/tests.rs b/vm/devices/virtio/virtio/src/tests.rs index 7dee922185..3bb471bc0c 100644 --- a/vm/devices/virtio/virtio/src/tests.rs +++ b/vm/devices/virtio/virtio/src/tests.rs @@ -846,20 +846,12 @@ impl VirtioDevice for TestDevice { .collect(); } - fn disable(&mut self) { - if self.workers.is_empty() { - return; + fn poll_disable(&mut self, cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + for worker in &mut self.workers { + std::task::ready!(worker.poll_stop(cx)); } - self.exit_event.notify(usize::MAX); - let mut workers = self.workers.drain(..).collect::>(); - self.driver - .spawn("shutdown-test-virtio-queues".to_owned(), async move { - futures::future::join_all(workers.iter_mut().map(async |worker| { - worker.stop().await; - })) - .await; - }) - .detach(); + self.workers.clear(); + std::task::Poll::Ready(()) } } diff --git a/vm/devices/virtio/virtio/src/transport/mmio.rs b/vm/devices/virtio/virtio/src/transport/mmio.rs index 2eaecef9d0..27c1f48cda 100644 --- a/vm/devices/virtio/virtio/src/transport/mmio.rs +++ b/vm/devices/virtio/virtio/src/transport/mmio.rs @@ -147,7 +147,9 @@ impl VirtioMmioDevice { impl Drop for VirtioMmioDevice { fn drop(&mut self) { - self.device.disable(); + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(&waker); + let _ = self.device.poll_disable(&mut cx); } } @@ -365,7 +367,9 @@ impl VirtioMmioDevice { self.config_generation = 0; if started { self.doorbells.clear(); - self.device.disable(); + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(&waker); + let _ = self.device.poll_disable(&mut cx); } self.interrupt_state.lock().update(false, !0); } diff --git a/vm/devices/virtio/virtio/src/transport/pci.rs b/vm/devices/virtio/virtio/src/transport/pci.rs index dc7b4a8f04..81aed86242 100644 --- a/vm/devices/virtio/virtio/src/transport/pci.rs +++ b/vm/devices/virtio/virtio/src/transport/pci.rs @@ -405,7 +405,9 @@ impl VirtioPciDevice { self.config_generation = 0; if started { self.doorbells.clear(); - self.device.disable(); + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(&waker); + let _ = self.device.poll_disable(&mut cx); } *self.interrupt_status.lock() = 0; } @@ -561,7 +563,9 @@ impl VirtioPciDevice { impl Drop for VirtioPciDevice { fn drop(&mut self) { // TODO conditionalize - self.device.disable(); + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(&waker); + let _ = self.device.poll_disable(&mut cx); } } diff --git a/vm/devices/virtio/virtio_net/src/lib.rs b/vm/devices/virtio/virtio_net/src/lib.rs index eda8620d6a..c808274b4c 100644 --- a/vm/devices/virtio/virtio_net/src/lib.rs +++ b/vm/devices/virtio/virtio_net/src/lib.rs @@ -41,6 +41,7 @@ use pal_async::wait::PolledWait; use std::future::pending; use std::mem::offset_of; use std::sync::Arc; +use std::task::Context; use std::task::Poll; use task_control::AsyncRun; use task_control::InspectTaskMut; @@ -346,10 +347,11 @@ impl VirtioDevice for Device { self.coordinator.start(); } - fn disable(&mut self) { + fn poll_disable(&mut self, _cx: &mut Context<'_>) -> Poll<()> { if let Some(send) = self.coordinator_send.take() { send.send(CoordinatorMessage::Disable); } + Poll::Ready(()) } } diff --git a/vm/devices/virtio/virtio_p9/Cargo.toml b/vm/devices/virtio/virtio_p9/Cargo.toml index 0ef2f340aa..f3876a192e 100644 --- a/vm/devices/virtio/virtio_p9/Cargo.toml +++ b/vm/devices/virtio/virtio_p9/Cargo.toml @@ -20,7 +20,6 @@ task_control.workspace = true anyhow.workspace = true async-trait.workspace = true event-listener.workspace = true -pal_async.workspace = true tracing.workspace = true [lints] diff --git a/vm/devices/virtio/virtio_p9/src/lib.rs b/vm/devices/virtio/virtio_p9/src/lib.rs index cbf1b8951a..725b26ecfd 100644 --- a/vm/devices/virtio/virtio_p9/src/lib.rs +++ b/vm/devices/virtio/virtio_p9/src/lib.rs @@ -10,9 +10,11 @@ pub mod resolver; use async_trait::async_trait; use guestmem::GuestMemory; use inspect::InspectMut; -use pal_async::task::Spawn; use plan9::Plan9FileSystem; use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::ready; use task_control::TaskControl; use virtio::DeviceTraits; use virtio::Resources; @@ -132,16 +134,12 @@ impl VirtioDevice for VirtioPlan9Device { )); } - fn disable(&mut self) { - let Some(mut worker) = self.worker.take() else { - return; - }; - self.exit_event.notify(usize::MAX); - self.driver - .spawn("shutdown-virtio-9p-queue".to_owned(), async move { - worker.stop().await; - }) - .detach(); + fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()> { + if let Some(worker) = &mut self.worker { + ready!(worker.poll_stop(cx)); + } + self.worker = None; + Poll::Ready(()) } } diff --git a/vm/devices/virtio/virtio_pmem/Cargo.toml b/vm/devices/virtio/virtio_pmem/Cargo.toml index c345876c4d..6649e952d9 100644 --- a/vm/devices/virtio/virtio_pmem/Cargo.toml +++ b/vm/devices/virtio/virtio_pmem/Cargo.toml @@ -14,8 +14,6 @@ inspect = { workspace = true, features = ["filepath"] } guestmem.workspace = true vmcore.workspace = true vm_resource.workspace = true - -pal_async.workspace = true sparse_mmap.workspace = true task_control.workspace = true diff --git a/vm/devices/virtio/virtio_pmem/src/lib.rs b/vm/devices/virtio/virtio_pmem/src/lib.rs index 75c7a14d42..36b21c815c 100644 --- a/vm/devices/virtio/virtio_pmem/src/lib.rs +++ b/vm/devices/virtio/virtio_pmem/src/lib.rs @@ -10,9 +10,10 @@ use anyhow::Context; use async_trait::async_trait; use guestmem::GuestMemory; use inspect::InspectMut; -use pal_async::task::Spawn; use std::fs; use std::sync::Arc; +use std::task::Poll; +use std::task::ready; use task_control::TaskControl; use virtio::DeviceTraits; use virtio::DeviceTraitsSharedMemory; @@ -126,15 +127,12 @@ impl VirtioDevice for Device { }; } - fn disable(&mut self) { - self.exit_event.notify(usize::MAX); - if let Some(mut worker) = self.worker.take() { - self.driver - .spawn("shutdown-virtio-pmem-queue".to_owned(), async move { - worker.stop().await; - }) - .detach(); + fn poll_disable(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> { + if let Some(worker) = &mut self.worker { + ready!(worker.poll_stop(cx)); } + self.worker = None; + Poll::Ready(()) } } diff --git a/vm/devices/virtio/virtiofs/Cargo.toml b/vm/devices/virtio/virtiofs/Cargo.toml index e54797a0f5..4d4af8e571 100644 --- a/vm/devices/virtio/virtiofs/Cargo.toml +++ b/vm/devices/virtio/virtiofs/Cargo.toml @@ -19,13 +19,11 @@ inspect.workspace = true lx.workspace = true lxutil.workspace = true pal.workspace = true -pal_async.workspace = true task_control.workspace = true anyhow.workspace = true async-trait.workspace = true event-listener.workspace = true -futures.workspace = true parking_lot.workspace = true tracing.workspace = true zerocopy.workspace = true diff --git a/vm/devices/virtio/virtiofs/src/virtio.rs b/vm/devices/virtio/virtiofs/src/virtio.rs index f9cd5ba4d8..0c5da827b2 100644 --- a/vm/devices/virtio/virtiofs/src/virtio.rs +++ b/vm/devices/virtio/virtiofs/src/virtio.rs @@ -7,10 +7,12 @@ use async_trait::async_trait; use guestmem::GuestMemory; use guestmem::MappedMemoryRegion; use inspect::InspectMut; -use pal_async::task::Spawn; use std::io; use std::io::Write; use std::sync::Arc; +use std::task::Context; +use std::task::Poll; +use std::task::ready; use task_control::TaskControl; use virtio::DeviceTraits; use virtio::DeviceTraitsSharedMemory; @@ -158,17 +160,12 @@ impl VirtioDevice for VirtioFsDevice { .collect(); } - fn disable(&mut self) { - self.exit_event.notify(usize::MAX); - let mut workers = self.workers.drain(..).collect::>(); - self.driver - .spawn("shutdown-virtiofs-queues".to_owned(), async move { - futures::future::join_all(workers.iter_mut().map(async |worker| { - worker.stop().await; - })) - .await; - }) - .detach(); + fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()> { + for worker in &mut self.workers { + ready!(worker.poll_stop(cx)); + } + self.workers.clear(); + Poll::Ready(()) } } From 5b49612d63df1cf1f7fe1f02cf1712d1bbd1c37b Mon Sep 17 00:00:00 2001 From: John Starks Date: Wed, 25 Feb 2026 20:41:59 +0000 Subject: [PATCH 3/6] virtio: implement PollDevice for async disable in transports Add PollDevice support to both MMIO and PCI virtio transports to enable asynchronous device disabling. When a guest writes status=0 to trigger a device reset, the transport now: 1. Initiates poll_disable with a noop waker 2. If the device completes synchronously (fast path), clears status immediately 3. If pending, sets a 'disabling' flag and returns -- the guest sees the old status bits (non-zero) until PollDevice drives completion The PollDevice::poll_device implementation drives the pending disable to completion via the chipset infrastructure's polling mechanism. Once ready, it clears the device status to 0 so the guest can observe the reset. Also implements ChangeDeviceState::reset() for both transports (previously TODO), which drives poll_disable to completion and resets all transport state. Updates Drop to only poll_disable when the device was actually started or is in the middle of disabling. --- vm/devices/virtio/virtio/src/tests.rs | 1 - .../virtio/virtio/src/transport/mmio.rs | 48 +++++++++++++++--- vm/devices/virtio/virtio/src/transport/pci.rs | 49 ++++++++++++++++--- 3 files changed, 84 insertions(+), 14 deletions(-) diff --git a/vm/devices/virtio/virtio/src/tests.rs b/vm/devices/virtio/virtio/src/tests.rs index 3bb471bc0c..66a9219818 100644 --- a/vm/devices/virtio/virtio/src/tests.rs +++ b/vm/devices/virtio/virtio/src/tests.rs @@ -33,7 +33,6 @@ use guestmem::GuestMemoryBackingError; use inspect::InspectMut; use pal_async::DefaultDriver; use pal_async::async_test; -use pal_async::task::Spawn; use pal_async::timer::PolledTimer; use pal_event::Event; use parking_lot::Mutex; diff --git a/vm/devices/virtio/virtio/src/transport/mmio.rs b/vm/devices/virtio/virtio/src/transport/mmio.rs index 27c1f48cda..d7d5d95cb7 100644 --- a/vm/devices/virtio/virtio/src/transport/mmio.rs +++ b/vm/devices/virtio/virtio/src/transport/mmio.rs @@ -11,6 +11,7 @@ use crate::spec::*; use chipset_device::ChipsetDevice; use chipset_device::io::IoResult; use chipset_device::mmio::MmioIntercept; +use chipset_device::poll_device::PollDevice; use device_emulators::ReadWriteRequestType; use device_emulators::read_as_u32_chunks; use device_emulators::write_as_u32_chunks; @@ -53,6 +54,7 @@ pub struct VirtioMmioDevice { #[inspect(skip)] queues: Vec, device_status: VirtioDeviceStatus, + disabling: bool, config_generation: u32, #[inspect(skip)] doorbells: VirtioDoorbells, @@ -129,6 +131,7 @@ impl VirtioMmioDevice { events, queues, device_status: VirtioDeviceStatus::new(), + disabling: false, config_generation: 0, doorbells: VirtioDoorbells::new(doorbell_registration), interrupt_state, @@ -147,9 +150,11 @@ impl VirtioMmioDevice { impl Drop for VirtioMmioDevice { fn drop(&mut self) { - let waker = std::task::Waker::noop(); - let mut cx = std::task::Context::from_waker(&waker); - let _ = self.device.poll_disable(&mut cx); + if self.device_status.driver_ok() || self.disabling { + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(&waker); + let _ = self.device.poll_disable(&mut cx); + } } } @@ -362,16 +367,24 @@ impl VirtioMmioDevice { // Device status 112 => { if val == 0 { + if self.disabling { + return; + } let started = self.device_status.driver_ok(); - self.device_status = VirtioDeviceStatus::new(); self.config_generation = 0; if started { self.doorbells.clear(); let waker = std::task::Waker::noop(); let mut cx = std::task::Context::from_waker(&waker); - let _ = self.device.poll_disable(&mut cx); + if self.device.poll_disable(&mut cx).is_pending() { + self.disabling = true; + return; + } } + // Fast path: disable completed synchronously. + self.device_status = VirtioDeviceStatus::new(); self.interrupt_state.lock().update(false, !0); + return; } let new_status = VirtioDeviceStatus::from(val as u8); @@ -484,7 +497,26 @@ impl ChangeDeviceState for VirtioMmioDevice { async fn stop(&mut self) {} async fn reset(&mut self) { - // TODO + if self.device_status.driver_ok() || self.disabling { + self.doorbells.clear(); + std::future::poll_fn(|cx| self.device.poll_disable(cx)).await; + } + self.device_status = VirtioDeviceStatus::new(); + self.disabling = false; + self.config_generation = 0; + self.interrupt_state.lock().update(false, !0); + } +} + +impl PollDevice for VirtioMmioDevice { + fn poll_device(&mut self, cx: &mut std::task::Context<'_>) { + if self.disabling { + if self.device.poll_disable(cx).is_ready() { + self.device_status = VirtioDeviceStatus::new(); + self.disabling = false; + self.interrupt_state.lock().update(false, !0); + } + } } } @@ -492,6 +524,10 @@ impl ChipsetDevice for VirtioMmioDevice { fn supports_mmio(&mut self) -> Option<&mut dyn MmioIntercept> { Some(self) } + + fn supports_poll_device(&mut self) -> Option<&mut dyn PollDevice> { + Some(self) + } } impl SaveRestore for VirtioMmioDevice { diff --git a/vm/devices/virtio/virtio/src/transport/pci.rs b/vm/devices/virtio/virtio/src/transport/pci.rs index 81aed86242..55e5352adf 100644 --- a/vm/devices/virtio/virtio/src/transport/pci.rs +++ b/vm/devices/virtio/virtio/src/transport/pci.rs @@ -17,6 +17,7 @@ use chipset_device::io::IoResult; use chipset_device::mmio::MmioIntercept; use chipset_device::mmio::RegisterMmioIntercept; use chipset_device::pci::PciConfigSpace; +use chipset_device::poll_device::PollDevice; use device_emulators::ReadWriteRequestType; use device_emulators::read_as_u32_chunks; use device_emulators::write_as_u32_chunks; @@ -84,6 +85,7 @@ pub struct VirtioPciDevice { interrupt_status: Arc>, #[inspect(hex)] device_status: VirtioDeviceStatus, + disabling: bool, config_generation: u32, config_space: ConfigSpaceType0Emulator, @@ -233,6 +235,7 @@ impl VirtioPciDevice { msix_vectors, interrupt_status: Arc::new(Mutex::new(0)), device_status: VirtioDeviceStatus::new(), + disabling: false, config_generation: 0, interrupt_kind, config_space, @@ -400,16 +403,24 @@ impl VirtioPciDevice { self.queue_select = val >> 16; let val = val & 0xff; if val == 0 { + if self.disabling { + return; + } let started = self.device_status.driver_ok(); - self.device_status = VirtioDeviceStatus::new(); self.config_generation = 0; if started { self.doorbells.clear(); let waker = std::task::Waker::noop(); let mut cx = std::task::Context::from_waker(&waker); - let _ = self.device.poll_disable(&mut cx); + if self.device.poll_disable(&mut cx).is_pending() { + self.disabling = true; + return; + } } + // Fast path: disable completed synchronously. + self.device_status = VirtioDeviceStatus::new(); *self.interrupt_status.lock() = 0; + return; } let new_status = VirtioDeviceStatus::from(val as u8); @@ -562,10 +573,11 @@ impl VirtioPciDevice { impl Drop for VirtioPciDevice { fn drop(&mut self) { - // TODO conditionalize - let waker = std::task::Waker::noop(); - let mut cx = std::task::Context::from_waker(&waker); - let _ = self.device.poll_disable(&mut cx); + if self.device_status.driver_ok() || self.disabling { + let waker = std::task::Waker::noop(); + let mut cx = std::task::Context::from_waker(&waker); + let _ = self.device.poll_disable(&mut cx); + } } } @@ -603,7 +615,26 @@ impl ChangeDeviceState for VirtioPciDevice { async fn stop(&mut self) {} async fn reset(&mut self) { - // TODO + if self.device_status.driver_ok() || self.disabling { + self.doorbells.clear(); + std::future::poll_fn(|cx| self.device.poll_disable(cx)).await; + } + self.device_status = VirtioDeviceStatus::new(); + self.disabling = false; + self.config_generation = 0; + *self.interrupt_status.lock() = 0; + } +} + +impl PollDevice for VirtioPciDevice { + fn poll_device(&mut self, cx: &mut std::task::Context<'_>) { + if self.disabling { + if self.device.poll_disable(cx).is_ready() { + self.device_status = VirtioDeviceStatus::new(); + self.disabling = false; + *self.interrupt_status.lock() = 0; + } + } } } @@ -615,6 +646,10 @@ impl ChipsetDevice for VirtioPciDevice { fn supports_pci(&mut self) -> Option<&mut dyn PciConfigSpace> { Some(self) } + + fn supports_poll_device(&mut self) -> Option<&mut dyn PollDevice> { + Some(self) + } } impl SaveRestore for VirtioPciDevice { From cea481bf8ca07487bfc6057dccb5e88e4f65d054 Mon Sep 17 00:00:00 2001 From: John Starks Date: Wed, 25 Feb 2026 20:50:46 +0000 Subject: [PATCH 4/6] virtio: remove Drop impls from transports The poll_disable calls in Drop were a half-measure -- polling with a noop waker initiates shutdown but cannot drive it to completion. Proper async teardown will require a different mechanism (e.g., an async method that takes ownership). Remove the Drop impls to keep things clean in the meantime. --- vm/devices/virtio/virtio/src/transport/mmio.rs | 10 ---------- vm/devices/virtio/virtio/src/transport/pci.rs | 10 ---------- 2 files changed, 20 deletions(-) diff --git a/vm/devices/virtio/virtio/src/transport/mmio.rs b/vm/devices/virtio/virtio/src/transport/mmio.rs index d7d5d95cb7..ca41eb3b83 100644 --- a/vm/devices/virtio/virtio/src/transport/mmio.rs +++ b/vm/devices/virtio/virtio/src/transport/mmio.rs @@ -148,16 +148,6 @@ impl VirtioMmioDevice { } } -impl Drop for VirtioMmioDevice { - fn drop(&mut self) { - if self.device_status.driver_ok() || self.disabling { - let waker = std::task::Waker::noop(); - let mut cx = std::task::Context::from_waker(&waker); - let _ = self.device.poll_disable(&mut cx); - } - } -} - impl VirtioMmioDevice { pub(crate) fn read_u32(&self, address: u64) -> u32 { let offset = (address & 0xfff) as u16; diff --git a/vm/devices/virtio/virtio/src/transport/pci.rs b/vm/devices/virtio/virtio/src/transport/pci.rs index 55e5352adf..165fa25c33 100644 --- a/vm/devices/virtio/virtio/src/transport/pci.rs +++ b/vm/devices/virtio/virtio/src/transport/pci.rs @@ -571,16 +571,6 @@ impl VirtioPciDevice { } } -impl Drop for VirtioPciDevice { - fn drop(&mut self) { - if self.device_status.driver_ok() || self.disabling { - let waker = std::task::Waker::noop(); - let mut cx = std::task::Context::from_waker(&waker); - let _ = self.device.poll_disable(&mut cx); - } - } -} - impl VirtioPciDevice { fn read_bar_u32(&mut self, bar: u8, offset: u16) -> u32 { match bar { From 8b5dfa0e2ebaabda367d47b4e348e2d267f32c41 Mon Sep 17 00:00:00 2001 From: John Starks Date: Fri, 6 Mar 2026 04:56:25 +0000 Subject: [PATCH 5/6] feedback --- vm/devices/virtio/virtio/src/tests.rs | 1 + vm/devices/virtio/virtio/src/transport/mmio.rs | 12 +++++++++++- vm/devices/virtio/virtio/src/transport/pci.rs | 13 ++++++++++++- vm/devices/virtio/virtio_p9/src/lib.rs | 1 + vm/devices/virtio/virtio_pmem/src/lib.rs | 1 + vm/devices/virtio/virtiofs/src/virtio.rs | 1 + 6 files changed, 27 insertions(+), 2 deletions(-) diff --git a/vm/devices/virtio/virtio/src/tests.rs b/vm/devices/virtio/virtio/src/tests.rs index 66a9219818..c43d92be3d 100644 --- a/vm/devices/virtio/virtio/src/tests.rs +++ b/vm/devices/virtio/virtio/src/tests.rs @@ -846,6 +846,7 @@ impl VirtioDevice for TestDevice { } fn poll_disable(&mut self, cx: &mut std::task::Context<'_>) -> std::task::Poll<()> { + self.exit_event.notify(usize::MAX); for worker in &mut self.workers { std::task::ready!(worker.poll_stop(cx)); } diff --git a/vm/devices/virtio/virtio/src/transport/mmio.rs b/vm/devices/virtio/virtio/src/transport/mmio.rs index ca41eb3b83..d46249c8ab 100644 --- a/vm/devices/virtio/virtio/src/transport/mmio.rs +++ b/vm/devices/virtio/virtio/src/transport/mmio.rs @@ -55,6 +55,7 @@ pub struct VirtioMmioDevice { queues: Vec, device_status: VirtioDeviceStatus, disabling: bool, + poll_waker: Option, config_generation: u32, #[inspect(skip)] doorbells: VirtioDoorbells, @@ -132,6 +133,7 @@ impl VirtioMmioDevice { queues, device_status: VirtioDeviceStatus::new(), disabling: false, + poll_waker: None, config_generation: 0, doorbells: VirtioDoorbells::new(doorbell_registration), interrupt_state, @@ -364,10 +366,17 @@ impl VirtioMmioDevice { self.config_generation = 0; if started { self.doorbells.clear(); + // Try the fast path: poll with a noop waker to see if + // the device can disable synchronously. let waker = std::task::Waker::noop(); - let mut cx = std::task::Context::from_waker(&waker); + let mut cx = std::task::Context::from_waker(waker); if self.device.poll_disable(&mut cx).is_pending() { self.disabling = true; + // Wake the real poll waker so that poll_device will + // re-poll with a real waker, replacing the noop one. + if let Some(waker) = self.poll_waker.take() { + waker.wake(); + } return; } } @@ -500,6 +509,7 @@ impl ChangeDeviceState for VirtioMmioDevice { impl PollDevice for VirtioMmioDevice { fn poll_device(&mut self, cx: &mut std::task::Context<'_>) { + self.poll_waker = Some(cx.waker().clone()); if self.disabling { if self.device.poll_disable(cx).is_ready() { self.device_status = VirtioDeviceStatus::new(); diff --git a/vm/devices/virtio/virtio/src/transport/pci.rs b/vm/devices/virtio/virtio/src/transport/pci.rs index 165fa25c33..87bf1ed66b 100644 --- a/vm/devices/virtio/virtio/src/transport/pci.rs +++ b/vm/devices/virtio/virtio/src/transport/pci.rs @@ -86,6 +86,8 @@ pub struct VirtioPciDevice { #[inspect(hex)] device_status: VirtioDeviceStatus, disabling: bool, + #[inspect(skip)] + poll_waker: Option, config_generation: u32, config_space: ConfigSpaceType0Emulator, @@ -236,6 +238,7 @@ impl VirtioPciDevice { interrupt_status: Arc::new(Mutex::new(0)), device_status: VirtioDeviceStatus::new(), disabling: false, + poll_waker: None, config_generation: 0, interrupt_kind, config_space, @@ -410,10 +413,17 @@ impl VirtioPciDevice { self.config_generation = 0; if started { self.doorbells.clear(); + // Try the fast path: poll with a noop waker to see if + // the device can disable synchronously. let waker = std::task::Waker::noop(); - let mut cx = std::task::Context::from_waker(&waker); + let mut cx = std::task::Context::from_waker(waker); if self.device.poll_disable(&mut cx).is_pending() { self.disabling = true; + // Wake the real poll waker so that poll_device will + // re-poll with a real waker, replacing the noop one. + if let Some(waker) = self.poll_waker.take() { + waker.wake(); + } return; } } @@ -618,6 +628,7 @@ impl ChangeDeviceState for VirtioPciDevice { impl PollDevice for VirtioPciDevice { fn poll_device(&mut self, cx: &mut std::task::Context<'_>) { + self.poll_waker = Some(cx.waker().clone()); if self.disabling { if self.device.poll_disable(cx).is_ready() { self.device_status = VirtioDeviceStatus::new(); diff --git a/vm/devices/virtio/virtio_p9/src/lib.rs b/vm/devices/virtio/virtio_p9/src/lib.rs index 725b26ecfd..9c432d8e92 100644 --- a/vm/devices/virtio/virtio_p9/src/lib.rs +++ b/vm/devices/virtio/virtio_p9/src/lib.rs @@ -135,6 +135,7 @@ impl VirtioDevice for VirtioPlan9Device { } fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()> { + self.exit_event.notify(usize::MAX); if let Some(worker) = &mut self.worker { ready!(worker.poll_stop(cx)); } diff --git a/vm/devices/virtio/virtio_pmem/src/lib.rs b/vm/devices/virtio/virtio_pmem/src/lib.rs index 36b21c815c..5b90a6d8a6 100644 --- a/vm/devices/virtio/virtio_pmem/src/lib.rs +++ b/vm/devices/virtio/virtio_pmem/src/lib.rs @@ -128,6 +128,7 @@ impl VirtioDevice for Device { } fn poll_disable(&mut self, cx: &mut std::task::Context<'_>) -> Poll<()> { + self.exit_event.notify(usize::MAX); if let Some(worker) = &mut self.worker { ready!(worker.poll_stop(cx)); } diff --git a/vm/devices/virtio/virtiofs/src/virtio.rs b/vm/devices/virtio/virtiofs/src/virtio.rs index 0c5da827b2..f03fedb41c 100644 --- a/vm/devices/virtio/virtiofs/src/virtio.rs +++ b/vm/devices/virtio/virtiofs/src/virtio.rs @@ -161,6 +161,7 @@ impl VirtioDevice for VirtioFsDevice { } fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()> { + self.exit_event.notify(usize::MAX); for worker in &mut self.workers { ready!(worker.poll_stop(cx)); } From 170b89cf45906ff2396ae79744fe1d1da5b70060 Mon Sep 17 00:00:00 2001 From: John Starks Date: Tue, 10 Mar 2026 03:17:22 +0000 Subject: [PATCH 6/6] fix --- vm/devices/virtio/virtio/src/transport/mmio.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/vm/devices/virtio/virtio/src/transport/mmio.rs b/vm/devices/virtio/virtio/src/transport/mmio.rs index d46249c8ab..6e987a84a7 100644 --- a/vm/devices/virtio/virtio/src/transport/mmio.rs +++ b/vm/devices/virtio/virtio/src/transport/mmio.rs @@ -55,6 +55,7 @@ pub struct VirtioMmioDevice { queues: Vec, device_status: VirtioDeviceStatus, disabling: bool, + #[inspect(skip)] poll_waker: Option, config_generation: u32, #[inspect(skip)]