Give things unique ID types, not aliases, to prevent mixups

backend/common/src/assign_id.rs

@@ -1,65 +1,52 @@
-use std::{fmt::Display, hash::Hash};
+use std::hash::Hash;
-use serde::{Serialize,Deserialize};
 use bimap::BiMap;
 
-#[derive(Clone,Copy,Debug,Hash,PartialEq,Eq,Serialize,Deserialize)]
-pub struct Id(usize);
-
-impl std::convert::From<Id> for usize {
-    fn from(id: Id) -> usize {
-        id.0
-    }
-}
-impl std::convert::From<usize> for Id {
-    fn from(n: usize) -> Id {
-        Id(n)
-    }
-}
-impl Display for Id {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        self.0.fmt(f)
-    }
-}
-
 /// A struct that allows you to assign an ID to an arbitrary set of
 /// details (so long as they are Eq+Hash+Clone), and then access
 /// the assigned ID given those details or access the details given
 /// the ID.
 #[derive(Debug)]
-pub struct AssignId<Details> {
+pub struct AssignId<Id, Details> {
-    current_id: Id,
+    current_id: usize,
-    mapping: BiMap<Id, Details>
+    mapping: BiMap<usize, Details>,
+    _id_type: std::marker::PhantomData<Id>
 }
 
-impl <Details> AssignId<Details> where Details: Eq + Hash {
+impl <Id, Details> AssignId<Id, Details>
+where
+    Details: Eq + Hash,
+    Id: From<usize> + Copy,
+    usize: From<Id>
+{
     pub fn new() -> Self {
         Self {
-            current_id: Id(0),
+            current_id: 0,
-            mapping: BiMap::new()
+            mapping: BiMap::new(),
+            _id_type: std::marker::PhantomData
         }
     }
 
     pub fn assign_id(&mut self, details: Details) -> Id {
         let this_id = self.current_id;
-        self.current_id.0 += 1;
+        self.current_id += 1;
         self.mapping.insert(this_id, details);
-        this_id
+        this_id.into()
     }
 
     pub fn get_details(&mut self, id: Id) -> Option<&Details> {
-        self.mapping.get_by_left(&id)
+        self.mapping.get_by_left(&id.into())
     }
 
     pub fn get_id(&mut self, details: &Details) -> Option<Id> {
-        self.mapping.get_by_right(details).map(|id| *id)
+        self.mapping.get_by_right(details).map(|&id| id.into())
     }
 
     pub fn remove_by_id(&mut self, id: Id) -> Option<Details> {
-        self.mapping.remove_by_left(&id).map(|(_,details)| details)
+        self.mapping.remove_by_left(&id.into()).map(|(_,details)| details)
     }
 
     pub fn remove_by_details(&mut self, details: &Details) -> Option<Id> {
-        self.mapping.remove_by_right(&details).map(|(id,_)| id)
+        self.mapping.remove_by_right(&details).map(|(id,_)| id.into())
     }
 
     pub fn clear(&mut self) {
@@ -67,6 +54,6 @@ impl <Details> AssignId<Details> where Details: Eq + Hash {
     }
 
     pub fn iter(&self) -> impl Iterator<Item = (Id, &Details)> {
-        self.mapping.iter().map(|(id, details)| (*id, details))
+        self.mapping.iter().map(|(&id, details)| (id.into(), details))
     }
 }

backend/common/src/id_type.rs

@@ -0,0 +1,56 @@
+/// Define a type that can be used as an ID, be converted from/to the inner type,
+/// and serialized/deserialized transparently into the inner type.
+#[macro_export]
+macro_rules! id_type {
+    ($( #[$attrs:meta] )* $vis:vis $ty:ident ( $inner:ident ) $(;)? ) => {
+        #[derive(Debug,Clone,Copy,PartialEq,Eq,Hash)]
+        $( #[$attrs] )*
+        $vis struct $ty($inner);
+
+        impl $ty {
+            #[allow(dead_code)]
+            pub fn new(inner: $inner) -> Self {
+                Self(inner)
+            }
+        }
+
+        impl From<$inner> for $ty {
+            fn from(inner: $inner) -> Self {
+                Self(inner)
+            }
+        }
+
+        impl From<$ty> for $inner {
+            fn from(ty: $ty) -> Self {
+                ty.0
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    #[test]
+    fn create_and_use_new_id_type() {
+        id_type!{
+            Foo(usize)
+        };
+        let _ = Foo::new(123);
+        let id = Foo::from(123);
+        let _: usize = id.into();
+
+        // Check that these don't lead to compile errors:
+        id_type!{
+            Bar(usize);
+        };
+        id_type!{
+            pub Wibble(u64)
+        };
+        id_type!{
+            /// We can have doc strings, too
+            pub(crate) Wobble(u16)
+        };
+    }
+
+}

backend/common/src/internal_messages.rs

@@ -2,12 +2,16 @@ use std::net::IpAddr;
 
 use crate::node::Payload;
 use crate::types::{NodeDetails, BlockHash};
-use crate::assign_id::Id;
+use crate::id_type;
 use serde::{Deserialize, Serialize};
 
-/// The shard-local ID of a given node, where a single connection
+id_type! {
-/// might send data on behalf of more than one chain.
+    /// The shard-local ID of a given node, where a single connection
-pub type LocalId = Id;
+    /// might send data on behalf of more than one chain.
+    #[derive(serde::Serialize, serde::Deserialize)]
+    pub ShardNodeId(usize);
+}
+
 
 /// Message sent from the shard to the backend core
 #[derive(Deserialize, Serialize, Debug, Clone)]
@@ -16,17 +20,17 @@ pub enum FromShardAggregator {
     AddNode {
     	ip: Option<IpAddr>,
     	node: NodeDetails,
-    	local_id: LocalId,
+    	local_id: ShardNodeId,
         genesis_hash: BlockHash
     },
     /// Send a message payload to update details for a node
     UpdateNode {
-        local_id: LocalId,
+        local_id: ShardNodeId,
         payload: Payload,
     },
     /// Inform the core that a node has been removed
     RemoveNode {
-        local_id: LocalId
+        local_id: ShardNodeId
     }
 }
 
@@ -34,7 +38,7 @@ pub enum FromShardAggregator {
 #[derive(Deserialize, Serialize, Debug, Clone)]
 pub enum FromTelemetryCore {
 	Mute {
-		local_id: LocalId,
+		local_id: ShardNodeId,
         reason: MuteReason
 	}
 }

backend/common/src/lib.rs

@@ -6,3 +6,4 @@ pub mod json;
 pub mod log_level;
 pub mod assign_id;
 pub mod most_seen;
+pub mod id_type;

backend/common/src/types.rs

@@ -4,10 +4,8 @@ use serde::{Deserialize, Serialize};
 use crate::util::{now, MeanList};
 use crate::json;
 
-pub type NodeId = usize;
 pub type BlockNumber = u64;
 pub type Timestamp = u64;
-pub type Address = Box<str>;
 pub use primitive_types::H256 as BlockHash;
 
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -106,21 +104,6 @@ pub struct NodeLocation {
     pub city: Box<str>,
 }
 
-// impl Serialize for NodeDetails {
-//     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-//     where
-//         S: Serializer,
-//     {
-//         let mut tup = serializer.serialize_tuple(6)?;
-//         tup.serialize_element(&self.name)?;
-//         tup.serialize_element(&self.implementation)?;
-//         tup.serialize_element(&self.version)?;
-//         tup.serialize_element(&self.validator)?;
-//         tup.serialize_element(&self.network_id)?;
-//         tup.end()
-//     }
-// }
-
 impl Serialize for NodeStats {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where

backend/common/src/util.rs

@@ -1,11 +1,9 @@
 mod dense_map;
 mod mean_list;
-mod null;
 mod num_stats;
 
 pub use dense_map::DenseMap;
 pub use mean_list::MeanList;
-pub use null::NullAny;
 pub use num_stats::NumStats;
 
 pub fn fnv<D: AsRef<[u8]>>(data: D) -> u64 {

backend/common/src/util/dense_map.rs

@@ -1,17 +1,22 @@
-pub type Id = usize;
+pub struct DenseMap<Id, T> {
-
-pub struct DenseMap<T> {
     /// List of retired indexes that can be re-used
-    retired: Vec<Id>,
+    retired: Vec<usize>,
     /// All items
     items: Vec<Option<T>>,
+    /// Our ID type
+    _id_ty: std::marker::PhantomData<Id>
 }
 
-impl<T> DenseMap<T> {
+impl<Id, T> DenseMap<Id, T>
+where
+    Id: From<usize> + Copy,
+    usize: From<Id>
+{
     pub fn new() -> Self {
         DenseMap {
             retired: Vec::new(),
             items: Vec::new(),
+            _id_ty: std::marker::PhantomData
         }
     }
 
@@ -25,11 +30,12 @@ impl<T> DenseMap<T> {
     {
         match self.retired.pop() {
             Some(id) => {
-                self.items[id] = Some(f(id));
+                let id_out = id.into();
-                id
+                self.items[id] = Some(f(id_out));
+                id_out
             }
             None => {
-                let id = self.items.len();
+                let id = self.items.len().into();
                 self.items.push(Some(f(id)));
                 id
             }
@@ -37,14 +43,17 @@ impl<T> DenseMap<T> {
     }
 
     pub fn get(&self, id: Id) -> Option<&T> {
+        let id: usize = id.into();
         self.items.get(id).and_then(|item| item.as_ref())
     }
 
     pub fn get_mut(&mut self, id: Id) -> Option<&mut T> {
+        let id: usize = id.into();
         self.items.get_mut(id).and_then(|item| item.as_mut())
     }
 
     pub fn remove(&mut self, id: Id) -> Option<T> {
+        let id: usize = id.into();
         let old = self.items.get_mut(id).and_then(|item| item.take());
 
         if old.is_some() {
@@ -60,14 +69,14 @@ impl<T> DenseMap<T> {
         self.items
             .iter()
             .enumerate()
-            .filter_map(|(id, item)| Some((id, item.as_ref()?)))
+            .filter_map(|(id, item)| Some((id.into(), item.as_ref()?)))
     }
 
     pub fn iter_mut(&mut self) -> impl Iterator<Item = (Id, &mut T)> + '_ {
         self.items
             .iter_mut()
             .enumerate()
-            .filter_map(|(id, item)| Some((id, item.as_mut()?)))
+            .filter_map(|(id, item)| Some((id.into(), item.as_mut()?)))
     }
 
     pub fn len(&self) -> usize {

backend/common/src/util/null.rs

@@ -1,136 +0,0 @@
-use serde::de::{Deserialize, Deserializer, IgnoredAny};
-use serde::ser::{Serialize, Serializer};
-
-/// Alternative to `serde::de::IgnoreAny` that implements `Serialize`.
-/// Will serialize to `null` in JSON, or empty data in bincode.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub struct NullAny;
-
-impl<'de> Deserialize<'de> for NullAny {
-    #[inline]
-    fn deserialize<D>(deserializer: D) -> Result<NullAny, D::Error>
-    where
-        D: Deserializer<'de>,
-    {
-        // `bincode` is going to throw an error here as it does not support `IgnoredAny`.
-        //
-        // When using `bincode` `NullAny` will always serialize to unit (aka no data), so
-        // this safely becomes a no-op.
-        let _ = deserializer.deserialize_ignored_any(IgnoredAny);
-
-        Ok(NullAny)
-    }
-}
-
-impl Serialize for NullAny {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        serializer.serialize_unit()
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use super::NullAny;
-    use bincode::Options;
-    use serde::{Deserialize, Serialize};
-
-    #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
-    struct Dummy {
-        ignore: NullAny,
-    }
-
-    #[test]
-    fn deserialize_json_null() {
-        let dummy: Dummy = serde_json::from_str(r#"{"ignore":null}"#).unwrap();
-
-        assert_eq!(dummy, Dummy { ignore: NullAny });
-    }
-
-    #[test]
-    fn deserialize_json_struct() {
-        let dummy: Dummy = serde_json::from_str(r#"{"ignore":{"foo":"bar"}}"#).unwrap();
-
-        assert_eq!(dummy, Dummy { ignore: NullAny });
-    }
-
-    #[test]
-    fn deserialize_json_struct_invalid() {
-        let dummy = serde_json::from_str::<Dummy>(r#"{"ignore":{"foo":"bar"}"#);
-
-        assert!(dummy.is_err());
-    }
-
-    #[test]
-    fn deserialize_json_vec_any() {
-        let raw = [NullAny; 10];
-        let json = r#"[null,true,false,10,{},[],[null],{"foo":"bar"},[9,9,9],"ten"]"#;
-
-        let deserialized: Vec<NullAny> = serde_json::from_str(json).unwrap();
-
-        assert_eq!(&raw[..], &deserialized);
-    }
-
-    #[test]
-    fn serialize_json_null() {
-        let dummy = Dummy { ignore: NullAny };
-
-        let json = serde_json::to_string(&dummy).unwrap();
-
-        assert_eq!(json, r#"{"ignore":null}"#);
-    }
-
-    #[test]
-    fn bincode_vec() {
-        let raw = vec![NullAny; 10];
-
-        let bytes = bincode::options().serialize(&raw).unwrap();
-
-        assert_eq!(bytes, &[10u8]);
-
-        let deserialized: Vec<NullAny> = bincode::options().deserialize(&bytes).unwrap();
-
-        assert_eq!(raw, deserialized);
-    }
-
-    #[test]
-    fn bincode_tuple() {
-        let raw = (NullAny, "Hello world".to_string());
-
-        let bytes = bincode::options().serialize(&raw).unwrap();
-
-        assert_eq!(bytes, b"\x0BHello world"); // 0B = 11 = length of string
-
-        let deserialized: (NullAny, String) = bincode::options().deserialize(&bytes).unwrap();
-
-        assert_eq!(raw, deserialized);
-    }
-
-    #[test]
-    fn json_vec() {
-        let raw = vec![NullAny; 10];
-
-        let json = serde_json::to_string(&raw).unwrap();
-
-        assert_eq!(json, "[null,null,null,null,null,null,null,null,null,null]");
-
-        let deserialized: Vec<NullAny> = serde_json::from_str(&json).unwrap();
-
-        assert_eq!(raw, deserialized);
-    }
-
-    #[test]
-    fn json_tuple() {
-        let raw = (NullAny, "Hello world".to_string());
-
-        let json = serde_json::to_string(&raw).unwrap();
-
-        assert_eq!(json, r#"[null,"Hello world"]"#);
-
-        let deserialized: (NullAny, String) = serde_json::from_str(&json).unwrap();
-
-        assert_eq!(raw, deserialized);
-    }
-}

backend/shard/src/aggregator.rs

@@ -1,4 +1,4 @@
-use common::{internal_messages::{self, LocalId}, node, assign_id::AssignId, types::BlockHash};
+use common::{internal_messages::{self, ShardNodeId}, node, assign_id::AssignId, types::BlockHash};
 use std::sync::Arc;
 use std::sync::atomic::AtomicU64;
 use futures::{channel::mpsc, future};
@@ -119,7 +119,7 @@ impl Aggregator {
         let mut to_local_id = AssignId::new();
 
         // Any messages coming from nodes that have been muted are ignored:
-        let mut muted: HashSet<LocalId> = HashSet::new();
+        let mut muted: HashSet<ShardNodeId> = HashSet::new();
 
         // Now, loop and receive messages to handle.
         while let Some(msg) = rx_from_external.next().await {

backend/telemetry/src/aggregator/aggregator.rs

@@ -5,12 +5,15 @@ use futures::{ future, Sink, SinkExt };
 use super::inner_loop;
 use crate::find_location::find_location;
 use crate::state::NodeId;
+use common::id_type;
 use std::net::Ipv4Addr;
 
-/// A unique Id is assigned per websocket connection (or more accurately,
+id_type! {
-/// per feed socket and per shard socket). This can be combined with the
+    /// A unique Id is assigned per websocket connection (or more accurately,
-/// [`LocalId`] of messages to give us a global ID.
+    /// per feed socket and per shard socket). This can be combined with the
-type ConnId = u64;
+    /// [`LocalId`] of messages to give us a global ID.
+    pub ConnId(u64)
+}
 
 #[derive(Clone)]
 pub struct Aggregator(Arc<AggregatorInternal>);
@@ -64,13 +67,13 @@ impl Aggregator {
     pub fn subscribe_shard(&self) -> impl Sink<inner_loop::FromShardWebsocket, Error = anyhow::Error> + Unpin {
         // Assign a unique aggregator-local ID to each connection that subscribes, and pass
         // that along with every message to the aggregator loop:
-        let shard_conn_id: ConnId = self.0.shard_conn_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        let shard_conn_id = self.0.shard_conn_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
         let tx_to_aggregator = self.0.tx_to_aggregator.clone();
 
         // Calling `send` on this Sink requires Unpin. There may be a nicer way than this,
         // but pinning by boxing is the easy solution for now:
         Box::pin(tx_to_aggregator.with(move |msg| async move {
-            Ok(inner_loop::ToAggregator::FromShardWebsocket(shard_conn_id, msg))
+            Ok(inner_loop::ToAggregator::FromShardWebsocket(shard_conn_id.into(), msg))
         }))
     }
 
@@ -78,13 +81,13 @@ impl Aggregator {
     pub fn subscribe_feed(&self) -> impl Sink<inner_loop::FromFeedWebsocket, Error = anyhow::Error> + Unpin {
         // Assign a unique aggregator-local ID to each connection that subscribes, and pass
         // that along with every message to the aggregator loop:
-        let feed_conn_id: ConnId = self.0.feed_conn_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
+        let feed_conn_id = self.0.feed_conn_id.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
         let tx_to_aggregator = self.0.tx_to_aggregator.clone();
 
         // Calling `send` on this Sink requires Unpin. There may be a nicer way than this,
         // but pinning by boxing is the easy solution for now:
         Box::pin(tx_to_aggregator.with(move |msg| async move {
-            Ok(inner_loop::ToAggregator::FromFeedWebsocket(feed_conn_id, msg))
+            Ok(inner_loop::ToAggregator::FromFeedWebsocket(feed_conn_id.into(), msg))
         }))
     }
 

backend/telemetry/src/aggregator/inner_loop.rs

@@ -1,7 +1,7 @@
 use common::{
     internal_messages::{
         self,
-        LocalId,
+        ShardNodeId,
         MuteReason
     },
     types::BlockHash,
@@ -16,11 +16,7 @@ use std::collections::{ HashMap, HashSet };
 use crate::state::{ self, State, NodeId };
 use crate::feed_message::{ self, FeedMessageSerializer };
 use crate::find_location;
-
+use super::aggregator::ConnId;
-/// A unique Id is assigned per websocket connection (or more accurately,
-/// per feed socket and per shard socket). This can be combined with the
-/// [`LocalId`] of messages to give us a global ID.
-type ConnId = u64;
 
 /// Incoming messages come via subscriptions, and end up looking like this.
 #[derive(Clone,Debug)]
@@ -40,19 +36,19 @@ pub enum FromShardWebsocket {
     },
     /// Tell the aggregator about a new node.
     Add {
-        local_id: LocalId,
+        local_id: ShardNodeId,
         ip: Option<std::net::IpAddr>,
         node: common::types::NodeDetails,
         genesis_hash: common::types::BlockHash
     },
     /// Update/pass through details about a node.
     Update {
-        local_id: LocalId,
+        local_id: ShardNodeId,
         payload: node::Payload
     },
     /// Tell the aggregator that a node has been removed when it disconnects.
     Remove {
-        local_id: LocalId,
+        local_id: ShardNodeId,
     },
     /// The shard is disconnected.
     Disconnected
@@ -63,7 +59,7 @@ pub enum FromShardWebsocket {
 pub enum ToShardWebsocket {
     /// Mute messages to the core by passing the shard-local ID of them.
     Mute {
-        local_id: LocalId,
+        local_id: ShardNodeId,
         reason: internal_messages::MuteReason
     }
 }
@@ -129,7 +125,7 @@ pub struct InnerLoop {
     node_state: State,
     /// We maintain a mapping between NodeId and ConnId+LocalId, so that we know
     /// which messages are about which nodes.
-    node_ids: BiMap<NodeId, (ConnId, LocalId)>,
+    node_ids: BiMap<NodeId, (ConnId, ShardNodeId)>,
 
     /// Keep track of how to send messages out to feeds.
     feed_channels: HashMap<ConnId, mpsc::UnboundedSender<ToFeedWebsocket>>,
@@ -194,7 +190,7 @@ impl InnerLoop {
         if let Some(loc) = location {
             let mut feed_message_serializer = FeedMessageSerializer::new();
             feed_message_serializer.push(feed_message::LocatedNode(
-                node_id,
+                node_id.get_chain_node_id().into(),
                 loc.latitude,
                 loc.longitude,
                 &loc.city
@@ -212,7 +208,8 @@ impl InnerLoop {
 
     /// Handle messages coming from shards.
     async fn handle_from_shard(&mut self, shard_conn_id: ConnId, msg: FromShardWebsocket) {
-        log::debug!("Message from shard ({}): {:?}", shard_conn_id, msg);
+        log::debug!("Message from shard ({:?}): {:?}", shard_conn_id, msg);
+
         match msg {
             FromShardWebsocket::Initialize { channel } => {
                 self.shard_channels.insert(shard_conn_id, channel);
@@ -249,7 +246,7 @@ impl InnerLoop {
 
                         // Tell chain subscribers about the node we've just added:
                         let mut feed_messages_for_chain = FeedMessageSerializer::new();
-                        feed_messages_for_chain.push(feed_message::AddedNode(node_id, &details.node));
+                        feed_messages_for_chain.push(feed_message::AddedNode(node_id.get_chain_node_id().into(), &details.node));
                         self.finalize_and_broadcast_to_chain_feeds(&genesis_hash, feed_messages_for_chain).await;
 
                         // Tell everybody about the new node count and potential rename:
@@ -272,7 +269,7 @@ impl InnerLoop {
                 let node_id = match self.node_ids.remove_by_right(&(shard_conn_id, local_id)) {
                     Some((node_id, _)) => node_id,
                     None => {
-                        log::error!("Cannot find ID for node with shard/connectionId of {}/{}", shard_conn_id, local_id);
+                        log::error!("Cannot find ID for node with shard/connectionId of {:?}/{:?}", shard_conn_id, local_id);
                         return
                     }
                 };
@@ -282,11 +279,22 @@ impl InnerLoop {
                 let node_id = match self.node_ids.get_by_right(&(shard_conn_id, local_id)) {
                     Some(id) => *id,
                     None => {
-                        log::error!("Cannot find ID for node with shard/connectionId of {}/{}", shard_conn_id, local_id);
+                        log::error!("Cannot find ID for node with shard/connectionId of {:?}/{:?}", shard_conn_id, local_id);
                         return
                     }
                 };
-                self.handle_from_shard_update(node_id, payload).await;
+
+                let mut feed_message_serializer = FeedMessageSerializer::new();
+                let broadcast_finality = self.node_state.update_node(node_id, payload, &mut feed_message_serializer);
+
+                if let Some(chain) = self.node_state.get_chain_by_node_id(node_id) {
+                    let genesis_hash = *chain.genesis_hash();
+                    if broadcast_finality {
+                        self.finalize_and_broadcast_to_chain_finality_feeds(&genesis_hash, feed_message_serializer).await;
+                    } else {
+                        self.finalize_and_broadcast_to_chain_feeds(&genesis_hash, feed_message_serializer).await;
+                    }
+                }
             },
             FromShardWebsocket::Disconnected => {
                 // Find all nodes associated with this shard connection ID:
@@ -302,24 +310,9 @@ impl InnerLoop {
         }
     }
 
-    async fn handle_from_shard_update(&mut self, node_id: NodeId, payload: node::Payload) {
-        let mut feed_message_serializer = FeedMessageSerializer::new();
-
-        let broadcast_finality = self.node_state.update_node(node_id, payload, &mut feed_message_serializer);
-
-        if let Some(chain) = self.node_state.get_chain_by_node_id(node_id) {
-            let genesis_hash = *chain.genesis_hash();
-            if broadcast_finality {
-                self.finalize_and_broadcast_to_chain_finality_feeds(&genesis_hash, feed_message_serializer).await;
-            } else {
-                self.finalize_and_broadcast_to_chain_feeds(&genesis_hash, feed_message_serializer).await;
-            }
-        }
-    }
-
     /// Handle messages coming from feeds.
     async fn handle_from_feed(&mut self, feed_conn_id: ConnId, msg: FromFeedWebsocket) {
-        log::debug!("Message from feed ({}): {:?}", feed_conn_id, msg);
+        log::debug!("Message from feed ({:?}): {:?}", feed_conn_id, msg);
         match msg {
             FromFeedWebsocket::Initialize { mut channel } => {
                 self.feed_channels.insert(feed_conn_id, channel.clone());
@@ -396,7 +389,9 @@ impl InnerLoop {
                     new_chain.finalized_block().height,
                     new_chain.finalized_block().hash
                 ));
-                for (idx, (node_id, node)) in new_chain.iter_nodes().enumerate() {
+                for (idx, (chain_node_id, node)) in new_chain.iter_nodes().enumerate() {
+                    let chain_node_id = chain_node_id.into();
+
                     // Send subscription confirmation and chain head before doing all the nodes,
                     // and continue sending batches of 32 nodes a time over the wire subsequently
                     if idx % 32 == 0 {
@@ -404,14 +399,14 @@ impl InnerLoop {
                             let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes)).await;
                         }
                     }
-                    feed_serializer.push(feed_message::AddedNode(node_id, node));
+                    feed_serializer.push(feed_message::AddedNode(chain_node_id, node));
                     feed_serializer.push(feed_message::FinalizedBlock(
-                        node_id,
+                        chain_node_id,
                         node.finalized().height,
                         node.finalized().hash,
                     ));
                     if node.stale() {
-                        feed_serializer.push(feed_message::StaleNode(node_id));
+                        feed_serializer.push(feed_message::StaleNode(chain_node_id));
                     }
                 }
                 if let Some(bytes) = feed_serializer.into_finalized() {
@@ -480,9 +475,9 @@ impl InnerLoop {
         self.node_ids.remove_by_left(&node_id);
 
         let removed_details = match self.node_state.remove_node(node_id) {
-            Ok(remove_details) => remove_details,
+            Some(remove_details) => remove_details,
-            Err(err) => {
+            None => {
-                log::error!("Error removing node {}: {}", node_id, err);
+                log::error!("Could not find node {:?}", node_id);
                 return
             }
         };
@@ -504,7 +499,7 @@ impl InnerLoop {
         // Assuming the chain hasn't gone away, tell chain subscribers about the node removal
         if removed_details.chain_node_count != 0 {
             feed_for_chain.push(
-                feed_message::RemovedNode(node_id)
+                feed_message::RemovedNode(node_id.get_chain_node_id().into())
             );
         }
     }

backend/telemetry/src/feed_message.rs

@@ -7,10 +7,13 @@ use std::mem;
 use crate::state::Node;
 use serde_json::to_writer;
 use common::types::{
-    Address, BlockDetails, BlockHash, BlockNumber, NodeHardware, NodeIO, NodeId, NodeStats,
+    BlockDetails, BlockHash, BlockNumber, NodeHardware, NodeIO, NodeStats,
     Timestamp
 };
 
+type Address = Box<str>;
+type FeedNodeId = usize;
+
 pub trait FeedMessage {
     const ACTION: u8;
 }
@@ -133,28 +136,28 @@ pub struct BestBlock(pub BlockNumber, pub Timestamp, pub Option<u64>);
 #[derive(Serialize)]
 pub struct BestFinalized(pub BlockNumber, pub BlockHash);
 
-pub struct AddedNode<'a>(pub NodeId, pub &'a Node);
+pub struct AddedNode<'a>(pub FeedNodeId, pub &'a Node);
 
 #[derive(Serialize)]
-pub struct RemovedNode(pub NodeId);
+pub struct RemovedNode(pub FeedNodeId);
 
 #[derive(Serialize)]
-pub struct LocatedNode<'a>(pub NodeId, pub f32, pub f32, pub &'a str);
+pub struct LocatedNode<'a>(pub FeedNodeId, pub f32, pub f32, pub &'a str);
 
 #[derive(Serialize)]
-pub struct ImportedBlock<'a>(pub NodeId, pub &'a BlockDetails);
+pub struct ImportedBlock<'a>(pub FeedNodeId, pub &'a BlockDetails);
 
 #[derive(Serialize)]
-pub struct FinalizedBlock(pub NodeId, pub BlockNumber, pub BlockHash);
+pub struct FinalizedBlock(pub FeedNodeId, pub BlockNumber, pub BlockHash);
 
 #[derive(Serialize)]
-pub struct NodeStatsUpdate<'a>(pub NodeId, pub &'a NodeStats);
+pub struct NodeStatsUpdate<'a>(pub FeedNodeId, pub &'a NodeStats);
 
 #[derive(Serialize)]
-pub struct NodeIOUpdate<'a>(pub NodeId, pub &'a NodeIO);
+pub struct NodeIOUpdate<'a>(pub FeedNodeId, pub &'a NodeIO);
 
 #[derive(Serialize)]
-pub struct Hardware<'a>(pub NodeId, pub &'a NodeHardware);
+pub struct Hardware<'a>(pub FeedNodeId, pub &'a NodeHardware);
 
 #[derive(Serialize)]
 pub struct TimeSync(pub u64);
@@ -203,7 +206,7 @@ pub struct AfgAuthoritySet(
 );
 
 #[derive(Serialize)]
-pub struct StaleNode(pub NodeId);
+pub struct StaleNode(pub FeedNodeId);
 
 impl FeedMessageWrite for AddedNode<'_> {
     fn write_to_feed(&self, ser: &mut FeedMessageSerializer) {

backend/telemetry/src/state/chain.rs

@@ -5,11 +5,17 @@ use common::util::{now, DenseMap, NumStats};
 use common::most_seen::MostSeen;
 use common::node::Payload;
 use once_cell::sync::Lazy;
+use common::id_type;
 
 use crate::feed_message::{self, FeedMessageSerializer};
+use crate::find_location;
 
 use super::node::Node;
-use super::NodeId;
+
+id_type!{
+    /// A Node ID that is unique to the chain it's in.
+    pub ChainNodeId(usize)
+}
 
 pub type Label = Box<str>;
 
@@ -20,7 +26,7 @@ pub struct Chain {
     /// the most commonly used label as nodes are added/removed.
     labels: MostSeen<Label>,
     /// Set of nodes that are in this chain
-    node_ids: HashSet<NodeId>,
+    nodes: DenseMap<ChainNodeId, Node>,
     /// Best block
     best: Block,
     /// Finalized block
@@ -38,6 +44,7 @@ pub struct Chain {
 pub enum AddNodeResult {
     Overquota,
     Added {
+        id: ChainNodeId,
         chain_renamed: bool
     }
 }
@@ -65,7 +72,7 @@ impl Chain {
     pub fn new(genesis_hash: BlockHash) -> Self {
         Chain {
             labels: MostSeen::default(),
-            node_ids: HashSet::new(),
+            nodes: DenseMap::new(),
             best: Block::zero(),
             finalized: Block::zero(),
             block_times: NumStats::new(50),
@@ -79,29 +86,34 @@ impl Chain {
     pub fn can_add_node(&self) -> bool {
         // Dynamically determine the max nodes based on the most common
         // label so far, in case it changes to something with a different limit.
-        self.node_ids.len() < max_nodes(self.labels.best())
+        self.nodes.len() < max_nodes(self.labels.best())
     }
 
-    /// Assign a node to this chain. If the function returns false, it
+    /// Assign a node to this chain.
-    /// means that the node could not be added as we're at quota.
+    pub fn add_node(&mut self, node: Node) -> AddNodeResult {
-    pub fn add_node(&mut self, node_id: NodeId, chain_label: &Box<str>) -> AddNodeResult {
         if !self.can_add_node() {
             return AddNodeResult::Overquota
         }
 
-        let label_result = self.labels.insert(chain_label);
+        let node_chain_label = &node.details().chain;
-        self.node_ids.insert(node_id);
+        let label_result = self.labels.insert(node_chain_label);
+        let node_id = self.nodes.add(node);
 
         AddNodeResult::Added {
+            id: node_id,
             chain_renamed: label_result.has_changed()
         }
     }
 
-    /// Remove a node from this chain. We expect the label it used for the chain so
+    /// Remove a node from this chain.
-    /// that we can keep track of which label is most popular.
+    pub fn remove_node(&mut self, node_id: ChainNodeId) -> RemoveNodeResult {
-    pub fn remove_node(&mut self, node_id: NodeId, chain_label: &Box<str>) -> RemoveNodeResult {
+        let node = match self.nodes.remove(node_id) {
-        let label_result = self.labels.remove(&chain_label);
+            Some(node) => node,
-        self.node_ids.remove(&node_id);
+            None => return RemoveNodeResult { chain_renamed: false }
+        };
+
+        let node_chain_label = &node.details().chain;
+        let label_result = self.labels.remove(node_chain_label);
 
         RemoveNodeResult {
             chain_renamed: label_result.has_changed()
@@ -110,25 +122,25 @@ impl Chain {
 
     /// Attempt to update the best block seen in this chain.
     /// Returns a boolean which denotes whether the output is for finalization feeds (true) or not (false).
-    pub fn update_node(&mut self, all_nodes: &mut DenseMap<Node>, nid: NodeId, payload: Payload, feed: &mut FeedMessageSerializer) -> bool {
+    pub fn update_node(&mut self, nid: ChainNodeId, payload: Payload, feed: &mut FeedMessageSerializer) -> bool {
 
         if let Some(block) = payload.best_block() {
-            self.handle_block(all_nodes, block, nid, feed);
+            self.handle_block(block, nid, feed);
         }
 
-        if let Some(node) = all_nodes.get_mut(nid) {
+        if let Some(node) = self.nodes.get_mut(nid) {
             match payload {
                 Payload::SystemInterval(ref interval) => {
                     if node.update_hardware(interval) {
-                        feed.push(feed_message::Hardware(nid, node.hardware()));
+                        feed.push(feed_message::Hardware(nid.into(), node.hardware()));
                     }
 
                     if let Some(stats) = node.update_stats(interval) {
-                        feed.push(feed_message::NodeStatsUpdate(nid, stats));
+                        feed.push(feed_message::NodeStatsUpdate(nid.into(), stats));
                     }
 
                     if let Some(io) = node.update_io(interval) {
-                        feed.push(feed_message::NodeIOUpdate(nid, io));
+                        feed.push(feed_message::NodeIOUpdate(nid.into(), io));
                     }
                 }
                 Payload::AfgAuthoritySet(authority) => {
@@ -187,7 +199,7 @@ impl Chain {
             if let Some(block) = payload.finalized_block() {
                 if let Some(finalized) = node.update_finalized(block) {
                     feed.push(feed_message::FinalizedBlock(
-                        nid,
+                        nid.into(),
                         finalized.height,
                         finalized.hash,
                     ));
@@ -203,14 +215,14 @@ impl Chain {
         false
     }
 
-    fn handle_block(&mut self, all_nodes: &mut DenseMap<Node>, block: &Block, nid: NodeId, feed: &mut FeedMessageSerializer) {
+    fn handle_block(&mut self, block: &Block, nid: ChainNodeId, feed: &mut FeedMessageSerializer) {
         let mut propagation_time = None;
         let now = now();
-        let nodes_len = self.node_ids.len();
+        let nodes_len = self.nodes.len();
 
-        self.update_stale_nodes(all_nodes, now, feed);
+        self.update_stale_nodes(now, feed);
 
-        let node = match all_nodes.get_mut(nid) {
+        let node = match self.nodes.get_mut(nid) {
             Some(node) => node,
             None => return,
         };
@@ -243,14 +255,14 @@ impl Chain {
             }
 
             if let Some(details) = node.update_details(now, propagation_time) {
-                feed.push(feed_message::ImportedBlock(nid, details));
+                feed.push(feed_message::ImportedBlock(nid.into(), details));
             }
         }
     }
 
     /// Check if the chain is stale (has not received a new best block in a while).
     /// If so, find a new best block, ignoring any stale nodes and marking them as such.
-    fn update_stale_nodes(&mut self, all_nodes: &mut DenseMap<Node>, now: u64, feed: &mut FeedMessageSerializer) {
+    fn update_stale_nodes(&mut self, now: u64, feed: &mut FeedMessageSerializer) {
 
         let threshold = now - STALE_TIMEOUT;
         let timestamp = match self.timestamp {
@@ -267,11 +279,7 @@ impl Chain {
         let mut finalized = Block::zero();
         let mut timestamp = None;
 
-        for &nid in self.node_ids.iter() {
+        for (nid, node) in self.nodes.iter_mut() {
-            let node = match all_nodes.get_mut(nid) {
-                Some(node) => node,
-                None => continue
-            };
             if !node.update_stale(threshold) {
                 if node.best().height > best.height {
                     best = *node.best();
@@ -282,7 +290,7 @@ impl Chain {
                     finalized = *node.finalized();
                 }
             } else {
-                feed.push(feed_message::StaleNode(nid));
+                feed.push(feed_message::StaleNode(nid.into()));
             }
         }
 
@@ -301,14 +309,26 @@ impl Chain {
         }
     }
 
+    pub fn update_node_location(&mut self, node_id: ChainNodeId, location: find_location::Location) -> bool {
+        if let Some(node) = self.nodes.get_mut(node_id) {
+            node.update_location(location);
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn get_node(&self, id: ChainNodeId) -> Option<&Node> {
+        self.nodes.get(id)
+    }
+    pub fn iter_nodes(&self) -> impl Iterator<Item=(ChainNodeId, &Node)> {
+        self.nodes.iter()
+    }
     pub fn label(&self) -> &str {
         &self.labels.best()
     }
-    pub fn node_ids(&self) -> impl Iterator<Item=NodeId> + '_ {
-        self.node_ids.iter().copied()
-    }
     pub fn node_count(&self) -> usize {
-        self.node_ids.len()
+        self.nodes.len()
     }
     pub fn best_block(&self) -> &Block {
         &self.best

backend/telemetry/src/state/state.rs

@@ -7,28 +7,39 @@ use common::node::Payload;
 use std::iter::IntoIterator;
 use crate::feed_message::FeedMessageSerializer;
 use crate::find_location;
+use common::id_type;
 
-use super::chain::{ self, Chain };
+use super::chain::{ self, Chain, ChainNodeId };
 
-pub type NodeId = usize;
+id_type!{
-pub type ChainId = usize;
+    /// A globally unique Chain ID.
+    pub ChainId(usize)
+}
+
+/// A "global" Node ID is a composite of the ID of the chain it's
+/// on, and it's chain local ID.
+#[derive(Debug,Clone,Copy,Hash,PartialEq,Eq)]
+pub struct NodeId(ChainId, ChainNodeId);
+
+impl NodeId {
+    pub fn get_chain_node_id(&self) -> ChainNodeId {
+        self.1
+    }
+}
 
 /// Our state constains node and chain information
 pub struct State {
-    // Store nodes and chains in a fairly compact format.
+    chains: DenseMap<ChainId, Chain>,
-    nodes: DenseMap<Node>,
-    chains: DenseMap<Chain>,
 
     // Find the right chain given various details.
     chains_by_genesis_hash: HashMap<BlockHash, ChainId>,
     chains_by_label: HashMap<Box<str>, ChainId>,
-    chains_by_node: HashMap<NodeId, ChainId>,
 
     /// Chain labels that we do not want to allow connecting.
     denylist: HashSet<String>,
 }
 
-/// Adding a node to a chain leads to this result
+/// Adding a node to a chain leads to this node_idult
 pub enum AddNodeResult<'a> {
     /// The chain is on the "deny list", so we can't add the node
     ChainOnDenyList,
@@ -38,6 +49,16 @@ pub enum AddNodeResult<'a> {
     NodeAddedToChain(NodeAddedToChain<'a>)
 }
 
+#[cfg(test)]
+impl <'a> AddNodeResult<'a> {
+    pub fn unwrap_id(&self) -> NodeId {
+        match &self {
+            AddNodeResult::NodeAddedToChain(d) => d.id,
+            _ => panic!("Attempt to unwrap_id on AddNodeResult that did not succeed")
+        }
+    }
+}
+
 pub struct NodeAddedToChain<'a> {
     /// The ID assigned to this node.
     pub id: NodeId,
@@ -65,25 +86,12 @@ pub struct RemovedNode {
     pub new_chain_label: Box<str>,
 }
 
-/// If removing a node goes wrong, we get this back
-#[derive(Debug, thiserror::Error)]
-pub enum RemoveNodeError {
-    /// The node that you tried to remove wasn't found
-    #[error("Node not found")]
-    NodeNotFound,
-    /// The chain associated to the node wasn't found
-    #[error("Node chain not found")]
-    NodeChainNotFound
-}
-
 impl State {
     pub fn new<T: IntoIterator<Item=String>>(denylist: T) -> State {
         State {
-            nodes: DenseMap::new(),
             chains: DenseMap::new(),
             chains_by_genesis_hash: HashMap::new(),
             chains_by_label: HashMap::new(),
-            chains_by_node: HashMap::new(),
             denylist: denylist.into_iter().collect(),
         }
     }
@@ -91,28 +99,27 @@ impl State {
     pub fn iter_chains(&self) -> impl Iterator<Item=StateChain<'_>> {
         self.chains
             .iter()
-            .map(move |(_,chain)| StateChain { state: self, chain })
+            .map(move |(_,chain)| StateChain { chain })
     }
 
     pub fn get_chain_by_node_id(&self, node_id: NodeId) -> Option<StateChain<'_>> {
-        self.chains_by_node
+        self.chains
-            .get(&node_id)
+            .get(node_id.0)
-            .and_then(|&chain_id| self.chains.get(chain_id))
+            .map(|chain| StateChain { chain })
-            .map(|chain| StateChain { state: self, chain })
     }
 
     pub fn get_chain_by_genesis_hash(&self, genesis_hash: &BlockHash) -> Option<StateChain<'_>> {
         self.chains_by_genesis_hash
             .get(genesis_hash)
             .and_then(|&chain_id| self.chains.get(chain_id))
-            .map(|chain| StateChain { state: self, chain })
+            .map(|chain| StateChain { chain })
     }
 
     pub fn get_chain_by_label(&self, label: &str) -> Option<StateChain<'_>> {
         self.chains_by_label
             .get(label)
             .and_then(|&chain_id| self.chains.get(chain_id))
-            .map(|chain| StateChain { state: self, chain })
+            .map(|chain| StateChain { chain })
     }
 
     pub fn add_node(&mut self, genesis_hash: BlockHash, node_details: NodeDetails) -> AddNodeResult<'_> {
@@ -137,36 +144,26 @@ impl State {
         let chain = self.chains.get_mut(chain_id)
             .expect("should be known to exist after the above (unless chains_by_genesis_hash out of sync)");
 
-        // What ID will the node have when it's added? We don't actually want
+        let node = Node::new(node_details);
-        // to add it until we know whether the chain will accept it, but we want
-        // an ID to give to the chain.
-        let node_id = self.nodes.next_id();
-        let node_chain_label = node_details.chain.clone();
         let old_chain_label = chain.label().into();
 
-        match chain.add_node(node_id, &node_chain_label) {
+        match chain.add_node(node) {
             chain::AddNodeResult::Overquota => {
                 AddNodeResult::ChainOverQuota
             },
-            chain::AddNodeResult::Added { chain_renamed } => {
+            chain::AddNodeResult::Added { id, chain_renamed } => {
                 let chain = &*chain;
 
-                // Actually add the node, and a reference to its chain,
-                // if the chain adds it successfully:
-                self.nodes.add(Node::new(node_details));
-                self.chains_by_node.insert(node_id, chain_id);
-
                 // Update the label we use to reference the chain if
                 // it changes (it'll always change first time a node's added):
                 if chain_renamed {
                     self.chains_by_label.remove(&old_chain_label);
-                    self.chains_by_label.insert(chain.label().to_string().into_boxed_str(), chain_id);
+                    self.chains_by_label.insert(chain.label().into(), chain_id);
                 }
 
-                let node = self.nodes.get(node_id).expect("node added above");
                 AddNodeResult::NodeAddedToChain(NodeAddedToChain {
-                    id: node_id,
+                    id: NodeId(chain_id, id),
-                    node: node,
+                    node: chain.get_node(id).expect("node added above"),
                     old_chain_label: old_chain_label,
                     new_chain_label: chain.label(),
                     chain_node_count: chain.node_count(),
@@ -177,24 +174,20 @@ impl State {
     }
 
     /// Remove a node
-    pub fn remove_node(&mut self, node_id: NodeId) -> Result<RemovedNode,RemoveNodeError> {
+    pub fn remove_node(&mut self, NodeId(chain_id, chain_node_id): NodeId) -> Option<RemovedNode> {
-        self.nodes.remove(node_id)
+        let chain = self.chains.get_mut(chain_id)?;
-            .ok_or(RemoveNodeError::NodeNotFound)?;
+        let old_chain_label = chain.label().into();
 
-        let chain_id = self.chains_by_node.remove(&node_id)
+        // Actually remove the node
-            .ok_or(RemoveNodeError::NodeChainNotFound)?;
+        let remove_result = chain.remove_node(chain_node_id);
 
-        let chain = self.chains.get_mut(chain_id)
+        // Get updated chain details.
-            .ok_or(RemoveNodeError::NodeChainNotFound)?;
+        let new_chain_label: Box<str> = chain.label().into();
-
-        let old_chain_label = chain.label().to_string().into_boxed_str();
-        let remove_result = chain.remove_node(node_id, &old_chain_label);
-        let new_chain_label = chain.label().to_string().into_boxed_str();
         let chain_node_count = chain.node_count();
-        let genesis_hash = *chain.genesis_hash();
 
         // Is the chain empty? Remove if so and clean up indexes to it
         if chain_node_count == 0 {
+            let genesis_hash = *chain.genesis_hash();
             self.chains_by_label.remove(&old_chain_label);
             self.chains_by_genesis_hash.remove(&genesis_hash);
             self.chains.remove(chain_id);
@@ -206,7 +199,7 @@ impl State {
             self.chains_by_label.insert(new_chain_label.clone(), chain_id);
         }
 
-        Ok(RemovedNode {
+        Some(RemovedNode {
             old_chain_label,
             new_chain_label,
             chain_node_count: chain_node_count,
@@ -216,25 +209,19 @@ impl State {
 
     /// Attempt to update the best block seen, given a node and block.
     /// Returns a boolean which denotes whether the output is for finalization feeds (true) or not (false).
-    pub fn update_node(&mut self, node_id: NodeId, payload: Payload, feed: &mut FeedMessageSerializer) -> bool {
+    pub fn update_node(&mut self, NodeId(chain_id, chain_node_id): NodeId, payload: Payload, feed: &mut FeedMessageSerializer) -> bool {
-        let chain_id = match self.chains_by_node.get(&node_id) {
-            Some(chain_id) => *chain_id,
-            None => { log::error!("Cannot find chain_id for node with ID {}", node_id); return false }
-        };
-
         let chain = match self.chains.get_mut(chain_id) {
             Some(chain) => chain,
-            None => { log::error!("Cannot find chain for node with ID {}", node_id); return false }
+            None => { log::error!("Cannot find chain for node with ID {:?}", chain_id); return false }
         };
 
-        chain.update_node(&mut self.nodes, node_id, payload, feed)
+        chain.update_node(chain_node_id, payload, feed)
     }
 
     /// Update the location for a node. Return `false` if the node was not found.
-    pub fn update_node_location(&mut self, node_id: NodeId, location: find_location::Location) -> bool {
+    pub fn update_node_location(&mut self, NodeId(chain_id, chain_node_id): NodeId, location: find_location::Location) -> bool {
-        if let Some(node) = self.nodes.get_mut(node_id) {
+        if let Some(chain) = self.chains.get_mut(chain_id) {
-            node.update_location(location);
+            chain.update_node_location(chain_node_id, location)
-            true
         } else {
             false
         }
@@ -247,7 +234,6 @@ impl State {
 /// aren't really intended for use outside of [`State`] methods. Any modification
 /// of a chain needs to go through [`State`].
 pub struct StateChain<'a> {
-    state: &'a State,
     chain: &'a Chain
 }
 
@@ -273,11 +259,8 @@ impl <'a> StateChain<'a> {
     pub fn finalized_block(&self) -> &'a Block {
         self.chain.finalized_block()
     }
-    pub fn iter_nodes(&self) -> impl Iterator<Item=(NodeId, &'a Node)> + 'a {
+    pub fn iter_nodes(&self) -> impl Iterator<Item=(ChainNodeId, &'a Node)> + 'a {
-        let state = self.state;
+        self.chain.iter_nodes()
-        self.chain.node_ids().filter_map(move |id| {
-            Some((id, state.nodes.get(id)?))
-        })
     }
 }
 
@@ -314,7 +297,7 @@ mod test {
             AddNodeResult::NodeAddedToChain(details) => details
         };
 
-        assert_eq!(add_node_result.id, 0);
+        assert_eq!(add_node_result.id, NodeId(0.into(), 0.into()));
         assert_eq!(&*add_node_result.old_chain_label, "");
         assert_eq!(&*add_node_result.new_chain_label, "Chain One");
         assert_eq!(add_node_result.chain_node_count, 1);
@@ -331,7 +314,7 @@ mod test {
             AddNodeResult::NodeAddedToChain(details) => details
         };
 
-        assert_eq!(add_node_result.id, 1);
+        assert_eq!(add_node_result.id, NodeId(0.into(), 1.into()));
         assert_eq!(&*add_node_result.old_chain_label, "Chain One");
         assert_eq!(&*add_node_result.new_chain_label, "Chain One");
         assert_eq!(add_node_result.chain_node_count, 2);
@@ -343,32 +326,38 @@ mod test {
         let mut state = State::new(None);
 
         let chain1_genesis = BlockHash::from_low_u64_be(1);
-        state.add_node(chain1_genesis, node("A", "Chain One")); // 0
+        let node_id0 = state
+            .add_node(chain1_genesis, node("A", "Chain One")) // 0
+            .unwrap_id();
 
-        assert_eq!(state.get_chain_by_node_id(0).expect("Chain should exist").label(), "Chain One");
+        assert_eq!(state.get_chain_by_node_id(node_id0).expect("Chain should exist").label(), "Chain One");
         assert!(state.get_chain_by_label("Chain One").is_some());
         assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_some());
 
-        state.add_node(chain1_genesis, node("B", "Chain Two")); // 1
+        let node_id1 = state
+            .add_node(chain1_genesis, node("B", "Chain Two")) // 1
+            .unwrap_id();
 
         // Chain name hasn't changed yet; "Chain One" as common as "Chain Two"..
-        assert_eq!(state.get_chain_by_node_id(0).expect("Chain should exist").label(), "Chain One");
+        assert_eq!(state.get_chain_by_node_id(node_id0).expect("Chain should exist").label(), "Chain One");
         assert!(state.get_chain_by_label("Chain One").is_some());
         assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_some());
 
-        state.add_node(chain1_genesis, node("B", "Chain Two")); // 2
+        let node_id2 = state
+            .add_node(chain1_genesis, node("B", "Chain Two"))
+            .unwrap_id(); // 2
 
         // Chain name has changed; "Chain Two" the winner now..
-        assert_eq!(state.get_chain_by_node_id(0).expect("Chain should exist").label(), "Chain Two");
+        assert_eq!(state.get_chain_by_node_id(node_id0).expect("Chain should exist").label(), "Chain Two");
         assert!(state.get_chain_by_label("Chain One").is_none());
         assert!(state.get_chain_by_label("Chain Two").is_some());
         assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_some());
 
-        state.remove_node(1).expect("Removal OK (id: 1)");
+        state.remove_node(node_id1).expect("Removal OK (id: 1)");
-        state.remove_node(2).expect("Removal OK (id: 2");
+        state.remove_node(node_id2).expect("Removal OK (id: 2)");
 
         // Removed both "Chain Two" nodes; dominant name now "Chain One" again..
-        assert_eq!(state.get_chain_by_node_id(0).expect("Chain should exist").label(), "Chain One");
+        assert_eq!(state.get_chain_by_node_id(node_id0).expect("Chain should exist").label(), "Chain One");
         assert!(state.get_chain_by_label("Chain One").is_some());
         assert!(state.get_chain_by_label("Chain Two").is_none());
         assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_some());
@@ -379,13 +368,15 @@ mod test {
         let mut state = State::new(None);
 
         let chain1_genesis = BlockHash::from_low_u64_be(1);
-        state.add_node(chain1_genesis, node("A", "Chain One")); // 0
+        let node_id = state
+            .add_node(chain1_genesis, node("A", "Chain One")) // 0
+            .unwrap_id();
 
         assert!(state.get_chain_by_label("Chain One").is_some());
         assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_some());
         assert_eq!(state.iter_chains().count(), 1);
 
-        state.remove_node(0);
+        state.remove_node(node_id);
 
         assert!(state.get_chain_by_label("Chain One").is_none());
         assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_none());