vsock: Enable live migrations (snapshot-restore)

jemoreira · jemoreira · commit d8fce2402eae · 2026-02-12T16:08:22.000-08:00
No device state is stored during the migration. When a buffer appears on
the event queue for the first time after the frontend requested a load
state operation the backend sends a VIRTIO_VSOCK_EVENT_TRANSPORT_RESET to
the driver and waits for the driver to read the config again (the signal
that the reset completed) before starting normal operations.

Signed-off-by: Jorge E. Moreira &lt;jemoreira@google.com&gt;
diff --git a/vhost-device-vsock/CHANGELOG.md b/vhost-device-vsock/CHANGELOG.md
@@ -3,6 +3,8 @@
 
 ### Added
 
+- [#936](https://github.com/rust-vmm/vhost-device/pull/936) vsock: Enable live migrations (snapshot-restore)
+
 ### Changed
 
 ### Fixed
diff --git a/vhost-device-vsock/src/thread_backend.rs b/vhost-device-vsock/src/thread_backend.rs
@@ -217,7 +217,7 @@ pub(crate) struct VsockThreadBackend {
     pub raw_pkts_queue: Arc<RwLock<RawPktsQ>>,
     /// Set of groups assigned to the device which it is allowed to communicate
     /// with.
-    groups_set: Arc<RwLock<HashSet<String>>>,
+    pub groups_set: Arc<RwLock<HashSet<String>>>,
 }
 
 impl VsockThreadBackend {
@@ -332,7 +332,7 @@ impl VsockThreadBackend {
             if dst_cid != VSOCK_HOST_CID {
                 let cid_map = self.cid_map.read().unwrap();
                 if cid_map.contains_key(&dst_cid) {
-                    let (sibling_raw_pkts_queue, sibling_groups_set, sibling_event_fd) =
+                    let (sibling_raw_pkts_queue_opt, sibling_groups_set, sibling_event_fd) =
                         cid_map.get(&dst_cid).unwrap();
 
                     if self
@@ -345,11 +345,18 @@ impl VsockThreadBackend {
                         return Ok(());
                     }
 
-                    sibling_raw_pkts_queue
-                        .write()
-                        .unwrap()
-                        .push_back(RawVsockPacket::from_vsock_packet(pkt)?);
-                    let _ = sibling_event_fd.write(1);
+                    match sibling_raw_pkts_queue_opt {
+                        Some(queue) => {
+                            queue
+                                .write()
+                                .unwrap()
+                                .push_back(RawVsockPacket::from_vsock_packet(pkt)?);
+                            let _ = sibling_event_fd.write(1);
+                        }
+                        None => {
+                            info!("vsock: dropping packet for cid: {dst_cid:?} due to inactive device");
+                        }
+                    }
                 } else {
                     warn!("vsock: dropping packet for unknown cid: {dst_cid:?}");
                 }
@@ -525,6 +532,7 @@ mod tests {
     #[cfg(feature = "backend_vsock")]
     use crate::vhu_vsock::VsockProxyInfo;
     use crate::vhu_vsock::{BackendType, VhostUserVsockBackend, VsockConfig, VSOCK_OP_RW};
+    use vhost_user_backend::VhostUserBackend;
 
     const DATA_LEN: usize = 16;
     const CONN_TX_BUF_SIZE: u32 = 64 * 1024;
@@ -698,11 +706,28 @@ mod tests {
         // SAFETY: Safe as hdr_raw and data_raw are guaranteed to be valid.
         let mut packet = unsafe { VsockPacket::new(hdr_raw, Some(data_raw)).unwrap() };
 
+        packet.set_type(VSOCK_TYPE_STREAM);
+        packet.set_src_cid(CID);
+        packet.set_dst_cid(SIBLING_CID);
+        packet.set_dst_port(SIBLING_LISTENING_PORT);
+        packet.set_op(VSOCK_OP_RW);
+        packet.set_len(DATA_LEN as u32);
+        packet
+            .data_slice()
+            .unwrap()
+            .copy_from(&[0x01u8, 0x12u8, 0x23u8, 0x34u8]);
+
+        // The packet will be dropped silently because the thread won't activate until the config
+        // is read.
+        vtp.send_pkt(&packet).unwrap();
         assert_eq!(
             vtp.recv_raw_pkt(&mut packet).unwrap_err().to_string(),
             Error::EmptyRawPktsQueue.to_string()
         );
 
+        sibling_backend.get_config(0, 8);
+        sibling2_backend.get_config(0, 8);
+
         packet.set_type(VSOCK_TYPE_STREAM);
         packet.set_src_cid(CID);
         packet.set_dst_cid(SIBLING_CID);
diff --git a/vhost-device-vsock/src/vhu_vsock.rs b/vhost-device-vsock/src/vhu_vsock.rs
@@ -2,14 +2,18 @@
 
 use std::{
     collections::{HashMap, HashSet},
+    fs::File,
     io::Result as IoResult,
     path::PathBuf,
     sync::{Arc, Mutex, RwLock},
 };
 
 use log::warn;
 use thiserror::Error as ThisError;
-use vhost::vhost_user::message::{VhostUserProtocolFeatures, VhostUserVirtioFeatures};
+use vhost::vhost_user::message::{
+    VhostTransferStateDirection, VhostTransferStatePhase, VhostUserProtocolFeatures,
+    VhostUserVirtioFeatures,
+};
 use vhost_user_backend::{VhostUserBackend, VringRwLock, VringT};
 use virtio_bindings::bindings::{
     virtio_config::{VIRTIO_F_NOTIFY_ON_EMPTY, VIRTIO_F_VERSION_1},
@@ -25,8 +29,14 @@ use vmm_sys_util::{
 
 use crate::{thread_backend::RawPktsQ, vhu_vsock_thread::*};
 
-pub(crate) type CidMap =
-    HashMap<u64, (Arc<RwLock<RawPktsQ>>, Arc<RwLock<HashSet<String>>>, EventFd)>;
+pub(crate) type CidMap = HashMap<
+    u64,
+    (
+        Option<Arc<RwLock<RawPktsQ>>>,
+        Arc<RwLock<HashSet<String>>>,
+        EventFd,
+    ),
+>;
 
 const NUM_QUEUES: usize = 3;
 
@@ -73,8 +83,11 @@ pub(crate) const VSOCK_FLAGS_SHUTDOWN_RCV: u32 = 1;
 /// data
 pub(crate) const VSOCK_FLAGS_SHUTDOWN_SEND: u32 = 2;
 
+/// Vsock events - `VSOCK_EVENT_TRANSPORT_RESET`: Communication has been interrupted
+pub(crate) const VSOCK_EVENT_TRANSPORT_RESET: u32 = 0;
+
 // Queue mask to select vrings.
-const QUEUE_MASK: u64 = 0b11;
+const QUEUE_MASK: u64 = 0b111;
 
 pub(crate) type Result<T> = std::result::Result<T, Error>;
 
@@ -142,6 +155,8 @@ pub(crate) enum Error {
     EmptyRawPktsQueue,
     #[error("CID already in use by another vsock device")]
     CidAlreadyInUse,
+    #[error("Failed to write to event virtqueue")]
+    EventQueueWrite,
 }
 
 impl std::convert::From<Error> for std::io::Error {
@@ -150,6 +165,11 @@ impl std::convert::From<Error> for std::io::Error {
     }
 }
 
+enum TransportResetState {
+    Pending,
+    Sent,
+}
+
 #[cfg(feature = "backend_vsock")]
 #[derive(Debug, PartialEq, Clone)]
 pub(crate) struct VsockProxyInfo {
@@ -262,6 +282,7 @@ pub(crate) struct VhostUserVsockBackend {
     queues_per_thread: Vec<u64>,
     exit_consumer: EventConsumer,
     exit_notifier: EventNotifier,
+    transport_reset_state: Arc<Mutex<Option<TransportResetState>>>,
 }
 
 impl VhostUserVsockBackend {
@@ -287,6 +308,7 @@ impl VhostUserVsockBackend {
             queues_per_thread,
             exit_consumer,
             exit_notifier,
+            transport_reset_state: Arc::new(Mutex::new(None)),
         })
     }
 }
@@ -311,7 +333,9 @@ impl VhostUserBackend for VhostUserVsockBackend {
     }
 
     fn protocol_features(&self) -> VhostUserProtocolFeatures {
-        VhostUserProtocolFeatures::MQ | VhostUserProtocolFeatures::CONFIG
+        VhostUserProtocolFeatures::MQ
+            | VhostUserProtocolFeatures::CONFIG
+            | VhostUserProtocolFeatures::DEVICE_STATE
     }
 
     fn set_event_idx(&self, enabled: bool) {
@@ -336,6 +360,7 @@ impl VhostUserBackend for VhostUserVsockBackend {
     ) -> IoResult<()> {
         let vring_rx = &vrings[0];
         let vring_tx = &vrings[1];
+        let vring_evt = &vrings[2];
 
         if evset != EventSet::IN {
             return Err(Error::HandleEventNotEpollIn.into());
@@ -350,7 +375,15 @@ impl VhostUserBackend for VhostUserVsockBackend {
                 thread.process_tx(vring_tx, evt_idx)?;
             }
             EVT_QUEUE_EVENT => {
-                warn!("Received an unexpected EVT_QUEUE_EVENT");
+                if let Some(state) = &mut *self.transport_reset_state.lock().unwrap() {
+                    if vring_evt.get_ref().get_queue().ready() {
+                        if matches!(*state, TransportResetState::Pending)
+                            && thread.reset_transport(vring_evt, evt_idx)?
+                        {
+                            *state = TransportResetState::Sent;
+                        }
+                    }
+                }
             }
             BACKEND_EVENT => {
                 thread.process_backend_evt(evset);
@@ -398,6 +431,19 @@ impl VhostUserBackend for VhostUserVsockBackend {
             return Vec::new();
         }
 
+        if offset + size == buf.len() {
+            // The last byte of the config is read when the driver is initializing or after it has
+            // processed a transport reset event. Either way, no transport reset will be pending
+            // after this.
+            *self.transport_reset_state.lock().unwrap() = None;
+            // Activate all threads once it's known a pending reset event is no longer pending.
+            for thread in self.threads.iter() {
+                if let Err(e) = thread.lock().unwrap().activate() {
+                    warn!("Failed to activate thread: {}", e);
+                }
+            }
+        }
+
         buf[offset..offset + size].to_vec()
     }
 
@@ -410,6 +456,22 @@ impl VhostUserBackend for VhostUserVsockBackend {
         let notifier = self.exit_notifier.try_clone().ok()?;
         Some((consumer, notifier))
     }
+
+    fn set_device_state_fd(
+        &self,
+        direction: VhostTransferStateDirection,
+        _phase: VhostTransferStatePhase,
+        _file: File,
+    ) -> std::result::Result<Option<File>, std::io::Error> {
+        if let VhostTransferStateDirection::LOAD = direction {
+            *self.transport_reset_state.lock().unwrap() = Some(TransportResetState::Pending);
+        }
+        Ok(None)
+    }
+
+    fn check_device_state(&self) -> std::result::Result<(), std::io::Error> {
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -445,17 +507,20 @@ mod tests {
         let vrings = [
             VringRwLock::new(mem.clone(), 0x1000).unwrap(),
             VringRwLock::new(mem.clone(), 0x2000).unwrap(),
+            VringRwLock::new(mem.clone(), 0x1000).unwrap(),
         ];
         vrings[0].set_queue_info(0x100, 0x200, 0x300).unwrap();
         vrings[0].set_queue_ready(true);
         vrings[1].set_queue_info(0x1100, 0x1200, 0x1300).unwrap();
         vrings[1].set_queue_ready(true);
+        vrings[2].set_queue_info(0x2100, 0x2200, 0x2300).unwrap();
+        vrings[2].set_queue_ready(true);
 
         backend.update_memory(mem).unwrap();
 
         let queues_per_thread = backend.queues_per_thread();
         assert_eq!(queues_per_thread.len(), 1);
-        assert_eq!(queues_per_thread[0], 0b11);
+        assert_eq!(queues_per_thread[0], 0b111);
 
         let config = backend.get_config(0, 8);
         assert_eq!(config.len(), 8);
@@ -578,6 +643,7 @@ mod tests {
         let vrings = [
             VringRwLock::new(mem.clone(), 0x1000).unwrap(),
             VringRwLock::new(mem.clone(), 0x2000).unwrap(),
+            VringRwLock::new(mem.clone(), 0x1000).unwrap(),
         ];
 
         backend.update_memory(mem).unwrap();
diff --git a/vhost-device-vsock/src/vhu_vsock_thread.rs b/vhost-device-vsock/src/vhu_vsock_thread.rs
