mirror of
https://github.com/pezkuwichain/pezkuwi-telemetry.git
synced 2026-06-09 20:21:01 +00:00
rename binaries for clarity, and first pass of connect_to_servers test util
This commit is contained in:
@@ -0,0 +1,117 @@
|
||||
use super::inner_loop;
|
||||
use crate::find_location::find_location;
|
||||
use crate::state::NodeId;
|
||||
use common::id_type;
|
||||
use futures::channel::mpsc;
|
||||
use futures::{future, Sink, SinkExt};
|
||||
use std::net::Ipv4Addr;
|
||||
use std::sync::atomic::AtomicU64;
|
||||
use std::sync::Arc;
|
||||
|
||||
id_type! {
|
||||
/// 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.
|
||||
pub struct ConnId(u64)
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct Aggregator(Arc<AggregatorInternal>);
|
||||
|
||||
struct AggregatorInternal {
|
||||
/// Shards that connect are each assigned a unique connection ID.
|
||||
/// This helps us know who to send messages back to (especially in
|
||||
/// conjunction with the [`LocalId`] that messages will come with).
|
||||
shard_conn_id: AtomicU64,
|
||||
/// Feeds that connect have their own unique connection ID, too.
|
||||
feed_conn_id: AtomicU64,
|
||||
/// Send messages in to the aggregator from the outside via this. This is
|
||||
/// stored here so that anybody holding an `Aggregator` handle can
|
||||
/// make use of it.
|
||||
tx_to_aggregator: mpsc::Sender<inner_loop::ToAggregator>,
|
||||
}
|
||||
|
||||
impl Aggregator {
|
||||
/// Spawn a new Aggregator. This connects to the telemetry backend
|
||||
pub async fn spawn(denylist: Vec<String>) -> anyhow::Result<Aggregator> {
|
||||
let (tx_to_aggregator, rx_from_external) = mpsc::channel(10);
|
||||
|
||||
// Kick off a locator task to locate nodes, which hands back a channel to make location requests
|
||||
let tx_to_locator = find_location(tx_to_aggregator.clone().with(|(node_id, msg)| {
|
||||
future::ok::<_, mpsc::SendError>(inner_loop::ToAggregator::FromFindLocation(
|
||||
node_id, msg,
|
||||
))
|
||||
}));
|
||||
|
||||
// Handle any incoming messages in our handler loop:
|
||||
tokio::spawn(Aggregator::handle_messages(
|
||||
rx_from_external,
|
||||
tx_to_locator,
|
||||
denylist,
|
||||
));
|
||||
|
||||
// Return a handle to our aggregator:
|
||||
Ok(Aggregator(Arc::new(AggregatorInternal {
|
||||
shard_conn_id: AtomicU64::new(1),
|
||||
feed_conn_id: AtomicU64::new(1),
|
||||
tx_to_aggregator,
|
||||
})))
|
||||
}
|
||||
|
||||
// This is spawned into a separate task and handles any messages coming
|
||||
// in to the aggregator. If nobody is tolding the tx side of the channel
|
||||
// any more, this task will gracefully end.
|
||||
async fn handle_messages(
|
||||
rx_from_external: mpsc::Receiver<inner_loop::ToAggregator>,
|
||||
tx_to_aggregator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>,
|
||||
denylist: Vec<String>,
|
||||
) {
|
||||
inner_loop::InnerLoop::new(rx_from_external, tx_to_aggregator, denylist)
|
||||
.handle()
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Return a sink that a shard can send messages into to be handled by the 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 = 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.into(),
|
||||
msg,
|
||||
))
|
||||
}))
|
||||
}
|
||||
|
||||
/// Return a sink that a feed can send messages into to be handled by the 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 = 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.into(),
|
||||
msg,
|
||||
))
|
||||
}))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,627 @@
|
||||
use super::aggregator::ConnId;
|
||||
use crate::feed_message::{self, FeedMessageSerializer};
|
||||
use crate::find_location;
|
||||
use crate::state::{self, NodeId, State};
|
||||
use bimap::BiMap;
|
||||
use common::{
|
||||
internal_messages::{self, MuteReason, ShardNodeId},
|
||||
node_message,
|
||||
node_types::BlockHash,
|
||||
time,
|
||||
};
|
||||
use futures::channel::mpsc;
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::{
|
||||
net::{IpAddr, Ipv4Addr},
|
||||
str::FromStr,
|
||||
};
|
||||
|
||||
/// Incoming messages come via subscriptions, and end up looking like this.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToAggregator {
|
||||
FromShardWebsocket(ConnId, FromShardWebsocket),
|
||||
FromFeedWebsocket(ConnId, FromFeedWebsocket),
|
||||
FromFindLocation(NodeId, find_location::Location),
|
||||
}
|
||||
|
||||
/// An incoming shard connection can send these messages to the aggregator.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum FromShardWebsocket {
|
||||
/// When the socket is opened, it'll send this first
|
||||
/// so that we have a way to communicate back to it.
|
||||
Initialize {
|
||||
channel: mpsc::Sender<ToShardWebsocket>,
|
||||
},
|
||||
/// Tell the aggregator about a new node.
|
||||
Add {
|
||||
local_id: ShardNodeId,
|
||||
ip: Option<std::net::IpAddr>,
|
||||
node: common::node_types::NodeDetails,
|
||||
genesis_hash: common::node_types::BlockHash,
|
||||
},
|
||||
/// Update/pass through details about a node.
|
||||
Update {
|
||||
local_id: ShardNodeId,
|
||||
payload: node_message::Payload,
|
||||
},
|
||||
/// Tell the aggregator that a node has been removed when it disconnects.
|
||||
Remove { local_id: ShardNodeId },
|
||||
/// The shard is disconnected.
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
/// The aggregator can these messages back to a shard connection.
|
||||
#[derive(Debug)]
|
||||
pub enum ToShardWebsocket {
|
||||
/// Mute messages to the core by passing the shard-local ID of them.
|
||||
Mute {
|
||||
local_id: ShardNodeId,
|
||||
reason: internal_messages::MuteReason,
|
||||
},
|
||||
}
|
||||
|
||||
/// An incoming feed connection can send these messages to the aggregator.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum FromFeedWebsocket {
|
||||
/// When the socket is opened, it'll send this first
|
||||
/// so that we have a way to communicate back to it.
|
||||
/// Unbounded so that slow feeds don't block aggregato
|
||||
/// progress.
|
||||
Initialize {
|
||||
channel: mpsc::UnboundedSender<ToFeedWebsocket>,
|
||||
},
|
||||
/// The feed can subscribe to a chain to receive
|
||||
/// messages relating to it.
|
||||
Subscribe { chain: Box<str> },
|
||||
/// The feed wants finality info for the chain, too.
|
||||
SendFinality,
|
||||
/// The feed doesn't want any more finality info for the chain.
|
||||
NoMoreFinality,
|
||||
/// An explicit ping message.
|
||||
Ping { value: Box<str> },
|
||||
/// The feed is disconnected.
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
// The frontend sends text based commands; parse them into these messages:
|
||||
impl FromStr for FromFeedWebsocket {
|
||||
type Err = anyhow::Error;
|
||||
fn from_str(s: &str) -> Result<Self, Self::Err> {
|
||||
let (cmd, value) = match s.find(':') {
|
||||
Some(idx) => (&s[..idx], s[idx + 1..].into()),
|
||||
None => return Err(anyhow::anyhow!("Expecting format `CMD:CHAIN_NAME`")),
|
||||
};
|
||||
match cmd {
|
||||
"ping" => Ok(FromFeedWebsocket::Ping { value }),
|
||||
"subscribe" => Ok(FromFeedWebsocket::Subscribe { chain: value }),
|
||||
"send-finality" => Ok(FromFeedWebsocket::SendFinality),
|
||||
"no-more-finality" => Ok(FromFeedWebsocket::NoMoreFinality),
|
||||
_ => return Err(anyhow::anyhow!("Command {} not recognised", cmd)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The aggregator can these messages back to a feed connection.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum ToFeedWebsocket {
|
||||
Bytes(Vec<u8>),
|
||||
}
|
||||
|
||||
/// Instances of this are responsible for handling incoming and
|
||||
/// outgoing messages in the main aggregator loop.
|
||||
pub struct InnerLoop {
|
||||
/// Messages from the outside world come into this:
|
||||
rx_from_external: mpsc::Receiver<ToAggregator>,
|
||||
|
||||
/// The state of our chains and nodes lives here:
|
||||
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, ShardNodeId)>,
|
||||
|
||||
/// Keep track of how to send messages out to feeds.
|
||||
feed_channels: HashMap<ConnId, mpsc::UnboundedSender<ToFeedWebsocket>>,
|
||||
/// Keep track of how to send messages out to shards.
|
||||
shard_channels: HashMap<ConnId, mpsc::Sender<ToShardWebsocket>>,
|
||||
|
||||
/// Which chain is a feed subscribed to?
|
||||
/// Feed Connection ID -> Chain Genesis Hash
|
||||
feed_conn_id_to_chain: HashMap<ConnId, BlockHash>,
|
||||
/// Which feeds are subscribed to a given chain (needs to stay in sync with above)?
|
||||
/// Chain Genesis Hash -> Feed Connection IDs
|
||||
chain_to_feed_conn_ids: HashMap<BlockHash, HashSet<ConnId>>,
|
||||
|
||||
/// These feeds want finality info, too.
|
||||
feed_conn_id_finality: HashSet<ConnId>,
|
||||
|
||||
/// Send messages here to make geographical location requests.
|
||||
tx_to_locator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>,
|
||||
}
|
||||
|
||||
impl InnerLoop {
|
||||
/// Create a new inner loop handler with the various state it needs.
|
||||
pub fn new(
|
||||
rx_from_external: mpsc::Receiver<ToAggregator>,
|
||||
tx_to_locator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>,
|
||||
denylist: Vec<String>,
|
||||
) -> Self {
|
||||
InnerLoop {
|
||||
rx_from_external,
|
||||
node_state: State::new(denylist),
|
||||
node_ids: BiMap::new(),
|
||||
feed_channels: HashMap::new(),
|
||||
shard_channels: HashMap::new(),
|
||||
feed_conn_id_to_chain: HashMap::new(),
|
||||
chain_to_feed_conn_ids: HashMap::new(),
|
||||
feed_conn_id_finality: HashSet::new(),
|
||||
tx_to_locator,
|
||||
}
|
||||
}
|
||||
|
||||
/// Start handling and responding to incoming messages.
|
||||
pub async fn handle(mut self) {
|
||||
while let Some(msg) = self.rx_from_external.next().await {
|
||||
match msg {
|
||||
ToAggregator::FromFeedWebsocket(feed_conn_id, msg) => {
|
||||
self.handle_from_feed(feed_conn_id, msg).await
|
||||
}
|
||||
ToAggregator::FromShardWebsocket(shard_conn_id, msg) => {
|
||||
self.handle_from_shard(shard_conn_id, msg).await
|
||||
}
|
||||
ToAggregator::FromFindLocation(node_id, location) => {
|
||||
self.handle_from_find_location(node_id, location).await
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle messages that come from the node geographical locator.
|
||||
async fn handle_from_find_location(
|
||||
&mut self,
|
||||
node_id: NodeId,
|
||||
location: find_location::Location,
|
||||
) {
|
||||
self.node_state
|
||||
.update_node_location(node_id, location.clone());
|
||||
|
||||
if let Some(loc) = location {
|
||||
let mut feed_message_serializer = FeedMessageSerializer::new();
|
||||
feed_message_serializer.push(feed_message::LocatedNode(
|
||||
node_id.get_chain_node_id().into(),
|
||||
loc.latitude,
|
||||
loc.longitude,
|
||||
&loc.city,
|
||||
));
|
||||
|
||||
let chain_genesis_hash = self
|
||||
.node_state
|
||||
.get_chain_by_node_id(node_id)
|
||||
.map(|chain| *chain.genesis_hash());
|
||||
|
||||
if let Some(chain_genesis_hash) = chain_genesis_hash {
|
||||
self.finalize_and_broadcast_to_chain_feeds(
|
||||
&chain_genesis_hash,
|
||||
feed_message_serializer,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// 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);
|
||||
|
||||
match msg {
|
||||
FromShardWebsocket::Initialize { channel } => {
|
||||
self.shard_channels.insert(shard_conn_id, channel);
|
||||
}
|
||||
FromShardWebsocket::Add {
|
||||
local_id,
|
||||
ip,
|
||||
node,
|
||||
genesis_hash,
|
||||
} => {
|
||||
match self.node_state.add_node(genesis_hash, node) {
|
||||
state::AddNodeResult::ChainOnDenyList => {
|
||||
if let Some(shard_conn) = self.shard_channels.get_mut(&shard_conn_id) {
|
||||
let _ = shard_conn
|
||||
.send(ToShardWebsocket::Mute {
|
||||
local_id,
|
||||
reason: MuteReason::ChainNotAllowed,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
state::AddNodeResult::ChainOverQuota => {
|
||||
if let Some(shard_conn) = self.shard_channels.get_mut(&shard_conn_id) {
|
||||
let _ = shard_conn
|
||||
.send(ToShardWebsocket::Mute {
|
||||
local_id,
|
||||
reason: MuteReason::Overquota,
|
||||
})
|
||||
.await;
|
||||
}
|
||||
}
|
||||
state::AddNodeResult::NodeAddedToChain(details) => {
|
||||
let node_id = details.id;
|
||||
|
||||
// Record ID <-> (shardId,localId) for future messages:
|
||||
self.node_ids.insert(node_id, (shard_conn_id, local_id));
|
||||
|
||||
// Don't hold onto details too long because we want &mut self later:
|
||||
let old_chain_label = details.old_chain_label.to_owned();
|
||||
let new_chain_label = details.new_chain_label.to_owned();
|
||||
let chain_node_count = details.chain_node_count;
|
||||
let has_chain_label_changed = details.has_chain_label_changed;
|
||||
|
||||
// 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.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:
|
||||
let mut feed_messages_for_all = FeedMessageSerializer::new();
|
||||
if has_chain_label_changed {
|
||||
feed_messages_for_all
|
||||
.push(feed_message::RemovedChain(&old_chain_label));
|
||||
}
|
||||
feed_messages_for_all
|
||||
.push(feed_message::AddedChain(&new_chain_label, chain_node_count));
|
||||
self.finalize_and_broadcast_to_all_feeds(feed_messages_for_all)
|
||||
.await;
|
||||
|
||||
// Ask for the grographical location of the node.
|
||||
// Currently we only geographically locate IPV4 addresses so ignore IPV6.
|
||||
if let Some(IpAddr::V4(ip_v4)) = ip {
|
||||
let _ = self.tx_to_locator.send((node_id, ip_v4)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
FromShardWebsocket::Remove { local_id } => {
|
||||
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
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
self.remove_nodes_and_broadcast_result(Some(node_id)).await;
|
||||
}
|
||||
FromShardWebsocket::Update { local_id, payload } => {
|
||||
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
|
||||
);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
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:
|
||||
let node_ids_to_remove: Vec<NodeId> = self
|
||||
.node_ids
|
||||
.iter()
|
||||
.filter(|(_, &(this_shard_conn_id, _))| shard_conn_id == this_shard_conn_id)
|
||||
.map(|(&node_id, _)| node_id)
|
||||
.collect();
|
||||
|
||||
// ... and remove them:
|
||||
self.remove_nodes_and_broadcast_result(node_ids_to_remove)
|
||||
.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);
|
||||
match msg {
|
||||
FromFeedWebsocket::Initialize { mut channel } => {
|
||||
self.feed_channels.insert(feed_conn_id, channel.clone());
|
||||
|
||||
// Tell the new feed subscription some basic things to get it going:
|
||||
let mut feed_serializer = FeedMessageSerializer::new();
|
||||
feed_serializer.push(feed_message::Version(31));
|
||||
for chain in self.node_state.iter_chains() {
|
||||
feed_serializer
|
||||
.push(feed_message::AddedChain(chain.label(), chain.node_count()));
|
||||
}
|
||||
|
||||
// Send this to the channel that subscribed:
|
||||
if let Some(bytes) = feed_serializer.into_finalized() {
|
||||
let _ = channel.send(ToFeedWebsocket::Bytes(bytes)).await;
|
||||
}
|
||||
}
|
||||
FromFeedWebsocket::Ping { value } => {
|
||||
let feed_channel = match self.feed_channels.get_mut(&feed_conn_id) {
|
||||
Some(chan) => chan,
|
||||
None => return,
|
||||
};
|
||||
|
||||
// Pong!
|
||||
let mut feed_serializer = FeedMessageSerializer::new();
|
||||
feed_serializer.push(feed_message::Pong(&value));
|
||||
if let Some(bytes) = feed_serializer.into_finalized() {
|
||||
let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes)).await;
|
||||
}
|
||||
}
|
||||
FromFeedWebsocket::Subscribe { chain } => {
|
||||
let feed_channel = match self.feed_channels.get_mut(&feed_conn_id) {
|
||||
Some(chan) => chan,
|
||||
None => return,
|
||||
};
|
||||
|
||||
// Unsubscribe from previous chain if subscribed to one:
|
||||
let old_genesis_hash = self.feed_conn_id_to_chain.remove(&feed_conn_id);
|
||||
if let Some(old_genesis_hash) = &old_genesis_hash {
|
||||
if let Some(map) = self.chain_to_feed_conn_ids.get_mut(old_genesis_hash) {
|
||||
map.remove(&feed_conn_id);
|
||||
}
|
||||
}
|
||||
|
||||
// Untoggle request for finality feeds:
|
||||
self.feed_conn_id_finality.remove(&feed_conn_id);
|
||||
|
||||
// Get old chain if there was one:
|
||||
let node_state = &self.node_state;
|
||||
let old_chain =
|
||||
old_genesis_hash.and_then(|hash| node_state.get_chain_by_genesis_hash(&hash));
|
||||
|
||||
// Get new chain, ignoring the rest if it doesn't exist.
|
||||
let new_chain = match self.node_state.get_chain_by_label(&chain) {
|
||||
Some(chain) => chain,
|
||||
None => return,
|
||||
};
|
||||
|
||||
// Send messages to the feed about this subscription:
|
||||
let mut feed_serializer = FeedMessageSerializer::new();
|
||||
if let Some(old_chain) = old_chain {
|
||||
feed_serializer.push(feed_message::UnsubscribedFrom(old_chain.label()));
|
||||
}
|
||||
feed_serializer.push(feed_message::SubscribedTo(new_chain.label()));
|
||||
feed_serializer.push(feed_message::TimeSync(time::now()));
|
||||
feed_serializer.push(feed_message::BestBlock(
|
||||
new_chain.best_block().height,
|
||||
new_chain.timestamp(),
|
||||
new_chain.average_block_time(),
|
||||
));
|
||||
feed_serializer.push(feed_message::BestFinalized(
|
||||
new_chain.finalized_block().height,
|
||||
new_chain.finalized_block().hash,
|
||||
));
|
||||
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 {
|
||||
if let Some(bytes) = feed_serializer.finalize() {
|
||||
let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes)).await;
|
||||
}
|
||||
}
|
||||
feed_serializer.push(feed_message::AddedNode(chain_node_id, node));
|
||||
feed_serializer.push(feed_message::FinalizedBlock(
|
||||
chain_node_id,
|
||||
node.finalized().height,
|
||||
node.finalized().hash,
|
||||
));
|
||||
if node.stale() {
|
||||
feed_serializer.push(feed_message::StaleNode(chain_node_id));
|
||||
}
|
||||
}
|
||||
if let Some(bytes) = feed_serializer.into_finalized() {
|
||||
let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes)).await;
|
||||
}
|
||||
|
||||
// Actually make a note of the new chain subsciption:
|
||||
let new_genesis_hash = *new_chain.genesis_hash();
|
||||
self.feed_conn_id_to_chain
|
||||
.insert(feed_conn_id, new_genesis_hash);
|
||||
self.chain_to_feed_conn_ids
|
||||
.entry(new_genesis_hash)
|
||||
.or_default()
|
||||
.insert(feed_conn_id);
|
||||
}
|
||||
FromFeedWebsocket::SendFinality => {
|
||||
self.feed_conn_id_finality.insert(feed_conn_id);
|
||||
}
|
||||
FromFeedWebsocket::NoMoreFinality => {
|
||||
self.feed_conn_id_finality.remove(&feed_conn_id);
|
||||
}
|
||||
FromFeedWebsocket::Disconnected => {
|
||||
// The feed has disconnected; clean up references to it:
|
||||
if let Some(chain) = self.feed_conn_id_to_chain.remove(&feed_conn_id) {
|
||||
self.chain_to_feed_conn_ids.remove(&chain);
|
||||
}
|
||||
self.feed_channels.remove(&feed_conn_id);
|
||||
self.feed_conn_id_finality.remove(&feed_conn_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove all of the node IDs provided and broadcast messages to feeds as needed.
|
||||
async fn remove_nodes_and_broadcast_result(
|
||||
&mut self,
|
||||
node_ids: impl IntoIterator<Item = NodeId>,
|
||||
) {
|
||||
// Group by chain to simplify the handling of feed messages:
|
||||
let mut node_ids_per_chain: HashMap<BlockHash, Vec<NodeId>> = HashMap::new();
|
||||
for node_id in node_ids.into_iter() {
|
||||
if let Some(chain) = self.node_state.get_chain_by_node_id(node_id) {
|
||||
node_ids_per_chain
|
||||
.entry(*chain.genesis_hash())
|
||||
.or_default()
|
||||
.push(node_id);
|
||||
}
|
||||
}
|
||||
|
||||
// Remove the nodes for each chain
|
||||
let mut feed_messages_for_all = FeedMessageSerializer::new();
|
||||
for (chain_label, node_ids) in node_ids_per_chain {
|
||||
let mut feed_messages_for_chain = FeedMessageSerializer::new();
|
||||
for node_id in node_ids {
|
||||
self.remove_node(
|
||||
node_id,
|
||||
&mut feed_messages_for_chain,
|
||||
&mut feed_messages_for_all,
|
||||
);
|
||||
}
|
||||
self.finalize_and_broadcast_to_chain_feeds(&chain_label, feed_messages_for_chain)
|
||||
.await;
|
||||
}
|
||||
self.finalize_and_broadcast_to_all_feeds(feed_messages_for_all)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Remove a single node by its ID, pushing any messages we'd want to send
|
||||
/// out to feeds onto the provided feed serializers. Doesn't actually send
|
||||
/// anything to the feeds; just updates state as needed.
|
||||
fn remove_node(
|
||||
&mut self,
|
||||
node_id: NodeId,
|
||||
feed_for_chain: &mut FeedMessageSerializer,
|
||||
feed_for_all: &mut FeedMessageSerializer,
|
||||
) {
|
||||
// Remove our top level association (this may already have been done).
|
||||
self.node_ids.remove_by_left(&node_id);
|
||||
|
||||
let removed_details = match self.node_state.remove_node(node_id) {
|
||||
Some(remove_details) => remove_details,
|
||||
None => {
|
||||
log::error!("Could not find node {:?}", node_id);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// The chain has been removed (no nodes left in it, or it was renamed):
|
||||
if removed_details.chain_node_count == 0 || removed_details.has_chain_label_changed {
|
||||
feed_for_all.push(feed_message::RemovedChain(&removed_details.old_chain_label));
|
||||
}
|
||||
|
||||
// If the chain still exists, tell everybody about the new label or updated node count:
|
||||
if removed_details.chain_node_count != 0 {
|
||||
feed_for_all.push(feed_message::AddedChain(
|
||||
&removed_details.new_chain_label,
|
||||
removed_details.chain_node_count,
|
||||
));
|
||||
}
|
||||
|
||||
// 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.get_chain_node_id().into(),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
/// Finalize a [`FeedMessageSerializer`] and broadcast the result to feeds for the chain.
|
||||
async fn finalize_and_broadcast_to_chain_feeds(
|
||||
&mut self,
|
||||
genesis_hash: &BlockHash,
|
||||
serializer: FeedMessageSerializer,
|
||||
) {
|
||||
if let Some(bytes) = serializer.into_finalized() {
|
||||
self.broadcast_to_chain_feeds(genesis_hash, ToFeedWebsocket::Bytes(bytes))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a message to all chain feeds.
|
||||
async fn broadcast_to_chain_feeds(
|
||||
&mut self,
|
||||
genesis_hash: &BlockHash,
|
||||
message: ToFeedWebsocket,
|
||||
) {
|
||||
if let Some(feeds) = self.chain_to_feed_conn_ids.get(genesis_hash) {
|
||||
for &feed_id in feeds {
|
||||
if let Some(chan) = self.feed_channels.get_mut(&feed_id) {
|
||||
let _ = chan.send(message.clone()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Finalize a [`FeedMessageSerializer`] and broadcast the result to all feeds
|
||||
async fn finalize_and_broadcast_to_all_feeds(&mut self, serializer: FeedMessageSerializer) {
|
||||
if let Some(bytes) = serializer.into_finalized() {
|
||||
self.broadcast_to_all_feeds(ToFeedWebsocket::Bytes(bytes))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a message to everybody.
|
||||
async fn broadcast_to_all_feeds(&mut self, message: ToFeedWebsocket) {
|
||||
for chan in self.feed_channels.values_mut() {
|
||||
let _ = chan.send(message.clone()).await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Finalize a [`FeedMessageSerializer`] and broadcast the result to chain finality feeds
|
||||
async fn finalize_and_broadcast_to_chain_finality_feeds(
|
||||
&mut self,
|
||||
genesis_hash: &BlockHash,
|
||||
serializer: FeedMessageSerializer,
|
||||
) {
|
||||
if let Some(bytes) = serializer.into_finalized() {
|
||||
self.broadcast_to_chain_finality_feeds(genesis_hash, ToFeedWebsocket::Bytes(bytes))
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a message to all chain finality feeds.
|
||||
async fn broadcast_to_chain_finality_feeds(
|
||||
&mut self,
|
||||
genesis_hash: &BlockHash,
|
||||
message: ToFeedWebsocket,
|
||||
) {
|
||||
if let Some(feeds) = self.chain_to_feed_conn_ids.get(genesis_hash) {
|
||||
// Get all feeds for the chain, but only broadcast to those feeds that
|
||||
// are also subscribed to receive finality updates.
|
||||
for &feed_id in feeds.union(&self.feed_conn_id_finality) {
|
||||
if let Some(chan) = self.feed_channels.get_mut(&feed_id) {
|
||||
let _ = chan.send(message.clone()).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
mod aggregator;
|
||||
mod inner_loop;
|
||||
|
||||
// Expose the various message types that can be worked with externally:
|
||||
pub use inner_loop::{FromFeedWebsocket, FromShardWebsocket, ToFeedWebsocket, ToShardWebsocket};
|
||||
|
||||
pub use aggregator::*;
|
||||
@@ -0,0 +1,234 @@
|
||||
//! This module provides a way of encoding the various messages that we'll
|
||||
//! send to subscribed feeds (browsers).
|
||||
|
||||
use serde::Serialize;
|
||||
use std::mem;
|
||||
|
||||
use crate::state::Node;
|
||||
use common::node_types::{
|
||||
BlockDetails, BlockHash, BlockNumber, NodeHardware, NodeIO, NodeStats, Timestamp,
|
||||
};
|
||||
use serde_json::to_writer;
|
||||
|
||||
type Address = Box<str>;
|
||||
type FeedNodeId = usize;
|
||||
|
||||
pub trait FeedMessage {
|
||||
const ACTION: u8;
|
||||
}
|
||||
|
||||
pub trait FeedMessageWrite: FeedMessage {
|
||||
fn write_to_feed(&self, ser: &mut FeedMessageSerializer);
|
||||
}
|
||||
|
||||
impl<T> FeedMessageWrite for T
|
||||
where
|
||||
T: FeedMessage + Serialize,
|
||||
{
|
||||
fn write_to_feed(&self, ser: &mut FeedMessageSerializer) {
|
||||
ser.write(self)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct FeedMessageSerializer {
|
||||
/// Current buffer,
|
||||
buffer: Vec<u8>,
|
||||
}
|
||||
|
||||
const BUFCAP: usize = 128;
|
||||
|
||||
impl FeedMessageSerializer {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
buffer: Vec::with_capacity(BUFCAP),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn push<Message>(&mut self, msg: Message)
|
||||
where
|
||||
Message: FeedMessageWrite,
|
||||
{
|
||||
let glue = match self.buffer.len() {
|
||||
0 => b'[',
|
||||
_ => b',',
|
||||
};
|
||||
|
||||
self.buffer.push(glue);
|
||||
self.write(&Message::ACTION);
|
||||
self.buffer.push(b',');
|
||||
msg.write_to_feed(self);
|
||||
}
|
||||
|
||||
fn write<S>(&mut self, value: &S)
|
||||
where
|
||||
S: Serialize,
|
||||
{
|
||||
let _ = to_writer(&mut self.buffer, value);
|
||||
}
|
||||
|
||||
/// Return the bytes we've serialized so far and prepare a new buffer. If you're
|
||||
/// finished serializing data, prefer [`FeedMessageSerializer::into_finalized`]
|
||||
pub fn finalize(&mut self) -> Option<Vec<u8>> {
|
||||
if self.buffer.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
self.buffer.push(b']');
|
||||
|
||||
let bytes = mem::replace(&mut self.buffer, Vec::with_capacity(BUFCAP));
|
||||
|
||||
Some(bytes)
|
||||
}
|
||||
|
||||
/// Return the bytes that we've serialized so far, consuming the serializer.
|
||||
pub fn into_finalized(mut self) -> Option<Vec<u8>> {
|
||||
if self.buffer.is_empty() {
|
||||
return None;
|
||||
}
|
||||
|
||||
self.buffer.push(b']');
|
||||
Some(self.buffer)
|
||||
}
|
||||
}
|
||||
|
||||
macro_rules! actions {
|
||||
($($action:literal: $t:ty,)*) => {
|
||||
$(
|
||||
impl FeedMessage for $t {
|
||||
const ACTION: u8 = $action;
|
||||
}
|
||||
)*
|
||||
}
|
||||
}
|
||||
|
||||
actions! {
|
||||
0: Version,
|
||||
1: BestBlock,
|
||||
2: BestFinalized,
|
||||
3: AddedNode<'_>,
|
||||
4: RemovedNode,
|
||||
5: LocatedNode<'_>,
|
||||
6: ImportedBlock<'_>,
|
||||
7: FinalizedBlock,
|
||||
8: NodeStatsUpdate<'_>,
|
||||
9: Hardware<'_>,
|
||||
10: TimeSync,
|
||||
11: AddedChain<'_>,
|
||||
12: RemovedChain<'_>,
|
||||
13: SubscribedTo<'_>,
|
||||
14: UnsubscribedFrom<'_>,
|
||||
15: Pong<'_>,
|
||||
16: AfgFinalized,
|
||||
17: AfgReceivedPrevote,
|
||||
18: AfgReceivedPrecommit,
|
||||
19: AfgAuthoritySet,
|
||||
20: StaleNode,
|
||||
21: NodeIOUpdate<'_>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct Version(pub usize);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct BestBlock(pub BlockNumber, pub Timestamp, pub Option<u64>);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct BestFinalized(pub BlockNumber, pub BlockHash);
|
||||
|
||||
pub struct AddedNode<'a>(pub FeedNodeId, pub &'a Node);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct RemovedNode(pub FeedNodeId);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct LocatedNode<'a>(pub FeedNodeId, pub f32, pub f32, pub &'a str);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct ImportedBlock<'a>(pub FeedNodeId, pub &'a BlockDetails);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct FinalizedBlock(pub FeedNodeId, pub BlockNumber, pub BlockHash);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct NodeStatsUpdate<'a>(pub FeedNodeId, pub &'a NodeStats);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct NodeIOUpdate<'a>(pub FeedNodeId, pub &'a NodeIO);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct Hardware<'a>(pub FeedNodeId, pub &'a NodeHardware);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct TimeSync(pub u64);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct AddedChain<'a>(pub &'a str, pub usize);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct RemovedChain<'a>(pub &'a str);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct SubscribedTo<'a>(pub &'a str);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct UnsubscribedFrom<'a>(pub &'a str);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct Pong<'a>(pub &'a str);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct AfgFinalized(pub Address, pub BlockNumber, pub BlockHash);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct AfgReceivedPrevote(
|
||||
pub Address,
|
||||
pub BlockNumber,
|
||||
pub BlockHash,
|
||||
pub Option<Address>,
|
||||
);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct AfgReceivedPrecommit(
|
||||
pub Address,
|
||||
pub BlockNumber,
|
||||
pub BlockHash,
|
||||
pub Option<Address>,
|
||||
);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct AfgAuthoritySet(
|
||||
pub Address,
|
||||
pub Address,
|
||||
pub Address,
|
||||
pub BlockNumber,
|
||||
pub BlockHash,
|
||||
);
|
||||
|
||||
#[derive(Serialize)]
|
||||
pub struct StaleNode(pub FeedNodeId);
|
||||
|
||||
impl FeedMessageWrite for AddedNode<'_> {
|
||||
fn write_to_feed(&self, ser: &mut FeedMessageSerializer) {
|
||||
let AddedNode(nid, node) = self;
|
||||
|
||||
let details = node.details();
|
||||
let details = (
|
||||
&details.name,
|
||||
&details.implementation,
|
||||
&details.version,
|
||||
&details.validator,
|
||||
&details.network_id,
|
||||
);
|
||||
|
||||
ser.write(&(
|
||||
nid,
|
||||
details,
|
||||
node.stats(),
|
||||
node.io(),
|
||||
node.hardware(),
|
||||
node.block_details(),
|
||||
&node.location(),
|
||||
&node.startup_time(),
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,211 @@
|
||||
use std::net::Ipv4Addr;
|
||||
use std::sync::Arc;
|
||||
|
||||
use futures::channel::mpsc;
|
||||
use futures::{Sink, SinkExt, StreamExt};
|
||||
use parking_lot::RwLock;
|
||||
use rustc_hash::FxHashMap;
|
||||
use serde::Deserialize;
|
||||
|
||||
use common::node_types::NodeLocation;
|
||||
use tokio::sync::Semaphore;
|
||||
|
||||
/// The returned location is optional; it may be None if not found.
|
||||
pub type Location = Option<Arc<NodeLocation>>;
|
||||
|
||||
/// This is responsible for taking an IP address and attempting
|
||||
/// to find a geographical location from this
|
||||
pub fn find_location<Id, R>(response_chan: R) -> mpsc::UnboundedSender<(Id, Ipv4Addr)>
|
||||
where
|
||||
R: Sink<(Id, Option<Arc<NodeLocation>>)> + Unpin + Send + Clone + 'static,
|
||||
Id: Clone + Send + 'static,
|
||||
{
|
||||
let (tx, mut rx) = mpsc::unbounded();
|
||||
|
||||
// cache entries
|
||||
let mut cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>> = FxHashMap::default();
|
||||
|
||||
// Default entry for localhost
|
||||
cache.insert(
|
||||
Ipv4Addr::new(127, 0, 0, 1),
|
||||
Some(Arc::new(NodeLocation {
|
||||
latitude: 52.516_6667,
|
||||
longitude: 13.4,
|
||||
city: "Berlin".into(),
|
||||
})),
|
||||
);
|
||||
|
||||
// Create a locator with our cache. This is used to obtain locations.
|
||||
let locator = Locator::new(cache);
|
||||
|
||||
// Spawn a loop to handle location requests
|
||||
tokio::spawn(async move {
|
||||
// Allow 4 requests at a time. acquiring a token will block while the
|
||||
// number of concurrent location requests is more than this.
|
||||
let semaphore = Arc::new(Semaphore::new(4));
|
||||
|
||||
loop {
|
||||
while let Some((id, ip_address)) = rx.next().await {
|
||||
let permit = semaphore.clone().acquire_owned().await.unwrap();
|
||||
let mut response_chan = response_chan.clone();
|
||||
let locator = locator.clone();
|
||||
|
||||
// Once we have acquired our permit, spawn a task to avoid
|
||||
// blocking this loop so that we can handle concurrent requests.
|
||||
tokio::spawn(async move {
|
||||
match locator.locate(ip_address).await {
|
||||
Ok(loc) => {
|
||||
let _ = response_chan.send((id, loc)).await;
|
||||
}
|
||||
Err(e) => {
|
||||
log::debug!("GET error for ip location: {:?}", e);
|
||||
}
|
||||
};
|
||||
|
||||
// ensure permit is moved into task by dropping it explicitly:
|
||||
drop(permit);
|
||||
});
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
tx
|
||||
}
|
||||
|
||||
/// This struct can be used to make location requests, given
|
||||
/// an IPV4 address.
|
||||
#[derive(Clone)]
|
||||
struct Locator {
|
||||
client: reqwest::Client,
|
||||
cache: Arc<RwLock<FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>>>>,
|
||||
}
|
||||
|
||||
impl Locator {
|
||||
pub fn new(cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>>) -> Self {
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
Locator {
|
||||
client,
|
||||
cache: Arc::new(RwLock::new(cache)),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn locate(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
|
||||
// Return location quickly if it's cached:
|
||||
let cached_loc = {
|
||||
let cache_reader = self.cache.read();
|
||||
cache_reader.get(&ip).map(|o| o.clone())
|
||||
};
|
||||
if let Some(loc) = cached_loc {
|
||||
return Ok(loc);
|
||||
}
|
||||
|
||||
// Look it up via the location services if not cached:
|
||||
let location = self.iplocate_ipapi_co(ip).await?;
|
||||
let location = match location {
|
||||
Some(location) => Ok(Some(location)),
|
||||
None => self.iplocate_ipinfo_io(ip).await,
|
||||
}?;
|
||||
|
||||
self.cache.write().insert(ip, location.clone());
|
||||
Ok(location)
|
||||
}
|
||||
|
||||
async fn iplocate_ipapi_co(
|
||||
&self,
|
||||
ip: Ipv4Addr,
|
||||
) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
|
||||
let location = self
|
||||
.query(&format!("https://ipapi.co/{}/json", ip))
|
||||
.await?
|
||||
.map(Arc::new);
|
||||
|
||||
Ok(location)
|
||||
}
|
||||
|
||||
async fn iplocate_ipinfo_io(
|
||||
&self,
|
||||
ip: Ipv4Addr,
|
||||
) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
|
||||
let location = self
|
||||
.query(&format!("https://ipinfo.io/{}/json", ip))
|
||||
.await?
|
||||
.and_then(|loc: IPApiLocate| loc.into_node_location().map(Arc::new));
|
||||
|
||||
Ok(location)
|
||||
}
|
||||
|
||||
async fn query<T>(&self, url: &str) -> Result<Option<T>, reqwest::Error>
|
||||
where
|
||||
for<'de> T: Deserialize<'de>,
|
||||
{
|
||||
match self.client.get(url).send().await?.json::<T>().await {
|
||||
Ok(result) => Ok(Some(result)),
|
||||
Err(err) => {
|
||||
log::debug!("JSON error for ip location: {:?}", err);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// This is the format returned from ipinfo.co, so we do
|
||||
/// a little conversion to get it into the shape we want.
|
||||
#[derive(Deserialize)]
|
||||
struct IPApiLocate {
|
||||
city: Box<str>,
|
||||
loc: Box<str>,
|
||||
}
|
||||
|
||||
impl IPApiLocate {
|
||||
fn into_node_location(self) -> Option<NodeLocation> {
|
||||
let IPApiLocate { city, loc } = self;
|
||||
|
||||
let mut loc = loc.split(',').map(|n| n.parse());
|
||||
|
||||
let latitude = loc.next()?.ok()?;
|
||||
let longitude = loc.next()?.ok()?;
|
||||
|
||||
// Guarantee that the iterator has been exhausted
|
||||
if loc.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(NodeLocation {
|
||||
latitude,
|
||||
longitude,
|
||||
city,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn ipapi_locate_to_node_location() {
|
||||
let ipapi = IPApiLocate {
|
||||
loc: "12.5,56.25".into(),
|
||||
city: "Foobar".into(),
|
||||
};
|
||||
|
||||
let location = ipapi.into_node_location().unwrap();
|
||||
|
||||
assert_eq!(location.latitude, 12.5);
|
||||
assert_eq!(location.longitude, 56.25);
|
||||
assert_eq!(&*location.city, "Foobar");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ipapi_locate_to_node_location_too_many() {
|
||||
let ipapi = IPApiLocate {
|
||||
loc: "12.5,56.25,1.0".into(),
|
||||
city: "Foobar".into(),
|
||||
};
|
||||
|
||||
let location = ipapi.into_node_location();
|
||||
|
||||
assert!(location.is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,310 @@
|
||||
mod aggregator;
|
||||
mod feed_message;
|
||||
mod find_location;
|
||||
mod state;
|
||||
|
||||
use std::net::SocketAddr;
|
||||
use std::str::FromStr;
|
||||
|
||||
use aggregator::{
|
||||
Aggregator, FromFeedWebsocket, FromShardWebsocket, ToFeedWebsocket, ToShardWebsocket,
|
||||
};
|
||||
use bincode::Options;
|
||||
use common::internal_messages;
|
||||
use futures::{channel::mpsc, SinkExt, StreamExt};
|
||||
use simple_logger::SimpleLogger;
|
||||
use structopt::StructOpt;
|
||||
use warp::filters::ws;
|
||||
use warp::Filter;
|
||||
|
||||
const VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
|
||||
const NAME: &str = "Substrate Telemetry Backend Core";
|
||||
const ABOUT: &str = "This is the Telemetry Backend Core that receives telemetry messages \
|
||||
from Substrate/Polkadot nodes and provides the data to a subsribed feed";
|
||||
|
||||
#[derive(StructOpt, Debug)]
|
||||
#[structopt(name = NAME, version = VERSION, author = AUTHORS, about = ABOUT)]
|
||||
struct Opts {
|
||||
/// This is the socket address that Telemetryis listening to. This is restricted to
|
||||
/// localhost (127.0.0.1) by default and should be fine for most use cases. If
|
||||
/// you are using Telemetry in a container, you likely want to set this to '0.0.0.0:8000'
|
||||
#[structopt(short = "l", long = "listen", default_value = "127.0.0.1:8000")]
|
||||
socket: std::net::SocketAddr,
|
||||
/// The desired log level; one of 'error', 'warn', 'info', 'debug' or 'trace', where
|
||||
/// 'error' only logs errors and 'trace' logs everything.
|
||||
#[structopt(required = false, long = "log", default_value = "info")]
|
||||
log_level: log::LevelFilter,
|
||||
/// Space delimited list of the names of chains that are not allowed to connect to
|
||||
/// telemetry. Case sensitive.
|
||||
#[structopt(required = false, long = "denylist")]
|
||||
denylist: Vec<String>,
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
let opts = Opts::from_args();
|
||||
|
||||
SimpleLogger::new()
|
||||
.with_level(opts.log_level)
|
||||
.init()
|
||||
.expect("Must be able to start a logger");
|
||||
|
||||
log::info!("Starting Telemetry Core version: {}", VERSION);
|
||||
|
||||
if let Err(e) = start_server(opts).await {
|
||||
log::error!("Error starting server: {}", e);
|
||||
}
|
||||
}
|
||||
|
||||
/// Declare our routes and start the server.
|
||||
async fn start_server(opts: Opts) -> anyhow::Result<()> {
|
||||
let shard_aggregator = Aggregator::spawn(opts.denylist).await?;
|
||||
let feed_aggregator = shard_aggregator.clone();
|
||||
|
||||
// Handle requests to /health by returning OK.
|
||||
let health_route = warp::path("health").map(|| "OK");
|
||||
|
||||
// Handle websocket requests from shards.
|
||||
let ws_shard_submit_route = warp::path("shard_submit")
|
||||
.and(warp::ws())
|
||||
.and(warp::filters::addr::remote())
|
||||
.map(move |ws: ws::Ws, addr: Option<SocketAddr>| {
|
||||
let tx_to_aggregator = shard_aggregator.subscribe_shard();
|
||||
log::info!("Opening /shard_submit connection from {:?}", addr);
|
||||
ws.on_upgrade(move |websocket| async move {
|
||||
let (mut tx_to_aggregator, websocket) =
|
||||
handle_shard_websocket_connection(websocket, tx_to_aggregator).await;
|
||||
log::info!("Closing /shard_submit connection from {:?}", addr);
|
||||
// Tell the aggregator that this connection has closed, so it can tidy up.
|
||||
let _ = tx_to_aggregator
|
||||
.send(FromShardWebsocket::Disconnected)
|
||||
.await;
|
||||
let _ = websocket.close().await;
|
||||
})
|
||||
});
|
||||
|
||||
// Handle websocket requests from frontends.
|
||||
let ws_feed_route = warp::path("feed")
|
||||
.and(warp::ws())
|
||||
.and(warp::filters::addr::remote())
|
||||
.map(move |ws: ws::Ws, addr: Option<SocketAddr>| {
|
||||
let tx_to_aggregator = feed_aggregator.subscribe_feed();
|
||||
log::info!("Opening /feed connection from {:?}", addr);
|
||||
|
||||
// We can decide how many messages can be buffered to be sent, but not specifically how
|
||||
// large those messages are cumulatively allowed to be:
|
||||
ws.max_send_queue(1_000 ).on_upgrade(move |websocket| async move {
|
||||
let (mut tx_to_aggregator, websocket) =
|
||||
handle_feed_websocket_connection(websocket, tx_to_aggregator).await;
|
||||
log::info!("Closing /feed connection from {:?}", addr);
|
||||
// Tell the aggregator that this connection has closed, so it can tidy up.
|
||||
let _ = tx_to_aggregator.send(FromFeedWebsocket::Disconnected).await;
|
||||
let _ = websocket.close().await;
|
||||
})
|
||||
});
|
||||
|
||||
// Merge the routes and start our server:
|
||||
let routes = ws_shard_submit_route.or(ws_feed_route).or(health_route);
|
||||
warp::serve(routes).run(opts.socket).await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// This handles messages coming to/from a shard connection
|
||||
async fn handle_shard_websocket_connection<S>(
|
||||
mut websocket: ws::WebSocket,
|
||||
mut tx_to_aggregator: S,
|
||||
) -> (S, ws::WebSocket)
|
||||
where
|
||||
S: futures::Sink<FromShardWebsocket, Error = anyhow::Error> + Unpin,
|
||||
{
|
||||
let (tx_to_shard_conn, mut rx_from_aggregator) = mpsc::channel(10);
|
||||
|
||||
// Tell the aggregator about this new connection, and give it a way to send messages to us:
|
||||
let init_msg = FromShardWebsocket::Initialize {
|
||||
channel: tx_to_shard_conn,
|
||||
};
|
||||
if let Err(e) = tx_to_aggregator.send(init_msg).await {
|
||||
log::error!("Error sending message to aggregator: {}", e);
|
||||
return (tx_to_aggregator, websocket);
|
||||
}
|
||||
|
||||
// Loop, handling new messages from the shard or from the aggregator:
|
||||
loop {
|
||||
tokio::select! {
|
||||
// AGGREGATOR -> SHARD
|
||||
msg = rx_from_aggregator.next() => {
|
||||
// End the loop when connection from aggregator ends:
|
||||
let msg = match msg {
|
||||
Some(msg) => msg,
|
||||
None => break
|
||||
};
|
||||
|
||||
let internal_msg = match msg {
|
||||
ToShardWebsocket::Mute { local_id, reason } => {
|
||||
internal_messages::FromTelemetryCore::Mute { local_id, reason }
|
||||
}
|
||||
};
|
||||
|
||||
let bytes = bincode::options()
|
||||
.serialize(&internal_msg)
|
||||
.expect("message to shard should serialize");
|
||||
|
||||
if let Err(e) = websocket.send(ws::Message::binary(bytes)).await {
|
||||
log::error!("Error sending message to shard; booting it: {}", e);
|
||||
break
|
||||
}
|
||||
}
|
||||
// SHARD -> AGGREGATOR
|
||||
msg = websocket.next() => {
|
||||
// End the loop when connection from shard ends:
|
||||
let msg = match msg {
|
||||
Some(msg) => msg,
|
||||
None => break
|
||||
};
|
||||
|
||||
let msg = match msg {
|
||||
Err(e) => {
|
||||
log::error!("Error receiving message from shard; booting it: {}", e);
|
||||
break;
|
||||
},
|
||||
Ok(msg) => msg
|
||||
};
|
||||
|
||||
// Close message? Break and allow connection to be dropped.
|
||||
if msg.is_close() {
|
||||
break;
|
||||
}
|
||||
|
||||
// If the message isn't something we want to handle, just ignore it.
|
||||
// This includes system messages like "pings" and such, so don't log anything.
|
||||
if !msg.is_binary() && !msg.is_text() {
|
||||
continue;
|
||||
}
|
||||
|
||||
let bytes = msg.as_bytes();
|
||||
let msg: internal_messages::FromShardAggregator = match bincode::options().deserialize(bytes) {
|
||||
Ok(msg) => msg,
|
||||
Err(e) => {
|
||||
log::error!("Failed to deserialize message from shard; booting it: {}", e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
// Convert and send to the aggregator:
|
||||
let aggregator_msg = match msg {
|
||||
internal_messages::FromShardAggregator::AddNode { ip, node, local_id, genesis_hash } => {
|
||||
FromShardWebsocket::Add { ip, node, genesis_hash, local_id }
|
||||
},
|
||||
internal_messages::FromShardAggregator::UpdateNode { payload, local_id } => {
|
||||
FromShardWebsocket::Update { local_id, payload }
|
||||
},
|
||||
internal_messages::FromShardAggregator::RemoveNode { local_id } => {
|
||||
FromShardWebsocket::Remove { local_id }
|
||||
},
|
||||
};
|
||||
if let Err(e) = tx_to_aggregator.send(aggregator_msg).await {
|
||||
log::error!("Failed to send message to aggregator; closing shard: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// loop ended; give socket back to parent:
|
||||
(tx_to_aggregator, websocket)
|
||||
}
|
||||
|
||||
/// This handles messages coming from a feed connection
|
||||
async fn handle_feed_websocket_connection<S>(
|
||||
mut websocket: ws::WebSocket,
|
||||
mut tx_to_aggregator: S,
|
||||
) -> (S, ws::WebSocket)
|
||||
where
|
||||
S: futures::Sink<FromFeedWebsocket, Error = anyhow::Error> + Unpin,
|
||||
{
|
||||
// unbounded channel so that slow feeds don't block aggregator progress:
|
||||
let (tx_to_feed_conn, mut rx_from_aggregator) = mpsc::unbounded();
|
||||
|
||||
// Tell the aggregator about this new connection, and give it a way to send messages to us:
|
||||
let init_msg = FromFeedWebsocket::Initialize {
|
||||
channel: tx_to_feed_conn,
|
||||
};
|
||||
if let Err(e) = tx_to_aggregator.send(init_msg).await {
|
||||
log::error!("Error sending message to aggregator: {}", e);
|
||||
return (tx_to_aggregator, websocket);
|
||||
}
|
||||
|
||||
// Loop, handling new messages from the shard or from the aggregator:
|
||||
loop {
|
||||
tokio::select! {
|
||||
|
||||
// AGGREGATOR -> FRONTEND (buffer messages to the UI)
|
||||
msg = rx_from_aggregator.next() => {
|
||||
// End the loop when connection from aggregator ends:
|
||||
let msg = match msg {
|
||||
Some(msg) => msg,
|
||||
None => break
|
||||
};
|
||||
|
||||
// Send messages to the client (currently the only message is
|
||||
// pre-serialized bytes that we send as binary):
|
||||
let bytes = match msg {
|
||||
ToFeedWebsocket::Bytes(bytes) => bytes
|
||||
};
|
||||
|
||||
log::debug!("Message to feed: {}", std::str::from_utf8(&bytes).unwrap_or("INVALID UTF8"));
|
||||
|
||||
if let Err(e) = websocket.send(ws::Message::binary(bytes)).await {
|
||||
log::warn!("Closing feed websocket due to error: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// FRONTEND -> AGGREGATOR (relay messages to the aggregator)
|
||||
msg = websocket.next() => {
|
||||
// End the loop when connection from feed ends:
|
||||
let msg = match msg {
|
||||
Some(msg) => msg,
|
||||
None => break
|
||||
};
|
||||
|
||||
// If we see any errors, log them and end our loop:
|
||||
let msg = match msg {
|
||||
Err(e) => {
|
||||
log::error!("Error in node websocket connection: {}", e);
|
||||
break;
|
||||
},
|
||||
Ok(msg) => msg
|
||||
};
|
||||
|
||||
// Close message? Break and allow connection to be dropped.
|
||||
if msg.is_close() {
|
||||
break;
|
||||
}
|
||||
|
||||
// We ignore all but text messages from the frontend:
|
||||
let text = match msg.to_str() {
|
||||
Ok(s) => s,
|
||||
Err(_) => continue
|
||||
};
|
||||
|
||||
// Parse the message into a command we understand and send it to the aggregator:
|
||||
let cmd = match FromFeedWebsocket::from_str(text) {
|
||||
Ok(cmd) => cmd,
|
||||
Err(e) => {
|
||||
log::warn!("Ignoring invalid command '{}' from the frontend: {}", text, e);
|
||||
continue
|
||||
}
|
||||
};
|
||||
if let Err(e) = tx_to_aggregator.send(cmd).await {
|
||||
log::error!("Failed to send message to aggregator; closing feed: {}", e);
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// loop ended; give socket back to parent:
|
||||
(tx_to_aggregator, websocket)
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
mod chain;
|
||||
mod node;
|
||||
|
||||
mod state;
|
||||
|
||||
pub use node::Node;
|
||||
pub use state::*;
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user