Skip to content

virtio: convert device disable to async poll-based pattern#2850

Open
jstarks wants to merge 6 commits intomicrosoft:mainfrom
jstarks:poll_disable
Open

virtio: convert device disable to async poll-based pattern#2850
jstarks wants to merge 6 commits intomicrosoft:mainfrom
jstarks:poll_disable

Conversation

@jstarks
Copy link
Member

@jstarks jstarks commented Feb 27, 2026

Summary

Convert the VirtioDevice::disable() method from a synchronous fire-and-forget pattern to an async poll-based poll_disable(), enabling transports to properly track device shutdown and report completion to the guest.

Motivation

The previous disable() implementation had each device spawn a detached task to stop its workers. These detached tasks could outlive the device and there was no way for the transport to know when shutdown was actually complete. The guest saw status=0 immediately regardless of whether workers had finished draining.

Changes

task_control: add poll_stop() method — Extracts the polling logic from stop() into a new poll_stop() that returns Poll<bool>, then rewrites stop() as poll_fn(|cx| self.poll_stop(cx)).await. This gives devices a non-async way to drive worker shutdown from poll_disable(). Includes unit tests.

virtio: change VirtioDevice::disable() to poll_disable() — Changes the trait method signature to fn poll_disable(&mut self, cx: &mut Context<'_>) -> Poll<()>. Device implementations (virtio_net, virtiofs, virtio_pmem, virtio_p9, TestDevice) now use TaskControl::poll_stop() directly instead of spawning detached tasks. Removes unused pal_async and futures dependencies from several crates.

virtio: implement PollDevice for async disable in transports — Both MMIO and PCI transports now implement PollDevice to drive pending disables via the chipset polling infrastructure:

  • On guest status=0 write: initiates poll_disable with a noop waker
  • Fast path: if the device completes synchronously, clears status immediately
  • Slow path: sets a disabling flag and returns — the guest sees non-zero status until PollDevice::poll_device drives it to completion
  • Implements ChangeDeviceState::reset() (previously a TODO) to async-drive poll_disable and reset all transport state

virtio: remove Drop impls from transports — The Drop-based poll_disable with a noop waker could initiate but never drive shutdown to completion. Removed in favor of proper async teardown via reset().

Files changed

Area Files
task_control support/task_control/src/lib.rs
virtio core vm/devices/virtio/virtio/src/common.rs, vm/devices/virtio/virtio/src/tests.rs
transports vm/devices/virtio/virtio/src/transport/mmio.rs, vm/devices/virtio/virtio/src/transport/pci.rs
devices vm/devices/virtio/virtio_net/src/lib.rs, vm/devices/virtio/virtio_p9/src/lib.rs, vm/devices/virtio/virtio_pmem/src/lib.rs, vm/devices/virtio/virtiofs/src/virtio.rs

@github-actions github-actions bot added the unsafe Related to unsafe code label Feb 27, 2026
@github-actions
Copy link

⚠️ Unsafe Code Detected

This PR modifies files containing unsafe Rust code. Extra scrutiny is required during review.

For more on why we check whole files, instead of just diffs, check out the Rustonomicon

@jstarks jstarks marked this pull request as ready for review March 2, 2026 04:01
@jstarks jstarks requested review from a team as code owners March 2, 2026 04:01
Copilot AI review requested due to automatic review settings March 2, 2026 04:01
Copy link
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR refactors virtio device shutdown from a synchronous disable() to a poll-driven poll_disable() so transports can keep guest-visible device status non-zero until device shutdown is actually complete, and it adds a TaskControl::poll_stop() helper to enable non-async shutdown driving.

Changes:

  • Add TaskControl::poll_stop() (and refactor stop() to use it), including unit tests.
  • Replace VirtioDevice::disable() with VirtioDevice::poll_disable() and update multiple virtio device implementations accordingly.
  • Update virtio MMIO/PCI transports to asynchronously drive disables via PollDevice, and implement ChangeDeviceState::reset() to fully reset transport/device state.

Reviewed changes

Copilot reviewed 12 out of 13 changed files in this pull request and generated 6 comments.

Show a summary per file
File Description
support/task_control/src/lib.rs Introduces poll_stop() and rewrites stop() to use poll-based driving; adds tests.
vm/devices/virtio/virtio/src/common.rs Updates the VirtioDevice trait to poll_disable() with documentation.
vm/devices/virtio/virtio/src/transport/mmio.rs Adds disable state tracking and PollDevice integration for async reset/disable.
vm/devices/virtio/virtio/src/transport/pci.rs Adds disable state tracking and PollDevice integration for async reset/disable.
vm/devices/virtio/virtio/src/tests.rs Updates the test device implementation to the new poll_disable() contract.
vm/devices/virtio/virtio_net/src/lib.rs Converts device shutdown hook to poll_disable() (currently synchronous completion).
vm/devices/virtio/virtio_p9/src/lib.rs Converts shutdown logic to poll_disable() and removes detached-task shutdown.
vm/devices/virtio/virtio_p9/Cargo.toml Removes unused async/spawn dependency.
vm/devices/virtio/virtio_pmem/src/lib.rs Converts shutdown logic to poll_disable() and removes detached-task shutdown.
vm/devices/virtio/virtio_pmem/Cargo.toml Removes unused async/spawn dependency.
vm/devices/virtio/virtiofs/src/virtio.rs Converts shutdown logic to poll_disable() and removes detached-task shutdown.
vm/devices/virtio/virtiofs/Cargo.toml Removes unused futures/spawn dependencies.
Cargo.lock Reflects dependency removals from the affected crates.

Comment on lines +125 to +130
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(())
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Previous disable() signaled exit_event before stopping the queue worker. VirtioQueueWorker uses this exit_event to break out of its select_biased! loop; without notifying it, poll_stop() will stop via cancellation, which may drop in-flight queue work instead of letting the worker exit/drain cleanly.

Notify self.exit_event.notify(usize::MAX) when beginning poll_disable() (ideally once per disable) before polling worker.poll_stop(cx).

Copilot uses AI. Check for mistakes.
Comment on lines +131 to +136
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(())
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Previous disable() notified exit_event before stopping the worker. Since VirtioQueueWorker uses the exit_event to exit its run loop, stopping without notifying can cancel in-flight processing rather than allowing a graceful drain/exit.

Notify self.exit_event.notify(usize::MAX) when beginning poll_disable() (ideally once per disable) before polling worker.poll_stop(cx).

Copilot uses AI. Check for mistakes.
Comment on lines +845 to +850
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::<Vec<_>>();
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(())
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TestDevice’s poll_disable() no longer notifies exit_event before stopping queue workers. VirtioQueueWorker uses exit_event to exit its run loop; stopping without notifying changes behavior to cancellation-based shutdown and can drop in-flight work, diverging from the documented “drain any in-flight IO” semantics.

Notify self.exit_event.notify(usize::MAX) when beginning poll_disable() (ideally once) before polling the workers to stop.

Copilot uses AI. Check for mistakes.
Comment on lines 359 to +366
if started {
self.doorbells.clear();
self.device.disable();
let waker = std::task::Waker::noop();
let mut cx = std::task::Context::from_waker(&waker);
if self.device.poll_disable(&mut cx).is_pending() {
self.disabling = true;
return;
}
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

poll_disable() is first polled with Waker::noop(). Any pending device shutdown that registers this waker (e.g., TaskControl::poll_stop storing outer_waker) will never wake the chipset poll loop, and PollDevice is only polled again when its associated waker is woken. This can leave the transport stuck in disabling=true with guest status never reaching 0.

Initiate disables by setting disabling=true and waking a real poll waker captured from PollDevice::poll_device, then drive device.poll_disable(cx) from poll_device using the provided context waker (or otherwise arrange to register a non-noop waker before returning Pending).

Copilot uses AI. Check for mistakes.
Comment on lines +415 to +418
if self.device.poll_disable(&mut cx).is_pending() {
self.disabling = true;
return;
}
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same issue as MMIO: calling device.poll_disable() with Waker::noop() can cause pending shutdown to register a noop outer_waker, so completion won't wake the chipset polling infrastructure (which only re-polls on wake). This can wedge the device in disabling=true and prevent the guest from ever observing status=0.

Prefer: mark disabling=true, wake a real poll waker captured from PollDevice::poll_device, and do the actual poll_disable(cx) from poll_device (or ensure a real waker is registered before returning Pending).

Suggested change
if self.device.poll_disable(&mut cx).is_pending() {
self.disabling = true;
return;
}
let _ = self.device.poll_disable(&mut cx);

Copilot uses AI. Check for mistakes.
Comment on lines +158 to +163
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(())
Copy link

Copilot AI Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VirtioQueueWorker’s run loop waits on the provided exit_event (see VirtioQueueWorker::run_queue), and the previous disable() implementation notified it before stopping workers. With the new poll_disable(), workers are stopped via poll_stop() but the exit_event is never notified, which can cancel in-flight queue processing instead of letting workers observe the reset and drain/exit cleanly.

Consider notifying self.exit_event.notify(usize::MAX) when starting a disable (ideally once per disable) before polling worker shutdown.

Copilot uses AI. Check for mistakes.
Brian-Perkins
Brian-Perkins previously approved these changes Mar 9, 2026
jstarks added 5 commits March 10, 2026 02:22
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().
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.
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.
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

unsafe Related to unsafe code

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants