vsock: Enable live migrations (snapshot-restore)

The [virtio
spec](https://docs.oasis-open.org/virtio/virtio/v1.3/csd01/virtio-v1.3-csd01.html#x1-4950001)
mandates that when restoring from a snapshot (such as when the VM was
migrated to a different host) the device must send a TRANSPORT RESET
event via the event virtqueue to the driver. The driver in turn, upon
receiving the event, must drop all existing connections while keeping
all listeners and read the CID again from the device configuration
space and update the listeners with the (potentially different) CID.

On the device side care must be taken to ensure no new connections are
established until that transport reset event is handled otherwise the
guest will silently drop them and the peer may not know about it. Given
that the driver must read the CID from the configuration space during
normal boot and after doing a reset that's the best signal for the
device to "activate" and start processing packets from the host and
sibling VMs. Feature negotiation also happens during both normal boot
and restore, but in the case of restore it happens before the transport
reset and therefore can't be used as reliable signal to "activate" the
device.

The device doesn't need to save/load any state for this. When loading a
snapshot the device simply notes that it must send a TRANSPORT RESET
event to the driver as soon as possible, which it then does when it
receives a kick on the event vring. Independently of whether the device
is started to restore a previous VM state or brand new (the device
actually doesn't know until it `set_device_state_fd` is called), the
device always starts in "inactive" state, meaning it will drop any
packets coming from any source. As mentioned before the device
"activates" once the driver has read the configuration space.

Unlike other VMMs, QEMU takes ownership of the event vring and handles
sending of the TRANSPORT RESET event itself. This change attempts to handle
both approaches by sending the TRANSPORT RESET event if the vring is
kicked, but doesn't treat it as precondition to activate the device,
instead choosing to activate unconditionally once the driver reads the
configuration space.

There is a race in this implementation, that occurs when the snapshot
was taken before the guest driver read the configuration and then it
reads immediately upon restore, before the transport reset event was
handled. While this could be mitigated by waiting for the reset event to
be sent AND the config to be read AFTER the event was sent, QEMU's
decision to take over the event queue makes that approach unviable.
Hopefully, it will be very unlikely a snapshot will be taken before the
driver is fully initialized as that would be of very little value in
practice.

Signed-off-by: Jorge E. Moreira <jemoreira@google.com>

Author: jemoreira

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

vhost-device-vsock/README.md

@@ -129,6 +129,18 @@ In this configuration:
 - The host must have vsock support (e.g., `vsock_loopback` kernel module loaded)
 - For testing, you can load the module with: `modprobe vsock_loopback`
 
+## Live migration
+
+This device implementation advertises support for live migrations by offering the VHOST_USER_PROTOCOL_F_DEVICE_STATE protocol feature, however this doesn't work with Qemu yet as it marks its vsock frontend as "unmigratable". This feature does work with CrosVm and potentially other virtual machine managers.
+
+The device itself doesn't save or restore any state during a live migration. It relies instead on the frontend to save the vring's states and negotiated features. It also expects the the frontend to "kick" the queues that have pending buffers in it since the driver probably kicked those queues before the migration and won't do it again.
+
+The state saving flow is trivial as the device doesn't save any state as mentioned.
+
+The state loading flow is a bit more complicated because the virtio-vsock spec mandates that the device must send a VIRTIO_VSOCK_EVENT_TRANSPORT_RESET event to the driver. During a restore the backend is started no differently than during a regular boot. When the frontend sends the VHOST_USER_SET_DEVICE_STATE_FD command with LOAD direction the backend doesn't load anything, but it takes note that a transport reset event needs to be sent to the driver via the event vring when possible. In order to make sure this event is sent when the queue is ready, the backend waits for the event queue to be kicked before sending the event. While these kicks usually come from the driver, this particular one is actually sent by the vhost-user frontend. This implementation depends on the frontend to kick all queues with pending buffers after a restore because the driver is unlikely to do so as it probably did it before the snapshot was taken.
+
+In response to the transport reset event the driver drops any existing connections and reads the configuration space again. To prevent the driver from dropping any new connections established after the restore the backend doesn't forward any packets from outside the VM to the driver until it has read the configuration space. In fact, because the backend doesn't know at start time whether this is a restore or a clean boot, it always waits until after the driver has read the configuration space to start forwarding packets between the outside world and the driver.
+
 ## Usage
 
 Run the vhost-device-vsock device with unix domain socket backend:

vhost-device-vsock/src/thread_backend.rs

@@ -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);

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};
 use virtio_bindings::bindings::{
     virtio_config::{VIRTIO_F_NOTIFY_ON_EMPTY, VIRTIO_F_VERSION_1},
@@ -24,8 +28,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;
 
@@ -72,8 +82,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>;
 
@@ -141,6 +154,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 {
@@ -261,6 +276,7 @@ pub(crate) struct VhostUserVsockBackend {
     queues_per_thread: Vec<u64>,
     exit_consumer: EventConsumer,
     exit_notifier: EventNotifier,
+    transport_reset_pending: Arc<Mutex<bool>>,
 }
 
 impl VhostUserVsockBackend {
@@ -286,6 +302,7 @@ impl VhostUserVsockBackend {
             queues_per_thread,
             exit_consumer,
             exit_notifier,
+            transport_reset_pending: Arc::new(Mutex::new(false)),
         })
     }
 }
@@ -310,7 +327,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) {
@@ -335,6 +354,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());
@@ -349,7 +369,11 @@ impl VhostUserBackend for VhostUserVsockBackend {
                 thread.process_tx(vring_tx, evt_idx)?;
             }
             EVT_QUEUE_EVENT => {
-                warn!("Received an unexpected EVT_QUEUE_EVENT");
+                let reset_pending = &mut *self.transport_reset_pending.lock().unwrap();
+                if *reset_pending {
+                    thread.reset_transport(vring_evt, evt_idx)?;
+                    *reset_pending = false;
+                }
             }
             BACKEND_EVENT => {
                 thread.process_backend_evt(evset);
@@ -389,6 +413,15 @@ 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. Activate all threads once it's known a reset event is not pending.
+            for thread in self.threads.iter() {
+                thread.lock().unwrap().activate();
+            }
+        }
+
         buf[offset..offset + size].to_vec()
     }
 
@@ -401,6 +434,23 @@ 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_pending.lock().unwrap() = true;
+        }
+        Ok(None)
+    }
+
+    fn check_device_state(&self) -> std::result::Result<(), std::io::Error> {
+        // We had nothing to read/write to the fd, so always return Ok.
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -436,17 +486,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);
@@ -569,6 +622,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();