Add iterator and accessor to EventManager

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

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")]

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());
 }