vhost-user-backend: Avoid losing vring kicks

By registering the file descriptors only when the queue is ready and
enabled. This fixes a race in which a buffer could be added to the queue
and the kick sent after registration but before it being enabled which
causes the kick to be lost and the event never delivered to the backend.
This was particularly prevalent when restoring a snapshot.

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

coverage_config_x86_64.json

@@ -1,5 +1,5 @@
 {
-  "coverage_score": 86.22,
+  "coverage_score": 86.82,
   "exclude_path": "vhost/src/vhost_kern/",
   "crate_features": "vhost/vhost-user-frontend,vhost/vhost-user-backend,vhost-user-backend/postcopy"
 }

vhost-user-backend/CHANGELOG.md

@@ -7,6 +7,8 @@
 ### Deprecated
 ### Fixed
 
+- [[#343]](https://github.com/rust-vmm/vhost/pull/343) Avoid losing vring kicks
+
 ## v0.21.0
 
 ### Changed

vhost-user-backend/src/backend.rs

@@ -634,6 +634,10 @@ pub mod tests {
 
             backend
         }
+
+        pub fn events(&self) -> u64 {
+            self.events
+        }
     }
 
     impl VhostUserBackendMut for MockVhostBackend {

vhost-user-backend/src/handler.rs

@@ -204,23 +204,43 @@ where
     }
 
     fn initialize_vring(&self, vring: &T::Vring, index: u8) -> VhostUserResult<()> {
-        assert!(vring.get_ref().get_kick().is_some());
+        vring.set_queue_ready(true);
+        self.update_vring_registration(vring, index)
+    }
 
-        if let Some(fd) = vring.get_ref().get_kick() {
+    /// Adds or removes the vring's kick fd to the epoll instance based on the vring status.
+    /// Ensures that notifications are handled only while the vring is both started and enabled
+    /// and that no notifications are lost.
+    fn update_vring_registration(&self, vring: &T::Vring, index: u8) -> VhostUserResult<()> {
+        let vring_state = vring.get_ref();
+        if let Some(fd) = vring_state.get_kick() {
             for (thread_index, queues_mask) in self.queues_per_thread.iter().enumerate() {
                 let shifted_queues_mask = queues_mask >> index;
                 if shifted_queues_mask & 1u64 == 1u64 {
                     let evt_idx = queues_mask.count_ones() - shifted_queues_mask.count_ones();
-                    self.handlers[thread_index]
+                    if vring_state.get_queue().ready() && vring_state.is_enabled() {
-                        .register_event(fd.as_raw_fd(), EventSet::IN, u64::from(evt_idx))
+                        if let Err(e) = self.handlers[thread_index].register_event(
-                        .map_err(VhostUserError::ReqHandlerError)?;
+                            fd.as_raw_fd(),
+                            EventSet::IN,
+                            u64::from(evt_idx),
+                        ) {
+                            if e.kind() != io::ErrorKind::AlreadyExists {
+                                // This could happen if we're asked by the frontend to enable an
+                                // already enabled queue, don't fail in that case.
+                                return Err(VhostUserError::ReqHandlerError(e));
+                            }
+                        }
+                    } else {
+                        let _ = self.handlers[thread_index].unregister_event(
+                            fd.as_raw_fd(),
+                            EventSet::IN,
+                            u64::from(evt_idx),
+                        );
+                    }
                     break;
                 }
             }
         }
-
-        vring.set_queue_ready(true);
-
         Ok(())
     }
 
@@ -256,8 +276,9 @@ where
 
     fn reset_device(&mut self) -> VhostUserResult<()> {
         // Disable all vrings
-        for vring in self.vrings.iter_mut() {
+        for (index, vring) in self.vrings.iter().enumerate() {
             vring.set_enabled(false);
+            self.update_vring_registration(vring, index as u8)?;
         }
 
         // Reset device state, retain protocol state
@@ -288,8 +309,9 @@ where
         // Note: If `VHOST_USER_F_PROTOCOL_FEATURES` has been negotiated we must leave
         // the vrings in their current state.
         if self.acked_features & VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits() == 0 {
-            for vring in self.vrings.iter_mut() {
+            for (index, vring) in self.vrings.iter().enumerate() {
                 vring.set_enabled(true);
+                self.update_vring_registration(vring, index as u8)?;
             }
         }
 
@@ -431,19 +453,7 @@ where
         // VHOST_USER_SET_VRING_KICK, and stop ring upon receiving
         // VHOST_USER_GET_VRING_BASE.
         vring.set_queue_ready(false);
-
+        self.update_vring_registration(vring, index as u8)?;
-        if let Some(fd) = vring.get_ref().get_kick() {
-            for (thread_index, queues_mask) in self.queues_per_thread.iter().enumerate() {
-                let shifted_queues_mask = queues_mask >> index;
-                if shifted_queues_mask & 1u64 == 1u64 {
-                    let evt_idx = queues_mask.count_ones() - shifted_queues_mask.count_ones();
-                    self.handlers[thread_index]
-                        .unregister_event(fd.as_raw_fd(), EventSet::IN, u64::from(evt_idx))
-                        .map_err(VhostUserError::ReqHandlerError)?;
-                    break;
-                }
-            }
-        }
 
         let next_avail = vring.queue_next_avail();
 
@@ -529,6 +539,7 @@ where
         // or after it has been disabled by VHOST_USER_SET_VRING_ENABLE
         // with parameter 0.
         vring.set_enabled(enable);
+        self.update_vring_registration(vring, index as u8)?;
 
         Ok(())
     }
@@ -788,3 +799,65 @@ impl<T: VhostUserBackend> Drop for VhostUserHandler<T> {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::backend::tests::MockVhostBackend;
+    use std::os::fd::IntoRawFd;
+    use std::os::unix::io::FromRawFd;
+    use std::sync::Mutex;
+    use std::thread;
+    use std::time::Duration;
+    use vhost::vhost_user::message::VhostUserVirtioFeatures;
+    use vm_memory::{GuestAddress, GuestMemoryAtomic, GuestMemoryMmap};
+    use vmm_sys_util::event::{new_event_consumer_and_notifier, EventFlag};
+
+    #[test]
+    fn test_no_lost_kicks() {
+        let mem = GuestMemoryAtomic::new(
+            GuestMemoryMmap::<()>::from_ranges(&[(GuestAddress(0x100000), 0x10000)]).unwrap(),
+        );
+        let backend = Arc::new(Mutex::new(MockVhostBackend::new()));
+        let mut handler = VhostUserHandler::new(backend.clone(), mem.clone()).unwrap();
+        handler
+            .set_features(VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits())
+            .unwrap();
+
+        // Simulate VMM initializing the vring
+        let vring_index = 0;
+
+        let (kick_consumer, notifier) =
+            new_event_consumer_and_notifier(EventFlag::empty()).unwrap();
+        // Safety: we know kick_consumer is valid.
+        let kick_consumer_file = unsafe { File::from_raw_fd(kick_consumer.into_raw_fd()) };
+        handler
+            .set_vring_kick(vring_index as u8, Some(kick_consumer_file))
+            .unwrap();
+
+        // Ring is NOT enabled yet by default (if protocol features negotiated)
+        handler.set_vring_enable(vring_index as u32, false).unwrap();
+
+        // Kick it
+        notifier.notify().unwrap();
+
+        // The worker thread is already running (started in VhostUserHandler::new).
+        // Give it some time to NOT process the event.
+        thread::sleep(Duration::from_millis(200));
+
+        let events = backend.lock().unwrap().events();
+        assert_eq!(
+            events, 0,
+            "Backend should NOT have been kicked while disabled"
+        );
+
+        // Now enable it.
+        handler.set_vring_enable(vring_index as u32, true).unwrap();
+
+        // Give it some time to process the NOW-registered event.
+        thread::sleep(Duration::from_millis(200));
+
+        let events = backend.lock().unwrap().events();
+        assert_eq!(events, 1, "Backend SHOULD have been kicked after enabling");
+    }
+}

vhost-user-backend/src/vring.rs

@@ -129,6 +129,11 @@ impl<M: GuestAddressSpace> VringState<M> {
         })
     }
 
+    /// Whether the vring is enabled or not
+    pub fn is_enabled(&self) -> bool {
+        self.enabled
+    }
+
     /// Get an immutable reference to the underlying raw `Queue` object.
     pub fn get_queue(&self) -> &Queue {
         &self.queue