rename binaries for clarity, and first pass of connect_to_servers test util

backend/telemetry_core/src/state/chain.rs

@@ -0,0 +1,370 @@
+use common::node_message::Payload;
+use common::node_types::{Block, Timestamp};
+use common::node_types::{BlockHash, BlockNumber};
+use common::{id_type, time, DenseMap, MostSeen, NumStats};
+use once_cell::sync::Lazy;
+use std::collections::HashSet;
+
+use crate::feed_message::{self, FeedMessageSerializer};
+use crate::find_location;
+
+use super::node::Node;
+
+id_type! {
+    /// A Node ID that is unique to the chain it's in.
+    pub struct ChainNodeId(usize)
+}
+
+pub type Label = Box<str>;
+
+const STALE_TIMEOUT: u64 = 2 * 60 * 1000; // 2 minutes
+
+pub struct Chain {
+    /// Labels that nodes use for this chain. We keep track of
+    /// the most commonly used label as nodes are added/removed.
+    labels: MostSeen<Label>,
+    /// Set of nodes that are in this chain
+    nodes: DenseMap<ChainNodeId, Node>,
+    /// Best block
+    best: Block,
+    /// Finalized block
+    finalized: Block,
+    /// Block times history, stored so we can calculate averages
+    block_times: NumStats<u64>,
+    /// Calculated average block time
+    average_block_time: Option<u64>,
+    /// When the best block first arrived
+    timestamp: Option<Timestamp>,
+    /// Genesis hash of this chain
+    genesis_hash: BlockHash,
+}
+
+pub enum AddNodeResult {
+    Overquota,
+    Added {
+        id: ChainNodeId,
+        chain_renamed: bool,
+    },
+}
+
+pub struct RemoveNodeResult {
+    pub chain_renamed: bool,
+}
+
+/// Labels of chains we consider "first party". These chains allow any
+/// number of nodes to connect.
+static FIRST_PARTY_NETWORKS: Lazy<HashSet<&'static str>> = Lazy::new(|| {
+    let mut set = HashSet::new();
+    set.insert("Polkadot");
+    set.insert("Kusama");
+    set.insert("Westend");
+    set.insert("Rococo");
+    set
+});
+
+/// Max number of nodes allowed to connect to the telemetry server.
+const THIRD_PARTY_NETWORKS_MAX_NODES: usize = 500;
+
+impl Chain {
+    /// Create a new chain with an initial label.
+    pub fn new(genesis_hash: BlockHash) -> Self {
+        Chain {
+            labels: MostSeen::default(),
+            nodes: DenseMap::new(),
+            best: Block::zero(),
+            finalized: Block::zero(),
+            block_times: NumStats::new(50),
+            average_block_time: None,
+            timestamp: None,
+            genesis_hash,
+        }
+    }
+
+    /// Can we add a node? If not, it's because the chain is at its quota.
+    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.nodes.len() < max_nodes(self.labels.best())
+    }
+
+    /// Assign a node to this chain.
+    pub fn add_node(&mut self, node: Node) -> AddNodeResult {
+        if !self.can_add_node() {
+            return AddNodeResult::Overquota;
+        }
+
+        let node_chain_label = &node.details().chain;
+        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.
+    pub fn remove_node(&mut self, node_id: ChainNodeId) -> RemoveNodeResult {
+        let node = match self.nodes.remove(node_id) {
+            Some(node) => node,
+            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(),
+        }
+    }
+
+    /// 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,
+        nid: ChainNodeId,
+        payload: Payload,
+        feed: &mut FeedMessageSerializer,
+    ) -> bool {
+        if let Some(block) = payload.best_block() {
+            self.handle_block(block, nid, feed);
+        }
+
+        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.into(), node.hardware()));
+                    }
+
+                    if let Some(stats) = node.update_stats(interval) {
+                        feed.push(feed_message::NodeStatsUpdate(nid.into(), stats));
+                    }
+
+                    if let Some(io) = node.update_io(interval) {
+                        feed.push(feed_message::NodeIOUpdate(nid.into(), io));
+                    }
+                }
+                Payload::AfgAuthoritySet(authority) => {
+                    node.set_validator_address(authority.authority_id.clone());
+                    return false;
+                }
+                Payload::AfgFinalized(finalized) => {
+                    if let Ok(finalized_number) = finalized.finalized_number.parse::<BlockNumber>()
+                    {
+                        if let Some(addr) = node.details().validator.clone() {
+                            feed.push(feed_message::AfgFinalized(
+                                addr,
+                                finalized_number,
+                                finalized.finalized_hash,
+                            ));
+                        }
+                    }
+                    return true;
+                }
+                Payload::AfgReceivedPrecommit(precommit) => {
+                    if let Ok(finalized_number) = precommit.target_number.parse::<BlockNumber>() {
+                        if let Some(addr) = node.details().validator.clone() {
+                            let voter = precommit.voter.clone();
+                            feed.push(feed_message::AfgReceivedPrecommit(
+                                addr,
+                                finalized_number,
+                                precommit.target_hash,
+                                voter,
+                            ));
+                        }
+                    }
+                    return true;
+                }
+                Payload::AfgReceivedPrevote(prevote) => {
+                    if let Ok(finalized_number) = prevote.target_number.parse::<BlockNumber>() {
+                        if let Some(addr) = node.details().validator.clone() {
+                            let voter = prevote.voter.clone();
+                            feed.push(feed_message::AfgReceivedPrevote(
+                                addr,
+                                finalized_number,
+                                prevote.target_hash,
+                                voter,
+                            ));
+                        }
+                    }
+                    return true;
+                }
+                Payload::AfgReceivedCommit(_) => {}
+                _ => (),
+            }
+
+            if let Some(block) = payload.finalized_block() {
+                if let Some(finalized) = node.update_finalized(block) {
+                    feed.push(feed_message::FinalizedBlock(
+                        nid.into(),
+                        finalized.height,
+                        finalized.hash,
+                    ));
+
+                    if finalized.height > self.finalized.height {
+                        self.finalized = *finalized;
+                        feed.push(feed_message::BestFinalized(
+                            finalized.height,
+                            finalized.hash,
+                        ));
+                    }
+                }
+            }
+        }
+
+        false
+    }
+
+    fn handle_block(&mut self, block: &Block, nid: ChainNodeId, feed: &mut FeedMessageSerializer) {
+        let mut propagation_time = None;
+        let now = time::now();
+        let nodes_len = self.nodes.len();
+
+        self.update_stale_nodes(now, feed);
+
+        let node = match self.nodes.get_mut(nid) {
+            Some(node) => node,
+            None => return,
+        };
+
+        if node.update_block(*block) {
+            if block.height > self.best.height {
+                self.best = *block;
+                log::debug!(
+                    "[{}] [nodes={}] new best block={}/{:?}",
+                    self.labels.best(),
+                    nodes_len,
+                    self.best.height,
+                    self.best.hash,
+                );
+                if let Some(timestamp) = self.timestamp {
+                    self.block_times.push(now - timestamp);
+                    self.average_block_time = Some(self.block_times.average());
+                }
+                self.timestamp = Some(now);
+                feed.push(feed_message::BestBlock(
+                    self.best.height,
+                    now,
+                    self.average_block_time,
+                ));
+                propagation_time = Some(0);
+            } else if block.height == self.best.height {
+                if let Some(timestamp) = self.timestamp {
+                    propagation_time = Some(now - timestamp);
+                }
+            }
+
+            if let Some(details) = node.update_details(now, propagation_time) {
+                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, now: u64, feed: &mut FeedMessageSerializer) {
+        let threshold = now - STALE_TIMEOUT;
+        let timestamp = match self.timestamp {
+            Some(ts) => ts,
+            None => return,
+        };
+
+        if timestamp > threshold {
+            // Timestamp is in range, nothing to do
+            return;
+        }
+
+        let mut best = Block::zero();
+        let mut finalized = Block::zero();
+        let mut timestamp = None;
+
+        for (nid, node) in self.nodes.iter_mut() {
+            if !node.update_stale(threshold) {
+                if node.best().height > best.height {
+                    best = *node.best();
+                    timestamp = Some(node.best_timestamp());
+                }
+
+                if node.finalized().height > finalized.height {
+                    finalized = *node.finalized();
+                }
+            } else {
+                feed.push(feed_message::StaleNode(nid.into()));
+            }
+        }
+
+        if self.best.height != 0 || self.finalized.height != 0 {
+            self.best = best;
+            self.finalized = finalized;
+            self.block_times.reset();
+            self.timestamp = timestamp;
+
+            feed.push(feed_message::BestBlock(
+                self.best.height,
+                timestamp.unwrap_or(now),
+                None,
+            ));
+            feed.push(feed_message::BestFinalized(
+                finalized.height,
+                finalized.hash,
+            ));
+        }
+    }
+
+    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_count(&self) -> usize {
+        self.nodes.len()
+    }
+    pub fn best_block(&self) -> &Block {
+        &self.best
+    }
+    pub fn timestamp(&self) -> Option<Timestamp> {
+        self.timestamp
+    }
+    pub fn average_block_time(&self) -> Option<u64> {
+        self.average_block_time
+    }
+    pub fn finalized_block(&self) -> &Block {
+        &self.finalized
+    }
+    pub fn genesis_hash(&self) -> &BlockHash {
+        &self.genesis_hash
+    }
+}
+
+/// First party networks (Polkadot, Kusama etc) are allowed any number of nodes.
+/// Third party networks are allowed `THIRD_PARTY_NETWORKS_MAX_NODES` nodes and
+/// no more.
+fn max_nodes(label: &str) -> usize {
+    if FIRST_PARTY_NETWORKS.contains(label) {
+        usize::MAX
+    } else {
+        THIRD_PARTY_NETWORKS_MAX_NODES
+    }
+}

backend/telemetry_core/src/state/mod.rs

@@ -0,0 +1,7 @@
+mod chain;
+mod node;
+
+mod state;
+
+pub use node::Node;
+pub use state::*;

backend/telemetry_core/src/state/node.rs

@@ -0,0 +1,207 @@
+use crate::find_location;
+use common::node_message::SystemInterval;
+use common::node_types::{
+    Block, BlockDetails, NodeDetails, NodeHardware, NodeIO, NodeLocation, NodeStats, Timestamp,
+};
+use common::time;
+
+/// Minimum time between block below broadcasting updates to the browser gets throttled, in ms.
+const THROTTLE_THRESHOLD: u64 = 100;
+/// Minimum time of intervals for block updates sent to the browser when throttled, in ms.
+const THROTTLE_INTERVAL: u64 = 1000;
+
+pub struct Node {
+    /// Static details
+    details: NodeDetails,
+    /// Basic stats
+    stats: NodeStats,
+    /// Node IO stats
+    io: NodeIO,
+    /// Best block
+    best: BlockDetails,
+    /// Finalized block
+    finalized: Block,
+    /// Timer for throttling block updates
+    throttle: u64,
+    /// Hardware stats over time
+    hardware: NodeHardware,
+    /// Physical location details
+    location: find_location::Location,
+    /// Flag marking if the node is stale (not syncing or producing blocks)
+    stale: bool,
+    /// Unix timestamp for when node started up (falls back to connection time)
+    startup_time: Option<Timestamp>,
+}
+
+impl Node {
+    pub fn new(mut details: NodeDetails) -> Self {
+        let startup_time = details
+            .startup_time
+            .take()
+            .and_then(|time| time.parse().ok());
+
+        Node {
+            details,
+            stats: NodeStats::default(),
+            io: NodeIO::default(),
+            best: BlockDetails::default(),
+            finalized: Block::zero(),
+            throttle: 0,
+            hardware: NodeHardware::default(),
+            location: None,
+            stale: false,
+            startup_time,
+        }
+    }
+
+    pub fn details(&self) -> &NodeDetails {
+        &self.details
+    }
+
+    pub fn stats(&self) -> &NodeStats {
+        &self.stats
+    }
+
+    pub fn io(&self) -> &NodeIO {
+        &self.io
+    }
+
+    pub fn best(&self) -> &Block {
+        &self.best.block
+    }
+
+    pub fn best_timestamp(&self) -> u64 {
+        self.best.block_timestamp
+    }
+
+    pub fn finalized(&self) -> &Block {
+        &self.finalized
+    }
+
+    pub fn hardware(&self) -> &NodeHardware {
+        &self.hardware
+    }
+
+    pub fn location(&self) -> Option<&NodeLocation> {
+        self.location.as_deref()
+    }
+
+    pub fn update_location(&mut self, location: find_location::Location) {
+        self.location = location;
+    }
+
+    pub fn block_details(&self) -> &BlockDetails {
+        &self.best
+    }
+
+    pub fn update_block(&mut self, block: Block) -> bool {
+        if block.height > self.best.block.height {
+            self.stale = false;
+            self.best.block = block;
+
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn update_details(
+        &mut self,
+        timestamp: u64,
+        propagation_time: Option<u64>,
+    ) -> Option<&BlockDetails> {
+        self.best.block_time = timestamp - self.best.block_timestamp;
+        self.best.block_timestamp = timestamp;
+        self.best.propagation_time = propagation_time;
+
+        if self.throttle < timestamp {
+            if self.best.block_time <= THROTTLE_THRESHOLD {
+                self.throttle = timestamp + THROTTLE_INTERVAL;
+            }
+
+            Some(&self.best)
+        } else {
+            None
+        }
+    }
+
+    pub fn update_hardware(&mut self, interval: &SystemInterval) -> bool {
+        let mut changed = false;
+
+        if let Some(upload) = interval.bandwidth_upload {
+            changed |= self.hardware.upload.push(upload);
+        }
+        if let Some(download) = interval.bandwidth_download {
+            changed |= self.hardware.download.push(download);
+        }
+        self.hardware.chart_stamps.push(time::now() as f64);
+
+        changed
+    }
+
+    pub fn update_stats(&mut self, interval: &SystemInterval) -> Option<&NodeStats> {
+        let mut changed = false;
+
+        if let Some(peers) = interval.peers {
+            if peers != self.stats.peers {
+                self.stats.peers = peers;
+                changed = true;
+            }
+        }
+        if let Some(txcount) = interval.txcount {
+            if txcount != self.stats.txcount {
+                self.stats.txcount = txcount;
+                changed = true;
+            }
+        }
+
+        if changed {
+            Some(&self.stats)
+        } else {
+            None
+        }
+    }
+
+    pub fn update_io(&mut self, interval: &SystemInterval) -> Option<&NodeIO> {
+        let mut changed = false;
+
+        if let Some(size) = interval.used_state_cache_size {
+            changed |= self.io.used_state_cache_size.push(size);
+        }
+
+        if changed {
+            Some(&self.io)
+        } else {
+            None
+        }
+    }
+
+    pub fn update_finalized(&mut self, block: Block) -> Option<&Block> {
+        if block.height > self.finalized.height {
+            self.finalized = block;
+            Some(self.finalized())
+        } else {
+            None
+        }
+    }
+
+    pub fn update_stale(&mut self, threshold: u64) -> bool {
+        if self.best.block_timestamp < threshold {
+            self.stale = true;
+        }
+
+        self.stale
+    }
+
+    pub fn stale(&self) -> bool {
+        self.stale
+    }
+
+    pub fn set_validator_address(&mut self, addr: Box<str>) {
+        self.details.validator = Some(addr);
+    }
+
+    pub fn startup_time(&self) -> Option<Timestamp> {
+        self.startup_time
+    }
+}

backend/telemetry_core/src/state/state.rs

@@ -0,0 +1,414 @@
+use super::node::Node;
+use crate::feed_message::FeedMessageSerializer;
+use crate::find_location;
+use common::node_message::Payload;
+use common::node_types::{Block, BlockHash, NodeDetails, Timestamp};
+use common::{id_type, DenseMap};
+use std::collections::{HashMap, HashSet};
+use std::iter::IntoIterator;
+
+use super::chain::{self, Chain, ChainNodeId};
+
+id_type! {
+    /// A globally unique Chain ID.
+    pub struct 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 {
+    chains: DenseMap<ChainId, Chain>,
+
+    // Find the right chain given various details.
+    chains_by_genesis_hash: HashMap<BlockHash, ChainId>,
+    chains_by_label: HashMap<Box<str>, ChainId>,
+
+    /// Chain labels that we do not want to allow connecting.
+    denylist: HashSet<String>,
+}
+
+/// 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,
+    /// The chain is over quota (too many nodes connected), so can't add the node
+    ChainOverQuota,
+    /// The node was added to the chain
+    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,
+    /// The old label of the chain.
+    pub old_chain_label: Box<str>,
+    /// The new label of the chain.
+    pub new_chain_label: &'a str,
+    /// The node that was added.
+    pub node: &'a Node,
+    /// Number of nodes in the chain. If 1, the chain was just added.
+    pub chain_node_count: usize,
+    /// Has the chain label been updated?
+    pub has_chain_label_changed: bool,
+}
+
+/// if removing a node is successful, we get this information back.
+pub struct RemovedNode {
+    /// How many nodes remain on the chain (0 if the chain was removed)
+    pub chain_node_count: usize,
+    /// Has the chain label been updated?
+    pub has_chain_label_changed: bool,
+    /// The old label of the chain.
+    pub old_chain_label: Box<str>,
+    /// The new label of the chain.
+    pub new_chain_label: Box<str>,
+}
+
+impl State {
+    pub fn new<T: IntoIterator<Item = String>>(denylist: T) -> State {
+        State {
+            chains: DenseMap::new(),
+            chains_by_genesis_hash: HashMap::new(),
+            chains_by_label: HashMap::new(),
+            denylist: denylist.into_iter().collect(),
+        }
+    }
+
+    pub fn iter_chains(&self) -> impl Iterator<Item = StateChain<'_>> {
+        self.chains
+            .iter()
+            .map(move |(_, chain)| StateChain { chain })
+    }
+
+    pub fn get_chain_by_node_id(&self, node_id: NodeId) -> Option<StateChain<'_>> {
+        self.chains.get(node_id.0).map(|chain| StateChain { 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 { 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 { chain })
+    }
+
+    pub fn add_node(
+        &mut self,
+        genesis_hash: BlockHash,
+        node_details: NodeDetails,
+    ) -> AddNodeResult<'_> {
+        if self.denylist.contains(&*node_details.chain) {
+            return AddNodeResult::ChainOnDenyList;
+        }
+
+        // Get the chain ID, creating a new empty chain if one doesn't exist.
+        // If we create a chain here, we are expecting that it will allow at
+        // least this node to be added, because we don't currently try and clean it up
+        // if the add fails.
+        let chain_id = match self.chains_by_genesis_hash.get(&genesis_hash) {
+            Some(id) => *id,
+            None => {
+                let chain_id = self.chains.add(Chain::new(genesis_hash));
+                self.chains_by_genesis_hash.insert(genesis_hash, chain_id);
+                chain_id
+            }
+        };
+
+        // Get the chain.
+        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)",
+        );
+
+        let node = Node::new(node_details);
+        let old_chain_label = chain.label().into();
+
+        match chain.add_node(node) {
+            chain::AddNodeResult::Overquota => AddNodeResult::ChainOverQuota,
+            chain::AddNodeResult::Added { id, chain_renamed } => {
+                let chain = &*chain;
+
+                // 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().into(), chain_id);
+                }
+
+                AddNodeResult::NodeAddedToChain(NodeAddedToChain {
+                    id: NodeId(chain_id, id),
+                    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(),
+                    has_chain_label_changed: chain_renamed,
+                })
+            }
+        }
+    }
+
+    /// Remove a node
+    pub fn remove_node(&mut self, NodeId(chain_id, chain_node_id): NodeId) -> Option<RemovedNode> {
+        let chain = self.chains.get_mut(chain_id)?;
+        let old_chain_label = chain.label().into();
+
+        // Actually remove the node
+        let remove_result = chain.remove_node(chain_node_id);
+
+        // Get updated chain details.
+        let new_chain_label: Box<str> = chain.label().into();
+        let chain_node_count = chain.node_count();
+
+        // 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);
+        }
+
+        // Make sure chains always referenced by their most common label:
+        if remove_result.chain_renamed {
+            self.chains_by_label.remove(&old_chain_label);
+            self.chains_by_label
+                .insert(new_chain_label.clone(), chain_id);
+        }
+
+        Some(RemovedNode {
+            old_chain_label,
+            new_chain_label,
+            chain_node_count: chain_node_count,
+            has_chain_label_changed: remove_result.chain_renamed,
+        })
+    }
+
+    /// 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,
+        NodeId(chain_id, chain_node_id): NodeId,
+        payload: Payload,
+        feed: &mut FeedMessageSerializer,
+    ) -> bool {
+        let chain = match self.chains.get_mut(chain_id) {
+            Some(chain) => chain,
+            None => {
+                log::error!("Cannot find chain for node with ID {:?}", chain_id);
+                return false;
+            }
+        };
+
+        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,
+        NodeId(chain_id, chain_node_id): NodeId,
+        location: find_location::Location,
+    ) -> bool {
+        if let Some(chain) = self.chains.get_mut(chain_id) {
+            chain.update_node_location(chain_node_id, location)
+        } else {
+            false
+        }
+    }
+}
+
+/// When we ask for a chain, we get this struct back. This ensures that we have
+/// a consistent public interface, and don't expose methods on [`Chain`] that
+/// aren't really intended for use outside of [`State`] methods. Any modification
+/// of a chain needs to go through [`State`].
+pub struct StateChain<'a> {
+    chain: &'a Chain,
+}
+
+impl<'a> StateChain<'a> {
+    pub fn label(&self) -> &'a str {
+        self.chain.label()
+    }
+    pub fn genesis_hash(&self) -> &'a BlockHash {
+        self.chain.genesis_hash()
+    }
+    pub fn node_count(&self) -> usize {
+        self.chain.node_count()
+    }
+    pub fn best_block(&self) -> &'a Block {
+        self.chain.best_block()
+    }
+    pub fn timestamp(&self) -> Timestamp {
+        self.chain.timestamp().unwrap_or(0)
+    }
+    pub fn average_block_time(&self) -> Option<u64> {
+        self.chain.average_block_time()
+    }
+    pub fn finalized_block(&self) -> &'a Block {
+        self.chain.finalized_block()
+    }
+    pub fn iter_nodes(&self) -> impl Iterator<Item = (ChainNodeId, &'a Node)> + 'a {
+        self.chain.iter_nodes()
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn node(name: &str, chain: &str) -> NodeDetails {
+        NodeDetails {
+            chain: chain.into(),
+            name: name.into(),
+            implementation: "Bar".into(),
+            version: "0.1".into(),
+            validator: None,
+            network_id: None,
+            startup_time: None,
+        }
+    }
+
+    #[test]
+    fn adding_a_node_returns_expected_response() {
+        let mut state = State::new(None);
+
+        let chain1_genesis = BlockHash::from_low_u64_be(1);
+
+        let add_result = state.add_node(chain1_genesis, node("A", "Chain One"));
+
+        let add_node_result = match add_result {
+            AddNodeResult::ChainOnDenyList => panic!("Chain not on deny list"),
+            AddNodeResult::ChainOverQuota => panic!("Chain not Overquota"),
+            AddNodeResult::NodeAddedToChain(details) => details,
+        };
+
+        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);
+        assert_eq!(add_node_result.has_chain_label_changed, true);
+
+        let add_result = state.add_node(chain1_genesis, node("A", "Chain One"));
+
+        let add_node_result = match add_result {
+            AddNodeResult::ChainOnDenyList => panic!("Chain not on deny list"),
+            AddNodeResult::ChainOverQuota => panic!("Chain not Overquota"),
+            AddNodeResult::NodeAddedToChain(details) => details,
+        };
+
+        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);
+        assert_eq!(add_node_result.has_chain_label_changed, false);
+    }
+
+    #[test]
+    fn adding_and_removing_nodes_updates_chain_label_mapping() {
+        let mut state = State::new(None);
+
+        let chain1_genesis = BlockHash::from_low_u64_be(1);
+        let node_id0 = state
+            .add_node(chain1_genesis, node("A", "Chain One")) // 0
+            .unwrap_id();
+
+        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());
+
+        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(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());
+
+        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(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(node_id1).expect("Removal OK (id: 1)");
+        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(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());
+    }
+
+    #[test]
+    fn chain_removed_when_last_node_is() {
+        let mut state = State::new(None);
+
+        let chain1_genesis = BlockHash::from_low_u64_be(1);
+        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(node_id);
+
+        assert!(state.get_chain_by_label("Chain One").is_none());
+        assert!(state.get_chain_by_genesis_hash(&chain1_genesis).is_none());
+        assert_eq!(state.iter_chains().count(), 0);
+    }
+}