Aggregate chains by network_id (#234)

* Aggregate chains by network_id

* Fix network_id handling

backend/src/aggregator.rs

@@ -10,6 +10,7 @@ use crate::types::{NodeDetails, NodeId};
 
 pub struct Aggregator {
     labels: HashMap<Label, ChainId>,
+    networks: HashMap<Label, ChainId>,
     chains: DenseMap<ChainEntry>,
     feeds: DenseMap<Addr<FeedConnector>>,
     serializer: FeedMessageSerializer,
@@ -18,6 +19,7 @@ pub struct Aggregator {
 pub struct ChainEntry {
     addr: Addr<Chain>,
     label: Label,
+    network_id: Option<Label>,
     nodes: usize,
 }
 
@@ -25,6 +27,7 @@ impl Aggregator {
     pub fn new() -> Self {
         Aggregator {
             labels: HashMap::new(),
+            networks: HashMap::new(),
             chains: DenseMap::new(),
             feeds: DenseMap::new(),
             serializer: FeedMessageSerializer::new(),
@@ -33,33 +36,52 @@ impl Aggregator {
 
     /// Get an address to the chain actor by name. If the address is not found,
     /// or the address is disconnected (actor dropped), create a new one.
-    pub fn lazy_chain(&mut self, label: Label, ctx: &mut <Self as Actor>::Context) -> &mut ChainEntry {
-        let (cid, found) = self.labels
-            .get(&label)
-            .map(|&cid| (cid, true))
-            .unwrap_or_else(|| {
+    pub fn lazy_chain(
+        &mut self,
+        label: &str,
+        network: &Option<Label>,
+        ctx: &mut <Self as Actor>::Context,
+    ) -> ChainId {
+        let cid = match self.get_chain_id(label, network.as_ref()) {
+            Some(cid) => cid,
+            None => {
                 self.serializer.push(feed::AddedChain(&label, 1));
 
                 let addr = ctx.address();
-                let label = label.clone();
-                let cid = self.chains.add_with(move |cid| {
+                let label: Label = label.into();
+                let cid = self.chains.add_with(|cid| {
                     ChainEntry {
                         addr: Chain::new(cid, addr, label.clone()).start(),
-                        label,
+                        label: label.clone(),
+                        network_id: network.clone(),
                         nodes: 1,
                     }
                 });
 
+                self.labels.insert(label, cid);
+
+                if let Some(network) = network {
+                    self.networks.insert(network.clone(), cid);
+                }
+
                 self.broadcast();
 
-                (cid, false)
-            });
+                cid
+            }
+        };
 
-        if !found {
-            self.labels.insert(label, cid);
+        cid
+    }
+
+    fn get_chain_id(&self, label: &str, network: Option<&Label>) -> Option<ChainId> {
+        let labels = &self.labels;
+        let networks = &self.networks;
+
+        if let Some(network) = network {
+            networks.get(&**network).or_else(|| labels.get(label)).copied()
+        } else {
+            labels.get(label).copied()
         }
-
-        self.chains.get_mut(cid).expect("Entry just created above; qed")
     }
 
     fn get_chain(&mut self, label: &str) -> Option<&mut ChainEntry> {
@@ -84,13 +106,16 @@ impl Actor for Aggregator {
 #[derive(Message)]
 pub struct AddNode {
     pub node: NodeDetails,
-    pub chain: Label,
+    pub network_id: Option<Label>,
     pub rec: Recipient<Initialize>,
 }
 
 /// Message sent from the Chain to the Aggregator when the Chain loses all nodes
 #[derive(Message)]
-pub struct DropChain(pub Label);
+pub struct DropChain(pub ChainId);
+
+#[derive(Message)]
+pub struct RenameChain(pub ChainId, pub Label);
 
 /// Message sent from the FeedConnector to the Aggregator when subscribing to a new chain
 pub struct Subscribe {
@@ -146,9 +171,20 @@ impl Handler<AddNode> for Aggregator {
     type Result = ();
 
     fn handle(&mut self, msg: AddNode, ctx: &mut Self::Context) {
-        let AddNode { node, chain, rec } = msg;
+        let AddNode { node, network_id, rec } = msg;
 
-        self.lazy_chain(chain, ctx).addr.do_send(chain::AddNode {
+        let cid = self.lazy_chain(&node.chain, &network_id, ctx);
+        let chain = self.chains.get_mut(cid).expect("Entry just created above; qed");
+
+        if let Some(network_id) = network_id {
+            // Attach network id to the chain if it was not done yet
+            if chain.network_id.is_none() {
+                chain.network_id = Some(network_id.clone());
+                self.networks.insert(network_id, cid);
+            }
+        }
+
+        chain.addr.do_send(chain::AddNode {
             node,
             rec,
         });
@@ -159,15 +195,47 @@ impl Handler<DropChain> for Aggregator {
     type Result = ();
 
     fn handle(&mut self, msg: DropChain, _: &mut Self::Context) {
-        let DropChain(label) = msg;
+        let DropChain(cid) = msg;
 
-        if let Some(cid) = self.labels.remove(&label) {
-            self.chains.remove(cid);
-            self.serializer.push(feed::RemovedChain(&label));
+        if let Some(entry) = self.chains.remove(cid) {
+            let label = &entry.label;
+            self.labels.remove(label);
+            if let Some(network) = entry.network_id {
+                self.networks.remove(&network);
+            }
+
+            self.serializer.push(feed::RemovedChain(label));
+            info!("Dropped chain [{}] from the aggregator", label);
             self.broadcast();
         }
 
-        info!("Dropped chain [{}] from the aggregator", label);
+    }
+}
+
+impl Handler<RenameChain> for Aggregator {
+    type Result = ();
+
+    fn handle(&mut self, msg: RenameChain, _: &mut Self::Context) {
+        let RenameChain(cid, new) = msg;
+
+        if let Some(entry) = self.chains.get_mut(cid) {
+            if entry.label == new {
+                return;
+            }
+
+            // Update UI
+            self.serializer.push(feed::RemovedChain(&entry.label));
+            self.serializer.push(feed::AddedChain(&new, entry.nodes));
+
+            // Update labels -> cid map
+            self.labels.remove(&entry.label);
+            self.labels.insert(new.clone(), cid);
+
+            // Update entry
+            entry.label = new;
+
+            self.broadcast();
+        }
     }
 }