Add iterator and accessor to EventManager

coverage_config_x86_64.json

@@ -1,5 +1,5 @@
 {
-  "coverage_score": 85.6,
+  "coverage_score": 87.0,
   "exclude_path": "tests/,benches/,utilities/",
   "crate_features": "remote_endpoint,test_utilities"
 }

src/lib.rs

@@ -168,6 +168,9 @@ pub trait SubscriberOps {
     /// Removes the subscriber corresponding to `subscriber_id` from the watch list.
     fn remove_subscriber(&mut self, subscriber_id: SubscriberId) -> Result<Self::Subscriber>;
 
+    /// Returns a reference to the subscriber corresponding to `subscriber_id`.
+    fn subscriber_ref(&self, subscriber_id: SubscriberId) -> Option<&Self::Subscriber>;
+
     /// Returns a mutable reference to the subscriber corresponding to `subscriber_id`.
     fn subscriber_mut(&mut self, subscriber_id: SubscriberId) -> Result<&mut Self::Subscriber>;
 
@@ -276,3 +279,22 @@ impl<T: MutEventSubscriber + ?Sized> MutEventSubscriber for Box<T> {
         self.deref_mut().init(ops);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_subscriber_id_derives() {
+        let a = SubscriberId(1);
+        let b = SubscriberId(1);
+        let c = SubscriberId(2);
+
+        assert_eq!(a, b);
+        assert_ne!(a, c);
+        assert_ne!(c, b);
+
+        let d = c.clone();
+        assert_eq!(c, d);
+    }
+}

src/manager.rs

@@ -1,6 +1,7 @@
 // Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
 
+use std::collections::hash_map::{Iter, IterMut};
 use std::mem::size_of;
 
 use vmm_sys_util::epoll::EpollEvent;
@@ -56,6 +57,11 @@ impl<T: MutEventSubscriber> SubscriberOps for EventManager<T> {
         Ok(subscriber)
     }
 
+    /// Return a reference to the subscriber associated with the provided id.
+    fn subscriber_ref(&self, subscriber_id: SubscriberId) -> Option<&T> {
+        self.subscribers.get_subscriber(subscriber_id)
+    }
+
     /// Return a mutable reference to the subscriber associated with the provided id.
     fn subscriber_mut(&mut self, subscriber_id: SubscriberId) -> Result<&mut T> {
         if self.subscribers.contains(subscriber_id) {
@@ -181,6 +187,16 @@ impl<S: MutEventSubscriber> EventManager<S> {
         #[cfg(feature = "remote_endpoint")]
         self.dispatch_endpoint_event(endpoint_event);
     }
+
+    /// An iterator visiting all subscribers in arbitrary order, with immutable references.
+    pub fn iter(&mut self) -> Iter<'_, SubscriberId, S> {
+        self.subscribers.iter()
+    }
+
+    /// An iterator visiting all subscribers in arbitrary order, with mutable references.
+    pub fn iter_mut(&mut self) -> IterMut<'_, SubscriberId, S> {
+        self.subscribers.iter_mut()
+    }
 }
 
 #[cfg(feature = "remote_endpoint")]
@@ -481,4 +497,31 @@ mod tests {
             Error::InvalidId
         );
     }
+
+    #[test]
+    #[cfg(feature = "remote_endpoint")]
+    fn test_endpoint() {
+        use std::thread;
+
+        let mut event_manager = EventManager::<DummySubscriber>::new().unwrap();
+        let dummy = DummySubscriber::new();
+        let endpoint = event_manager.remote_endpoint();
+        let kicker = event_manager.remote_endpoint();
+
+        let thread_handle = thread::spawn(move || {
+            event_manager.run().unwrap();
+            event_manager.run().unwrap();
+        });
+
+        dummy.event_fd_1.write(1).unwrap();
+
+        let token = endpoint
+            .call_blocking(|sub_ops| -> Result<SubscriberId> { Ok(sub_ops.add_subscriber(dummy)) })
+            .unwrap();
+        assert_eq!(token, SubscriberId(1));
+
+        kicker.kick().unwrap();
+
+        thread_handle.join().unwrap();
+    }
 }

src/subscribers.rs

@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
 
 use super::SubscriberId;
-use std::collections::HashMap;
+use std::collections::hash_map::{HashMap, Iter, IterMut};
 
 // Internal structure used to keep the set of subscribers registered with an EventManger.
 // This structure is a thin wrapper over a `HashMap` in which the keys are uniquely
@@ -51,4 +51,19 @@ impl<T> Subscribers<T> {
     pub(crate) fn get_mut_unchecked(&mut self, subscriber_id: SubscriberId) -> &mut T {
         self.subscribers.get_mut(&subscriber_id).unwrap()
     }
+
+    // Return a reference to the subriber represented by `subscriber_id`.
+    pub(crate) fn get_subscriber(&self, subscriber_id: SubscriberId) -> Option<&T> {
+        self.subscribers.get(&subscriber_id)
+    }
+
+    /// An iterator visiting all subscribers in arbitrary order, with immutable references.
+    pub(crate) fn iter(&mut self) -> Iter<'_, SubscriberId, T> {
+        self.subscribers.iter()
+    }
+
+    /// An iterator visiting all subscribers in arbitrary order, with mutable references.
+    pub(crate) fn iter_mut(&mut self) -> IterMut<'_, SubscriberId, T> {
+        self.subscribers.iter_mut()
+    }
 }

tests/basic_event_manager.rs

@@ -35,6 +35,15 @@ impl App {
         let _ = self.event_manager.run_with_timeout(100);
     }
 
+    fn iter(&mut self) {
+        for (k, _) in self.event_manager.iter_mut() {
+            assert!(self.subscribers_id.contains(k));
+        }
+        for (k, _) in self.event_manager.iter() {
+            assert!(self.subscribers_id.contains(k));
+        }
+    }
+
     fn inject_events_for(&mut self, subscriber_index: &[usize]) {
         for i in subscriber_index {
             let subscriber = self
@@ -58,7 +67,7 @@ impl App {
     fn get_counters(&mut self) -> Vec<u64> {
         let mut result = Vec::<u64>::new();
         for subscriber_id in &self.subscribers_id {
-            let subscriber = self.event_manager.subscriber_mut(*subscriber_id).unwrap();
+            let subscriber = self.event_manager.subscriber_ref(*subscriber_id).unwrap();
             result.push(subscriber.counter());
         }
 
@@ -71,6 +80,7 @@ impl App {
         for id in &self.subscribers_id {
             let _ = self.event_manager.remove_subscriber(*id);
         }
+        self.subscribers_id.clear();
     }
 }
 
@@ -91,6 +101,7 @@ fn test_single_threaded() {
     let triggered_subscribers: Vec<usize> = vec![1, 3, 50, 97];
     app.inject_events_for(&triggered_subscribers);
     app.run();
+    app.iter();
 
     let counters = app.get_counters();
     for i in 0..100 {
@@ -104,8 +115,13 @@ fn test_single_threaded() {
         assert_eq!(counters[i], 1 & (triggered_subscribers.contains(&i) as u64));
     }
 
+    let id = app.subscribers_id[0];
+
     // Once the app does not need events anymore, the cleanup needs to be called.
     // This is particularly important when the app continues the execution, but event monitoring
     // is not necessary.
     app.cleanup();
+
+    assert!(app.event_manager.subscriber_ref(id).is_none());
+    assert!(app.event_manager.subscriber_mut(id).is_err());
 }