virtio-devices: Retry complete vhost-user reconnect

A Unix stream connection can succeed while the process owning the
listener is exiting. The resulting connection can then fail during
SET_OWNER, GET_FEATURES, or later reinitialization.

Retry the complete connect and reinitialization transaction when it
returns a transport failure within a 60-second retry window. Each retry
uses a fresh frontend. Connection errors retain their existing handling,
non-transport errors fail immediately, and waits remain interruptible by
the worker kill event.

Signed-off-by: Peter Delevoryas <pdel@meta.com>
Assisted-by: Codex:GPT-5

virtio-devices/src/vhost_user/blk.rs

@@ -64,6 +64,7 @@ impl Blk {
             num_queues as u64,
             false,
             &exit_evt,
+            |_| Ok(()),
         )?;
 
         let (

virtio-devices/src/vhost_user/fs.rs

@@ -89,8 +89,14 @@ impl Fs {
         let num_queues = NUM_QUEUE_OFFSET + req_num_queues;
 
         // Connect to the vhost-user socket.
-        let mut vu =
+        let mut vu = VhostUserHandle::connect_vhost_user(
-            VhostUserHandle::connect_vhost_user(false, path, num_queues as u64, false, &exit_evt)?;
+            false,
+            path,
+            num_queues as u64,
+            false,
+            &exit_evt,
+            |_| Ok(()),
+        )?;
 
         let (
             avail_features,

virtio-devices/src/vhost_user/generic_vhost_user.rs

@@ -76,8 +76,14 @@ impl GenericVhostUser {
         let num_queues = request_queue_sizes.len();
 
         // Connect to the vhost-user socket.
-        let mut vu =
+        let mut vu = VhostUserHandle::connect_vhost_user(
-            VhostUserHandle::connect_vhost_user(false, path, num_queues as u64, false, &exit_evt)?;
+            false,
+            path,
+            num_queues as u64,
+            false,
+            &exit_evt,
+            |_| Ok(()),
+        )?;
 
         let (
             avail_features,

virtio-devices/src/vhost_user/mod.rs

@@ -330,48 +330,43 @@ impl<S: VhostUserFrontendReqHandler> VhostUserEpollHandler<S> {
             epoll::Events::EPOLLHUP,
         )?;
 
+        let queues = self
+            .queues
+            .iter()
+            .map(|(i, q, e)| (*i, vm_virtio::clone_queue(q), e.try_clone().unwrap()))
+            .collect::<Vec<_>>();
+
         let mut vhost_user = match VhostUserHandle::connect_vhost_user(
             self.server,
             &self.socket_path,
             self.queues.len() as u64,
             true,
             &self.kill_evt,
+            |vhost_user| {
+                vhost_user.reinitialize_vhost_user(
+                    self.mem.memory().deref(),
+                    &queues,
+                    self.virtio_interrupt.as_ref(),
+                    self.acked_features,
+                    self.acked_protocol_features,
+                    &self.backend_req_handler,
+                    self.inflight.as_mut(),
+                )
+            },
         ) {
             Ok(vu) => vu,
-            // Kill event fired during the connect retry loop; abandon the
+            // Kill event fired during the reconnect retry loop; abandon the
             // reconnect attempt. The EpollHelper observes the same kill
             // event and will tear down on its next iteration.
             Err(Error::ConnectKilled) => return Ok(()),
             Err(e) => {
                 return Err(EpollHelperError::IoError(io::Error::other(format!(
-                    "failed connecting vhost-user backend for socket {}: {e:?}",
+                    "failed reconnecting vhost-user backend for socket {}: {e:?}",
                     self.socket_path
                 ))));
             }
         };
 
-        let queues = self
-            .queues
-            .iter()
-            .map(|(i, q, e)| (*i, vm_virtio::clone_queue(q), e.try_clone().unwrap()))
-            .collect::<Vec<_>>();
-        // Initialize the backend
-        vhost_user
-            .reinitialize_vhost_user(
-                self.mem.memory().deref(),
-                &queues,
-                self.virtio_interrupt.as_ref(),
-                self.acked_features,
-                self.acked_protocol_features,
-                &self.backend_req_handler,
-                self.inflight.as_mut(),
-            )
-            .map_err(|e| {
-                EpollHelperError::IoError(io::Error::other(format!(
-                    "failed reconnecting vhost-user backend: {e:?}"
-                )))
-            })?;
-
         helper.add_event_custom(
             vhost_user.socket_handle().as_raw_fd(),
             HUP_CONNECTION_EVENT,
@@ -875,3 +870,62 @@ impl VhostUserCommon {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod unit_tests {
+    use std::io::{Read, Write};
+    use std::os::unix::net::{UnixListener, UnixStream};
+    use std::thread;
+
+    use vhost::vhost_user::message::{FrontendReq, VhostUserHeaderFlag};
+    use vmm_sys_util::tempdir::TempDir;
+
+    use super::*;
+
+    fn read_request(stream: &mut UnixStream, expected_request: FrontendReq) {
+        let mut header = [0u8; 12];
+        stream.read_exact(&mut header).unwrap();
+        assert_eq!(
+            u32::from_ne_bytes(header[0..4].try_into().unwrap()),
+            u32::from(expected_request)
+        );
+        assert_eq!(u32::from_ne_bytes(header[8..12].try_into().unwrap()), 0);
+    }
+
+    #[test]
+    fn connect_retries_get_features_disconnect_with_fresh_socket() {
+        let temp_dir = TempDir::new_with_prefix("/tmp/vhost-user-reconnect-").unwrap();
+        let socket_path = temp_dir.as_path().join("backend.sock");
+        let listener = UnixListener::bind(&socket_path).unwrap();
+        let backend = thread::spawn(move || {
+            let (mut first, _) = listener.accept().unwrap();
+            read_request(&mut first, FrontendReq::SET_OWNER);
+            read_request(&mut first, FrontendReq::GET_FEATURES);
+            drop(first);
+
+            let (mut second, _) = listener.accept().unwrap();
+            read_request(&mut second, FrontendReq::SET_OWNER);
+            read_request(&mut second, FrontendReq::GET_FEATURES);
+
+            let mut reply = Vec::with_capacity(20);
+            reply.extend_from_slice(&u32::from(FrontendReq::GET_FEATURES).to_ne_bytes());
+            reply.extend_from_slice(&(VhostUserHeaderFlag::REPLY.bits() | 1).to_ne_bytes());
+            reply.extend_from_slice(&8u32.to_ne_bytes());
+            reply.extend_from_slice(&0u64.to_ne_bytes());
+            second.write_all(&reply).unwrap();
+        });
+
+        let kill_evt = EventFd::new(libc::EFD_NONBLOCK).unwrap();
+        VhostUserHandle::connect_vhost_user(
+            false,
+            socket_path.to_str().unwrap(),
+            1,
+            false,
+            &kill_evt,
+            |_| Ok(()),
+        )
+        .unwrap();
+
+        backend.join().unwrap();
+    }
+}

virtio-devices/src/vhost_user/net.rs

@@ -94,6 +94,7 @@ impl Net {
             num_queues as u64,
             false,
             &exit_evt,
+            |_| Ok(()),
         )?;
 
         let (

virtio-devices/src/vhost_user/vu_common_ctrl.rs

@@ -59,6 +59,7 @@ struct VringInfo {
 #[derive(Clone)]
 pub struct VhostUserHandle {
     vu: Frontend,
+    backend_features: u64,
     ready: bool,
     supports_migration: bool,
     supports_device_state: bool,
@@ -123,16 +124,7 @@ impl VhostUserHandle {
         avail_features: u64,
         avail_protocol_features: VhostUserProtocolFeatures,
     ) -> Result<(u64, u64)> {
-        // Set vhost-user owner.
+        let acked_features = avail_features & self.backend_features;
-        self.vu.set_owner().map_err(Error::VhostUserSetOwner)?;
-
-        // Get features from backend, do negotiation to get a feature collection which
-        // both VMM and backend support.
-        let backend_features = self
-            .vu
-            .get_features()
-            .map_err(Error::VhostUserGetFeatures)?;
-        let acked_features = avail_features & backend_features;
 
         let acked_protocol_features =
             if acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits() != 0 {
@@ -333,11 +325,6 @@ impl VhostUserHandle {
         acked_features: u64,
         acked_protocol_features: u64,
     ) -> Result<()> {
-        self.vu.set_owner().map_err(Error::VhostUserSetOwner)?;
-        self.vu
-            .get_features()
-            .map_err(Error::VhostUserGetFeatures)?;
-
         if acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits() != 0
             && let Some(acked_protocol_features) =
                 VhostUserProtocolFeatures::from_bits(acked_protocol_features)
@@ -418,6 +405,7 @@ impl VhostUserHandle {
         num_queues: u64,
         unlink_socket: bool,
         kill_evt: &EventFd,
+        mut initialize: impl FnMut(&mut Self) -> Result<()>,
     ) -> Result<Self> {
         if server {
             if unlink_socket {
@@ -429,8 +417,9 @@ impl VhostUserHandle {
             info!("Waiting for incoming vhost-user connection...");
             let (stream, _) = listener.accept().map_err(Error::AcceptConnection)?;
 
-            Ok(VhostUserHandle {
+            let mut vhost_user = Self {
                 vu: Frontend::from_stream(stream, num_queues),
+                backend_features: 0,
                 ready: false,
                 supports_migration: false,
                 supports_device_state: false,
@@ -438,7 +427,17 @@ impl VhostUserHandle {
                 acked_features: 0,
                 vrings_info: None,
                 queue_indexes: Vec::new(),
-            })
+            };
+            vhost_user
+                .vu
+                .set_owner()
+                .map_err(Error::VhostUserSetOwner)?;
+            vhost_user.backend_features = vhost_user
+                .vu
+                .get_features()
+                .map_err(Error::VhostUserGetFeatures)?;
+            initialize(&mut vhost_user)?;
+            Ok(vhost_user)
         } else {
             const RETRY_INTERVAL: Duration = Duration::from_millis(100);
             const CONNECT_TIMEOUT: Duration = Duration::from_secs(60);
@@ -475,44 +474,71 @@ impl VhostUserHandle {
             let mut events = [EpollEvent::default(); 1];
 
             loop {
-                let err = match Frontend::connect(socket_path, num_queues) {
+                let connection = Frontend::connect(socket_path, num_queues)
-                    Ok(m) => {
+                    .map(|vu| Self {
-                        return Ok(VhostUserHandle {
+                        vu,
-                            vu: m,
+                        backend_features: 0,
-                            ready: false,
+                        ready: false,
-                            supports_migration: false,
+                        supports_migration: false,
-                            supports_device_state: false,
+                        supports_device_state: false,
-                            shm_log: None,
+                        shm_log: None,
-                            acked_features: 0,
+                        acked_features: 0,
-                            vrings_info: None,
+                        vrings_info: None,
-                            queue_indexes: Vec::new(),
+                        queue_indexes: Vec::new(),
-                        });
+                    })
-                    }
+                    .map_err(Error::VhostUserConnect);
-                    Err(e) => e,
-                };
 
-                let retryable = match &err {
+                let (err, connect_failed) = match connection {
-                    VhostError::VhostUserProtocol(VhostUserError::SocketConnect(io_err)) => {
+                    Ok(mut vhost_user) => match vhost_user
-                        matches!(
+                        .vu
-                            io_err.kind(),
+                        .set_owner()
-                            io::ErrorKind::NotFound
+                        .map_err(Error::VhostUserSetOwner)
-                                | io::ErrorKind::Interrupted
+                        .and_then(|()| {
-                                | io::ErrorKind::ConnectionRefused
+                            vhost_user.backend_features = vhost_user
-                        )
+                                .vu
-                    }
+                                .get_features()
-                    _ => false,
+                                .map_err(Error::VhostUserGetFeatures)?;
-                };
+                            initialize(&mut vhost_user)
+                        }) {
+                        Ok(()) => return Ok(vhost_user),
+                        Err(e) if e.is_transport_lost() => (e, false),
+                        Err(e) => return Err(e),
+                    },
+                    Err(Error::VhostUserConnect(err)) => {
+                        let retryable = match &err {
+                            VhostError::VhostUserProtocol(VhostUserError::SocketConnect(
+                                io_err,
+                            )) => {
+                                matches!(
+                                    io_err.kind(),
+                                    io::ErrorKind::NotFound
+                                        | io::ErrorKind::Interrupted
+                                        | io::ErrorKind::ConnectionRefused
+                                )
+                            }
+                            _ => false,
+                        };
 
-                if !retryable {
+                        if !retryable {
-                    error!(
+                            error!(
-                        "Failed connecting to vhost-user backend for socket {socket_path}: {err:?}"
+                                "Failed connecting to vhost-user backend for socket {socket_path}: {err:?}"
-                    );
+                            );
-                    return Err(Error::VhostUserConnect(err));
+                            return Err(Error::VhostUserConnect(err));
-                }
+                        }
+
+                        (Error::VhostUserConnect(err), true)
+                    }
+                    Err(e) => return Err(e),
+                };
 
                 if start.elapsed() >= CONNECT_TIMEOUT {
-                    error!("Timed out waiting for vhost-user connection on socket {socket_path}");
+                    if connect_failed {
-                    return Err(Error::VhostUserConnectTimeout);
+                        error!(
+                            "Timed out waiting for vhost-user connection on socket {socket_path}"
+                        );
+                        return Err(Error::VhostUserConnectTimeout);
+                    }
+                    return Err(err);
                 }
 
                 loop {