Don't remove all feeds subscribed to a chain when one disconnects (#383)

* Only remove the feed that disconnected to not break the rest...

* use multimap struct to avoid sync issues between feed and chain

* add a remove test, too

* cargo fmt

* fix name of test

* move multimap to common so we can doctest it and add 'unique' to name

* cargo fmt

* Return old key if value moved to make uniqueness more obvious

backend/common/src/lib.rs

@@ -30,6 +30,7 @@ mod dense_map;
 mod either_sink;
 mod mean_list;
 mod most_seen;
+mod multi_map_unique;
 mod num_stats;
 
 // Export a bunch of common bits at the top level for ease of import:
@@ -38,4 +39,5 @@ pub use dense_map::DenseMap;
 pub use either_sink::EitherSink;
 pub use mean_list::MeanList;
 pub use most_seen::MostSeen;
+pub use multi_map_unique::MultiMapUnique;
 pub use num_stats::NumStats;

backend/common/src/multi_map_unique.rs

@@ -0,0 +1,166 @@
+// Source code for the Substrate Telemetry Server.
+// Copyright (C) 2021 Parity Technologies (UK) Ltd.
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+use std::collections::{HashMap, HashSet};
+use std::hash::Hash;
+
+/// A map where each key can contain multiple values. We enforce that a value
+/// only ever belongs to one key at a time (the latest key it was inserted
+/// against).
+pub struct MultiMapUnique<K, V> {
+    value_to_key: HashMap<V, K>,
+    key_to_values: HashMap<K, HashSet<V>>,
+}
+
+impl<K, V> MultiMapUnique<K, V> {
+    /// Construct a new MultiMap
+    pub fn new() -> Self {
+        Self {
+            value_to_key: HashMap::new(),
+            key_to_values: HashMap::new(),
+        }
+    }
+
+    /// Return the set of values associated with a key.
+    pub fn get_values(&self, key: &K) -> Option<&HashSet<V>>
+    where
+        K: Eq + Hash,
+    {
+        self.key_to_values.get(key)
+    }
+
+    /// Remove a value from the MultiMap, returning the key it was found
+    /// under, if it was found at all.
+    ///
+    /// ```
+    /// let mut m = common::MultiMapUnique::new();
+    ///
+    /// m.insert("a", 1);
+    /// m.insert("a", 2);
+    ///
+    /// m.insert("b", 3);
+    /// m.insert("b", 4);
+    ///
+    /// assert_eq!(m.num_keys(), 2);
+    /// assert_eq!(m.num_values(), 4);
+    ///
+    /// m.remove_value(&1);
+    ///
+    /// assert_eq!(m.num_keys(), 2);
+    /// assert_eq!(m.num_values(), 3);
+    ///
+    /// m.remove_value(&2);
+    ///
+    /// assert_eq!(m.num_keys(), 1);
+    /// assert_eq!(m.num_values(), 2);
+    /// ```
+    pub fn remove_value(&mut self, value: &V) -> Option<K>
+    where
+        V: Eq + Hash,
+        K: Eq + Hash,
+    {
+        if let Some(key) = self.value_to_key.remove(value) {
+            if let Some(m) = self.key_to_values.get_mut(&key) {
+                m.remove(value);
+                if m.is_empty() {
+                    self.key_to_values.remove(&key);
+                }
+            }
+            return Some(key);
+        }
+        None
+    }
+
+    /// Insert a key+value pair into the multimap. Multiple different
+    /// values can exist for a single key, but only one of each value can
+    /// exist in the MultiMap.
+    ///
+    /// If a previous value did exist, the old key it was inserted against
+    /// is returned.
+    ///
+    /// ```
+    /// let mut m = common::MultiMapUnique::new();
+    ///
+    /// let old_key = m.insert("a", 1);
+    /// assert_eq!(old_key, None);
+    ///
+    /// let old_key = m.insert("b", 1);
+    /// assert_eq!(old_key, Some("a"));
+    ///
+    /// let old_key = m.insert("c", 1);
+    /// assert_eq!(old_key, Some("b"));
+    ///
+    /// assert_eq!(m.num_keys(), 1);
+    /// assert_eq!(m.num_values(), 1);
+    ///
+    /// // The value `1` must be unique in the map, so it only exists
+    /// // in the last location it was inserted:
+    /// assert!(m.get_values(&"a").is_none());
+    /// assert!(m.get_values(&"b").is_none());
+    /// assert_eq!(m.get_values(&"c").unwrap().iter().collect::<Vec<_>>(), vec![&1]);
+    /// ```
+    pub fn insert(&mut self, key: K, value: V) -> Option<K>
+    where
+        V: Clone + Eq + Hash,
+        K: Clone + Eq + Hash,
+    {
+        // Ensure that the value doesn't exist elsewhere already;
+        // values must be unique and only belong to one key:
+        let old_key = self.remove_value(&value);
+
+        self.value_to_key.insert(value.clone(), key.clone());
+        self.key_to_values.entry(key).or_default().insert(value);
+
+        old_key
+    }
+
+    /// Number of values stored in the map
+    pub fn num_values(&self) -> usize {
+        self.value_to_key.len()
+    }
+
+    /// Number of keys stored in the map
+    pub fn num_keys(&self) -> usize {
+        self.key_to_values.len()
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+
+    #[test]
+    fn multiple_values_allowed_per_key() {
+        let mut m = MultiMapUnique::new();
+
+        m.insert("a", 1);
+        m.insert("a", 2);
+
+        m.insert("b", 3);
+        m.insert("b", 4);
+
+        assert_eq!(m.num_keys(), 2);
+        assert_eq!(m.num_values(), 4);
+
+        let a_vals = m.get_values(&"a").expect("a vals");
+        assert!(a_vals.contains(&1));
+        assert!(a_vals.contains(&2));
+
+        let b_vals = m.get_values(&"b").expect("b vals");
+        assert!(b_vals.contains(&3));
+        assert!(b_vals.contains(&4));
+    }
+}