vsock: Enable live migrations (snapshot-restore)

No device state is stored. 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 <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/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 {

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},
@@ -73,8 +77,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 +149,10 @@ pub(crate) enum Error {
     EmptyRawPktsQueue,
     #[error("CID already in use by another vsock device")]
     CidAlreadyInUse,
+    #[error("Event virtqueue is empty")]
+    EmptyEventQueue,
+    #[error("Failed to write to event virtqueue")]
+    EventQueueWrite,
 }
 
 impl std::convert::From<Error> for std::io::Error {
@@ -150,6 +161,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 +278,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 +304,7 @@ impl VhostUserVsockBackend {
             queues_per_thread,
             exit_consumer,
             exit_notifier,
+            transport_reset_state: Arc::new(Mutex::new(None)),
         })
     }
 }
@@ -311,7 +329,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 +356,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 +371,12 @@ 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 matches!(*state, TransportResetState::Pending) {
+                        thread.reset_transport(vring_evt, evt_idx)?;
+                        *state = TransportResetState::Sent;
+                    }
+                }
             }
             BACKEND_EVENT => {
                 thread.process_backend_evt(evset);
@@ -398,6 +424,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 +449,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)]
@@ -455,7 +510,7 @@ mod tests {
 
         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);

vhost-device-vsock/src/vhu_vsock_thread.rs

@@ -3,7 +3,7 @@
 use std::{
     collections::{HashMap, HashSet},
     fs::File,
-    io::{self, BufRead, BufReader},
+    io::{self, BufRead, BufReader, Write},
     iter::FromIterator,
     num::Wrapping,
     ops::Deref,
@@ -20,7 +20,7 @@ use std::{
 
 use log::{error, warn};
 use vhost_user_backend::{VringEpollHandler, VringRwLock, VringT};
-use virtio_queue::QueueOwnedT;
+use virtio_queue::{QueueOwnedT, QueueT};
 use virtio_vsock::packet::{VsockPacket, PKT_HEADER_SIZE};
 use vm_memory::{GuestAddressSpace, GuestMemoryAtomic, GuestMemoryMmap};
 use vmm_sys_util::{
@@ -35,7 +35,7 @@ use crate::{
     thread_backend::*,
     vhu_vsock::{
         BackendType, CidMap, ConnMapKey, Error, Result, VhostUserVsockBackend, BACKEND_EVENT,
-        SIBLING_VM_EVENT, VSOCK_HOST_CID,
+        SIBLING_VM_EVENT, VSOCK_EVENT_TRANSPORT_RESET, VSOCK_HOST_CID,
     },
     vsock_conn::*,
 };
@@ -89,6 +89,9 @@ pub(crate) struct VhostUserVsockThread {
     /// Used to alternate between the RX queues to prevent the starvation of one
     /// by the other.
     last_processed: RxQueueType,
+    /// Backend and sibling vm connections will be dropped until this is set to true, typically
+    /// when communication with the driver has been established.
+    active: bool,
 }
 
 impl VhostUserVsockThread {
@@ -100,6 +103,9 @@ impl VhostUserVsockThread {
         groups: Vec<String>,
         cid_map: Arc<RwLock<CidMap>>,
     ) -> Result<Self> {
+        if cid_map.read().unwrap().contains_key(&guest_cid) {
+            return Err(Error::CidAlreadyInUse);
+        }
         let mut host_listeners_map = HashMap::new();
         match &backend_info {
             BackendType::UnixDomainSocket(uds_path) => {
@@ -138,25 +144,10 @@ impl VhostUserVsockThread {
             epoll_fd,
             guest_cid,
             tx_buffer_size,
-            groups_set.clone(),
+            groups_set,
             cid_map.clone(),
         );
 
-        {
-            let mut cid_map = cid_map.write().unwrap();
-            if cid_map.contains_key(&guest_cid) {
-                return Err(Error::CidAlreadyInUse);
-            }
-
-            cid_map.insert(
-                guest_cid,
-                (
-                    thread_backend.raw_pkts_queue.clone(),
-                    groups_set,
-                    sibling_event_fd.try_clone().unwrap(),
-                ),
-            );
-        }
         let (sender, receiver) = mpsc::channel::<EventData>();
         thread::spawn(move || loop {
             // TODO: Understand why doing the following in the background thread works.
@@ -180,6 +171,7 @@ impl VhostUserVsockThread {
             tx_buffer_size,
             sibling_event_fd,
             last_processed: RxQueueType::Standard,
+            active: false,
         };
 
         for host_raw_fd in thread.host_listeners_map.keys() {
@@ -259,6 +251,27 @@ impl VhostUserVsockThread {
         Ok(())
     }
 
+    /// Start processing requests from the host and sibling VMs
+    pub fn activate(&mut self) -> Result<()> {
+        self.active = true;
+
+        // Add the guest cid to the map to become accessible to sibling VMs
+        let mut cid_map = self.thread_backend.cid_map.write().unwrap();
+        if cid_map.contains_key(&self.guest_cid) {
+            return Err(Error::CidAlreadyInUse);
+        }
+        cid_map.insert(
+            self.guest_cid,
+            (
+                self.thread_backend.raw_pkts_queue.clone(),
+                self.thread_backend.groups_set.clone(),
+                self.sibling_event_fd.try_clone().unwrap(),
+            ),
+        );
+
+        Ok(())
+    }
+
     /// Return raw file descriptor of the epoll file.
     fn get_epoll_fd(&self) -> RawFd {
         self.epoll_file.as_raw_fd()
@@ -310,7 +323,7 @@ impl VhostUserVsockThread {
             match listener {
                 ListenerType::Unix(unix_listener) => {
                     let conn = unix_listener.accept().map_err(Error::UnixAccept);
-                    if self.mem.is_some() {
+                    if self.active && self.mem.is_some() {
                         conn.and_then(|(stream, _)| {
                             stream
                                 .set_nonblocking(true)
@@ -748,14 +761,20 @@ impl VhostUserVsockThread {
                 }
             };
 
-            if self.thread_backend.send_pkt(&pkt).is_err() {
-                vring
-                    .get_mut()
-                    .get_queue_mut()
-                    .iter(mem)
-                    .unwrap()
-                    .go_to_previous_position();
-                break;
+            if self.active {
+                if self.thread_backend.send_pkt(&pkt).is_err() {
+                    vring
+                        .get_mut()
+                        .get_queue_mut()
+                        .iter(mem)
+                        .unwrap()
+                        .go_to_previous_position();
+                    break;
+                }
+            } else {
+                // A tx packet can only arrive before we're active if it was sent before a
+                // transport reset completed.
+                warn!("vsock: Dropping packet received before reset completes");
             }
 
             // TODO: Check if the protocol requires read length to be correct
@@ -796,6 +815,39 @@ impl VhostUserVsockThread {
         }
         Ok(())
     }
+
+    pub fn reset_transport(&mut self, vring: &VringRwLock, event_idx: bool) -> Result<()> {
+        let atomic_mem = match &self.mem {
+            Some(m) => m,
+            None => return Err(Error::NoMemoryConfigured),
+        };
+
+        let mut vring_mut = vring.get_mut();
+        let queue = vring_mut.get_queue_mut();
+
+        let avail_desc = queue
+            .pop_descriptor_chain(atomic_mem.memory())
+            .ok_or(Error::EmptyEventQueue)?;
+
+        let mem = atomic_mem.clone().memory();
+        let head_idx = avail_desc.head_index();
+        let mut writer = avail_desc
+            .writer(&mem)
+            .map_err(|_| Error::EventQueueWrite)?;
+        writer
+            .write_all(&VSOCK_EVENT_TRANSPORT_RESET.to_le_bytes())
+            .map_err(|_| Error::EventQueueWrite)?;
+        self.sender
+            .send(EventData {
+                vring: vring.clone(),
+                event_idx,
+                head_idx,
+                used_len: std::mem::size_of::<u32>(),
+            })
+            .unwrap();
+
+        Ok(())
+    }
 }
 
 impl Drop for VhostUserVsockThread {