virtio-devices: Introduce WorkerThreads handle

Take the recent thread refactor further. Bundle a device's worker
JoinHandles together with the kill event that stops them into a single
WorkerThreads value, owned by VirtioCommon. Its Drop signals the workers
to exit, unparks any parked for migration, and joins them.

This makes a detached/leaked worker unrepresentable. reset(),
wait_for_epoll_threads() and VhostUserCommon::shutdown() now happen when
dropping the WorkerThreads, and the unpark-before-join teardown now
lives in one place.

Assisted-by: Claude:Opus-4.8
Signed-off-by: Dylan Reid <dgreid@fb.com>
Signed-off-by: Rob Bradford <rbradford@meta.com>

virtio-devices/src/balloon.rs

@@ -583,16 +583,6 @@ impl Balloon {
     }
 }
 
-impl Drop for Balloon {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Balloon {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/block.rs

@@ -1035,16 +1035,6 @@ impl Block {
     }
 }
 
-impl Drop for Block {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Block {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/console.rs

@@ -690,16 +690,6 @@ impl Console {
     }
 }
 
-impl Drop for Console {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Console {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/device.rs

@@ -225,17 +225,70 @@ pub trait DmaRemapping {
     ) -> std::result::Result<u64, std::io::Error>;
 }
 
+/// Owns a device's worker threads plus the kill event that stops them.
+///
+/// Dropping signals every worker to exit, unparks any that are parked, and joins them.
+pub struct WorkerThreads {
+    /// shared kill event, a single write wakes all of them.
+    kill_evt: EventFd,
+    // true if the device is paused.
+    paused: Arc<AtomicBool>,
+    // The running worker thread's handles.
+    threads: Vec<thread::JoinHandle<()>>,
+}
+
+impl WorkerThreads {
+    fn new(kill_evt: EventFd, paused: Arc<AtomicBool>) -> Self {
+        WorkerThreads {
+            kill_evt,
+            paused,
+            threads: Vec::new(),
+        }
+    }
+
+    /// Borrow access to the kill eventfd
+    fn kill_evt(&self) -> &EventFd {
+        &self.kill_evt
+    }
+
+    /// Signal the workers to exit without joining; they are joined later when
+    /// this is dropped.
+    pub(crate) fn signal_exit(&self) -> std::io::Result<()> {
+        self.kill_evt.write(1)
+    }
+
+    /// Unpark every worker so threads parked while paused resume their loop.
+    fn unpark(&self) {
+        for t in &self.threads {
+            t.thread().unpark();
+        }
+    }
+}
+
+impl Drop for WorkerThreads {
+    fn drop(&mut self) {
+        // Signal the workers to exit, wake any parked so they observe it, then join.
+        let _ = self.kill_evt.write(1);
+        self.paused.store(false, Ordering::SeqCst);
+        self.unpark();
+        for t in self.threads.drain(..) {
+            if let Err(e) = t.join() {
+                error!("Error joining thread: {e:?}");
+            }
+        }
+    }
+}
+
 /// Structure to handle device state common to all devices
 #[derive(Default)]
 pub struct VirtioCommon {
     pub avail_features: u64,
     pub acked_features: u64,
-    pub kill_evt: Option<EventFd>,
     pub interrupt_cb: Option<Arc<dyn VirtioInterrupt>>,
     pub pause_evt: Option<EventFd>,
     pub paused: Arc<AtomicBool>,
     pub paused_sync: Option<Arc<Barrier>>,
-    pub epoll_threads: Option<Vec<thread::JoinHandle<()>>>,
+    pub workers: Option<WorkerThreads>,
     pub queue_sizes: Vec<u16>,
     pub queue_evts: Vec<EventFd>,
     pub device_type: u32,
@@ -289,7 +342,9 @@ impl VirtioCommon {
             error!("failed creating kill EventFd: {e}");
             ActivateError::BadActivate
         })?;
-        self.kill_evt = Some(kill_evt);
+        // Create the worker collection up front so it owns the kill event;
+        // handlers clone it via dup_eventfds() before any worker is spawned.
+        self.workers = Some(WorkerThreads::new(kill_evt, self.paused.clone()));
 
         let pause_evt = EventFd::new(EFD_NONBLOCK).map_err(|e| {
             error!("failed creating pause EventFd: {e}");
@@ -306,35 +361,21 @@ impl VirtioCommon {
 
     pub fn reset(&mut self) {
         self.queue_evts.clear();
+        self.pause_evt = None;
 
-        // Resume the virtio thread if it was paused. Reset must always
+        // Clear paused explicitly; the workers' Drop does so too, but only
-        // converge to fresh state, so a resume failure is logged but doesn't
+        // when they exist, and reset may run before activate().
-        // skip the rest of the teardown.
+        self.paused.store(false, Ordering::SeqCst);
-        if self.pause_evt.take().is_some()
-            && let Err(e) = self.resume()
-        {
-            error!("Failed to resume paused device during reset: {e:?}");
-        }
 
-        if let Some(kill_evt) = self.kill_evt.take() {
+        // Dropping the workers signals kill_evt, unparks any thread parked
-            // Ignore the result because there is nothing we can do about it.
+        // for migration, and joins them.
-            let _ = kill_evt.write(1);
+        self.workers = None;
-        }
-
-        if let Some(mut threads) = self.epoll_threads.take() {
-            for t in threads.drain(..) {
-                if let Err(e) = t.join() {
-                    error!("Error joining thread: {e:?}");
-                }
-            }
-        }
 
         // Drop the interrupt callback clone
         self.interrupt_cb = None;
     }
 
-    /// Spawn a worker, push its handle into `self.epoll_threads`, and on
+    /// Spawn a worker; on failure, reset the device to join prior workers.
-    /// spawn failure run `self.reset()` to join prior workers.
     #[expect(clippy::too_many_arguments)]
     pub fn spawn_worker<F>(
         &mut self,
@@ -349,18 +390,23 @@ impl VirtioCommon {
     where
         F: FnOnce() -> Result<(), EpollHelperError> + Send + 'static,
     {
-        let mut threads = self.epoll_threads.take().unwrap_or_default();
+        // Scope the borrow of `workers` so it ends before the reset() below.
-        let res = spawn_virtio_thread(
+        let res = {
-            name,
+            let Some(workers) = self.workers.as_mut() else {
-            seccomp_action,
+                error!("spawn_worker called before activate()");
-            thread_type,
+                return Err(ActivateError::BadActivate);
-            &mut threads,
+            };
-            exit_evt,
+            spawn_virtio_thread(
-            device_status,
+                name,
-            interrupt_cb,
+                seccomp_action,
-            f,
+                thread_type,
-        );
+                &mut workers.threads,
-        self.epoll_threads = Some(threads);
+                exit_evt,
+                device_status,
+                interrupt_cb,
+                f,
+            )
+        };
         if let Err(e) = res {
             self.reset();
             return Err(e);
@@ -376,20 +422,19 @@ impl VirtioCommon {
         }
     }
 
-    // Wait for the worker thread to finish and return
+    // Dropping the workers signals, unparks, and joins them. Idempotent.
     pub fn wait_for_epoll_threads(&mut self) {
-        if let Some(mut threads) = self.epoll_threads.take() {
+        self.workers = None;
-            for t in threads.drain(..) {
-                if let Err(e) = t.join() {
-                    error!("Error joining thread: {e:?}");
-                }
-            }
-        }
     }
 
     pub fn dup_eventfds(&self) -> (EventFd, EventFd) {
         (
-            self.kill_evt.as_ref().unwrap().try_clone().unwrap(),
+            self.workers
+                .as_ref()
+                .unwrap()
+                .kill_evt()
+                .try_clone()
+                .unwrap(),
             self.pause_evt.as_ref().unwrap().try_clone().unwrap(),
         )
     }
@@ -446,10 +491,8 @@ impl Pausable for VirtioCommon {
             VirtioDeviceType::from(self.device_type)
         );
         self.paused.store(false, Ordering::SeqCst);
-        if let Some(epoll_threads) = &self.epoll_threads {
+        if let Some(workers) = &self.workers {
-            for t in epoll_threads.iter() {
+            workers.unpark();
-                t.thread().unpark();
-            }
         }
 
         // Signal each activated queue eventfd so workers process restored queues
@@ -495,19 +538,23 @@ mod unit_tests {
         }
     }
 
-    fn make_common_with_kill_evt() -> (VirtioCommon, EventFd) {
+    /// VirtioCommon with its worker collection created (as activate() does)
+    /// and a kill_evt clone for the spawned worker to watch.
+    fn make_common_with_workers() -> (VirtioCommon, EventFd) {
         let kill_evt = EventFd::new(EFD_NONBLOCK).unwrap();
         let kill_evt_clone = kill_evt.try_clone().unwrap();
+        let common = VirtioCommon::default();
+        let workers = WorkerThreads::new(kill_evt, common.paused.clone());
         let common = VirtioCommon {
-            kill_evt: Some(kill_evt),
+            workers: Some(workers),
-            ..Default::default()
+            ..common
         };
         (common, kill_evt_clone)
     }
 
     #[test]
-    fn spawn_worker_appends_to_epoll_threads() {
+    fn spawn_worker_appends_to_workers() {
-        let (mut common, kill_evt_clone) = make_common_with_kill_evt();
+        let (mut common, kill_evt_clone) = make_common_with_workers();
         let started = Arc::new(AtomicUsize::new(0));
         let started_clone = started.clone();
 
@@ -530,15 +577,60 @@ mod unit_tests {
             )
             .unwrap();
 
-        let threads = common.epoll_threads.as_ref().expect("epoll_threads set");
+        let workers = common.workers.as_ref().expect("workers set");
-        assert_eq!(threads.len(), 1);
+        assert_eq!(workers.threads.len(), 1);
 
-        // reset() joins the worker; this confirms the spawn_worker
+        // reset() drops the WorkerThreads and joins the worker, exercising
-        // -> reset() chain (used on spawn failure) actually drains
+        // the spawn-failure cleanup path on a real running thread.
-        // a real, running worker rather than dropping it detached.
         common.reset();
-        assert!(common.epoll_threads.is_none());
+        assert!(common.workers.is_none());
-        assert!(common.kill_evt.is_none());
         assert_eq!(started.load(Ordering::SeqCst), 1);
     }
+
+    #[test]
+    fn dropping_common_joins_workers() {
+        let (mut common, kill_evt_clone) = make_common_with_workers();
+        let started = Arc::new(AtomicUsize::new(0));
+        let started_clone = started.clone();
+
+        let exit_evt = EventFd::new(EFD_NONBLOCK).unwrap();
+        let status = Arc::new(AtomicU8::new(0));
+
+        common
+            .spawn_worker(
+                "test",
+                &SeccompAction::Allow,
+                Thread::VirtioBlock,
+                &exit_evt,
+                status,
+                Arc::new(NoopInterrupt),
+                move || {
+                    started_clone.fetch_add(1, Ordering::SeqCst);
+                    let _ = kill_evt_clone.read();
+                    Ok(())
+                },
+            )
+            .unwrap();
+
+        // Dropping `common` alone must join the worker via WorkerThreads' Drop.
+        drop(common);
+        assert_eq!(started.load(Ordering::SeqCst), 1);
+    }
+
+    #[test]
+    fn reset_clears_paused_without_workers() {
+        // reset() before any worker was spawned must still clear paused, or
+        // the next activation's workers would park immediately and never run.
+        let mut common = VirtioCommon {
+            pause_evt: Some(EventFd::new(EFD_NONBLOCK).unwrap()),
+            ..Default::default()
+        };
+        common.paused.store(true, Ordering::SeqCst);
+        assert!(common.workers.is_none());
+
+        common.reset();
+
+        assert!(!common.paused.load(Ordering::SeqCst));
+        assert!(common.pause_evt.is_none());
+    }
 }

virtio-devices/src/iommu.rs

@@ -1252,16 +1252,6 @@ impl Iommu {
     }
 }
 
-impl Drop for Iommu {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Iommu {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/mem.rs

@@ -922,16 +922,6 @@ impl Mem {
     }
 }
 
-impl Drop for Mem {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Mem {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/net.rs

@@ -644,17 +644,6 @@ impl Net {
     }
 }
 
-impl Drop for Net {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        // Needed to ensure all references to tap FDs are dropped (#4868)
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Net {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/pmem.rs

@@ -349,16 +349,6 @@ impl Pmem {
     }
 }
 
-impl Drop for Pmem {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Pmem {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/rng.rs

@@ -228,16 +228,6 @@ impl Rng {
     }
 }
 
-impl Drop for Rng {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Rng {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/rtc.rs

@@ -599,16 +599,6 @@ impl Rtc {
     }
 }
 
-impl Drop for Rtc {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl VirtioDevice for Rtc {
     fn device_type(&self) -> u32 {
         self.common.device_type

virtio-devices/src/vhost_user/mod.rs

@@ -613,19 +613,8 @@ impl VhostUserCommon {
     }
 
     pub fn shutdown(&mut self) {
-        // Signal workers to exit, unpause them (they may be parked
+        // Join the workers so they drop their Arc<VhostUserHandle> and the
-        // if the VM was paused for migration), then wait for them
+        // socket closes, letting the migration destination reconnect.
-        // to finish so they drop their Arc<VhostUserHandle> and the
-        // socket fully closes for the destination to reconnect.
-        if let Some(kill_evt) = self.virtio_common.kill_evt.take() {
-            let _ = kill_evt.write(1);
-        }
-        self.virtio_common.paused.store(false, Ordering::SeqCst);
-        if let Some(threads) = self.virtio_common.epoll_threads.as_ref() {
-            for t in threads {
-                t.thread().unpark();
-            }
-        }
         self.virtio_common.wait_for_epoll_threads();
 
         // Remove socket path if needed
@@ -862,8 +851,8 @@ impl VhostUserCommon {
 
         // Make sure the device thread is killed in order to prevent from
         // reconnections to the socket.
-        if let Some(kill_evt) = self.virtio_common.kill_evt.take() {
+        if let Some(workers) = self.virtio_common.workers.as_ref() {
-            kill_evt.write(1).map_err(|e| {
+            workers.signal_exit().map_err(|e| {
                 MigratableError::CompleteMigration(anyhow!(
                     "Error killing vhost-user thread: {e:?}"
                 ))

virtio-devices/src/vsock/device.rs

@@ -426,19 +426,6 @@ where
     }
 }
 
-impl<B> Drop for Vsock<B>
-where
-    B: VsockBackend,
-{
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 impl<B> VirtioDevice for Vsock<B>
 where
     B: VsockBackend + Sync + 'static,

virtio-devices/src/watchdog.rs

@@ -270,16 +270,6 @@ impl Watchdog {
     }
 }
 
-impl Drop for Watchdog {
-    fn drop(&mut self) {
-        if let Some(kill_evt) = self.common.kill_evt.take() {
-            // Ignore the result because there is nothing we can do about it.
-            let _ = kill_evt.write(1);
-        }
-        self.common.wait_for_epoll_threads();
-    }
-}
-
 fn timerfd_create() -> Result<RawFd, io::Error> {
     // SAFETY: FFI call, trivially safe
     let res = unsafe { libc::timerfd_create(libc::CLOCK_MONOTONIC, 0) };