pezkuwichain/pezkuwi-telemetry
1
0
Fork 0
mirror of https://github.com/pezkuwichain/pezkuwi-telemetry.git synced 2026-04-26 01:47:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

use flume throughout telemetry_core

Browse Source
This commit is contained in:
James Wilson
2021-08-10 10:26:13 +01:00
parent 8268cf2afe
commit 703a9ddc4e
9 changed files with 98 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/telemetry_core/src/aggregator/aggregator.rs
+8 -9
View File
@@ -18,7 +18,6 @@ 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;
@@ -54,17 +53,17 @@ struct AggregatorInternal {
/// 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::UnboundedSender<inner_loop::ToAggregator>,
tx_to_aggregator: flume::Sender<inner_loop::ToAggregator>,
}
impl Aggregator {
/// Spawn a new Aggregator. This connects to the telemetry backend
pub async fn spawn(opts: AggregatorOpts) -> anyhow::Result<Aggregator> {
let (tx_to_aggregator, rx_from_external) = mpsc::unbounded();
let (tx_to_aggregator, rx_from_external) = flume::unbounded();
// 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(
let tx_to_locator = find_location(tx_to_aggregator.clone().into_sink().with(|(node_id, msg)| {
future::ok::<_, flume::SendError<_>>(inner_loop::ToAggregator::FromFindLocation(
node_id, msg,
))
}));
@@ -89,8 +88,8 @@ impl Aggregator {
/// 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::UnboundedReceiver<inner_loop::ToAggregator>,
tx_to_aggregator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>,
rx_from_external: flume::Receiver<inner_loop::ToAggregator>,
tx_to_aggregator: flume::Sender<(NodeId, Ipv4Addr)>,
max_queue_len: usize,
denylist: Vec<String>,
) {
@@ -114,7 +113,7 @@ impl Aggregator {
// 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 {
Box::pin(tx_to_aggregator.into_sink().with(move |msg| async move {
Ok(inner_loop::ToAggregator::FromShardWebsocket(
shard_conn_id.into(),
msg,
@@ -141,7 +140,7 @@ impl Aggregator {
// but pinning by boxing is the easy solution for now:
(
feed_conn_id,
Box::pin(tx_to_aggregator.with(move |msg| async move {
Box::pin(tx_to_aggregator.into_sink().with(move |msg| async move {
Ok(inner_loop::ToAggregator::FromFeedWebsocket(
feed_conn_id.into(),
msg,
backend/telemetry_core/src/aggregator/inner_loop.rs
+27 -25
View File
@@ -20,13 +20,11 @@ use crate::find_location;
use crate::state::{self, NodeId, State};
use bimap::BiMap;
use common::{
channel::metered_unbounded,
internal_messages::{self, MuteReason, ShardNodeId},
node_message,
node_types::BlockHash,
time,
};
use futures::channel::mpsc;
use futures::StreamExt;
use std::collections::{HashMap, HashSet};
use std::{
@@ -48,7 +46,7 @@ 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::UnboundedSender<ToShardWebsocket>,
channel: flume::Sender<ToShardWebsocket>,
},
/// Tell the aggregator about a new node.
Add {
@@ -86,7 +84,7 @@ pub enum FromFeedWebsocket {
/// Unbounded so that slow feeds don't block aggregato
/// progress.
Initialize {
channel: mpsc::UnboundedSender<ToFeedWebsocket>,
channel: flume::Sender<ToFeedWebsocket>,
},
/// The feed can subscribe to a chain to receive
/// messages relating to it.
@@ -135,9 +133,9 @@ pub struct InnerLoop {
node_ids: BiMap<NodeId, (ConnId, ShardNodeId)>,
/// Keep track of how to send messages out to feeds.
feed_channels: HashMap<ConnId, mpsc::UnboundedSender<ToFeedWebsocket>>,
feed_channels: HashMap<ConnId, flume::Sender<ToFeedWebsocket>>,
/// Keep track of how to send messages out to shards.
shard_channels: HashMap<ConnId, mpsc::UnboundedSender<ToShardWebsocket>>,
shard_channels: HashMap<ConnId, flume::Sender<ToShardWebsocket>>,
/// Which chain is a feed subscribed to?
/// Feed Connection ID -> Chain Genesis Hash
@@ -150,7 +148,7 @@ pub struct InnerLoop {
feed_conn_id_finality: HashSet<ConnId>,
/// Send messages here to make geographical location requests.
tx_to_locator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>,
tx_to_locator: flume::Sender<(NodeId, Ipv4Addr)>,
/// How big can the queue of messages coming in to the aggregator get before messages
/// are prioritised and dropped to try and get back on track.
@@ -160,7 +158,7 @@ pub struct InnerLoop {
impl InnerLoop {
/// Create a new inner loop handler with the various state it needs.
pub fn new(
tx_to_locator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>,
tx_to_locator: flume::Sender<(NodeId, Ipv4Addr)>,
denylist: Vec<String>,
max_queue_len: usize,
) -> Self {
@@ -177,14 +175,16 @@ impl InnerLoop {
}
}
/// Start handling and responding to incoming messages. Owing to unbounded channels, we actually
/// only have a single `.await` (in this function). This helps to make it clear that the aggregator loop
/// will be able to make progress quickly without any potential yield points.
pub async fn handle(mut self, mut rx_from_external: mpsc::UnboundedReceiver<ToAggregator>) {
/// Start handling and responding to incoming messages.
pub async fn handle(mut self, rx_from_external: flume::Receiver<ToAggregator>) {
let max_queue_len = self.max_queue_len;
let (metered_tx, mut metered_rx) = metered_unbounded();
let (metered_tx, metered_rx) = flume::unbounded();
// Actually handle all of our messages, but before we get here, we
// check the length of the queue below to decide whether or not to
// pass the message on to this.
tokio::spawn(async move {
let mut metered_rx = metered_rx.into_stream();
while let Some(msg) = metered_rx.next().await {
match msg {
ToAggregator::FromFeedWebsocket(feed_conn_id, msg) => {
@@ -215,9 +215,11 @@ impl InnerLoop {
});
});
let mut rx_from_external = rx_from_external.into_stream();
while let Some(msg) = rx_from_external.next().await {
// ignore node updates if we have too many messages to handle, in an attempt
// to reduce the queue length back to something reasonable.
// to reduce the queue length back to something reasonable, lest it get out of
// control and start consuming a load of memory.
if metered_tx.len() > max_queue_len {
if matches!(
msg,
@@ -227,7 +229,7 @@ impl InnerLoop {
}
}
if let Err(e) = metered_tx.unbounded_send(msg) {
if let Err(e) = metered_tx.send(msg) {
log::error!("Cannot send message into aggregator: {}", e);
break;
}
@@ -277,7 +279,7 @@ impl InnerLoop {
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.unbounded_send(ToShardWebsocket::Mute {
let _ = shard_conn.send(ToShardWebsocket::Mute {
local_id,
reason: MuteReason::ChainNotAllowed,
});
@@ -285,7 +287,7 @@ impl InnerLoop {
}
state::AddNodeResult::ChainOverQuota => {
if let Some(shard_conn) = self.shard_channels.get_mut(&shard_conn_id) {
let _ = shard_conn.unbounded_send(ToShardWebsocket::Mute {
let _ = shard_conn.send(ToShardWebsocket::Mute {
local_id,
reason: MuteReason::Overquota,
});
@@ -326,7 +328,7 @@ impl InnerLoop {
// Ask for the grographical location of the node.
// Currently we only geographically locate IPV4 addresses so ignore IPV6.
if let IpAddr::V4(ip_v4) = ip {
let _ = self.tx_to_locator.unbounded_send((node_id, ip_v4));
let _ = self.tx_to_locator.send((node_id, ip_v4));
}
}
}
@@ -409,7 +411,7 @@ impl InnerLoop {
// Send this to the channel that subscribed:
if let Some(bytes) = feed_serializer.into_finalized() {
let _ = channel.unbounded_send(ToFeedWebsocket::Bytes(bytes));
let _ = channel.send(ToFeedWebsocket::Bytes(bytes));
}
}
FromFeedWebsocket::Ping { value } => {
@@ -422,7 +424,7 @@ impl InnerLoop {
let mut feed_serializer = FeedMessageSerializer::new();
feed_serializer.push(feed_message::Pong(&value));
if let Some(bytes) = feed_serializer.into_finalized() {
let _ = feed_channel.unbounded_send(ToFeedWebsocket::Bytes(bytes));
let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes));
}
}
FromFeedWebsocket::Subscribe { chain } => {
@@ -470,7 +472,7 @@ impl InnerLoop {
new_chain.finalized_block().hash,
));
if let Some(bytes) = feed_serializer.into_finalized() {
let _ = feed_channel.unbounded_send(ToFeedWebsocket::Bytes(bytes));
let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes));
}
// If many (eg 10k) nodes are connected, serializing all of their info takes time.
@@ -505,7 +507,7 @@ impl InnerLoop {
})
.collect();
for bytes in all_feed_messages {
let _ = feed_channel.unbounded_send(ToFeedWebsocket::Bytes(bytes));
let _ = feed_channel.send(ToFeedWebsocket::Bytes(bytes));
}
// Actually make a note of the new chain subsciption:
@@ -620,7 +622,7 @@ impl InnerLoop {
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.unbounded_send(message.clone());
let _ = chan.send(message.clone());
}
}
}
@@ -636,7 +638,7 @@ impl InnerLoop {
/// Send a message to everybody.
fn broadcast_to_all_feeds(&mut self, message: ToFeedWebsocket) {
for chan in self.feed_channels.values_mut() {
let _ = chan.unbounded_send(message.clone());
let _ = chan.send(message.clone());
}
}
@@ -662,7 +664,7 @@ impl InnerLoop {
// 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.unbounded_send(message.clone());
let _ = chan.send(message.clone());
}
}
}
backend/telemetry_core/src/find_location.rs
+3 -3
View File
@@ -17,7 +17,6 @@
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;
@@ -31,12 +30,13 @@ 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)>
pub fn find_location<Id, R>(response_chan: R) -> flume::Sender<(Id, Ipv4Addr)>
where
R: Sink<(Id, Option<Arc<NodeLocation>>)> + Unpin + Send + Clone + 'static,
Id: Clone + Send + 'static,
{
let (tx, mut rx) = mpsc::unbounded();
let (tx, rx) = flume::unbounded();
let mut rx = rx.into_stream();
// cache entries
let mut cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>> = FxHashMap::default();
backend/telemetry_core/src/main.rs
+5 -4
View File
@@ -29,7 +29,7 @@ use bincode::Options;
use common::http_utils;
use common::internal_messages;
use common::ready_chunks_all::ReadyChunksAll;
use futures::{channel::mpsc, SinkExt, StreamExt};
use futures::{SinkExt, StreamExt};
use hyper::{Method, Response};
use simple_logger::SimpleLogger;
use structopt::StructOpt;
@@ -201,7 +201,8 @@ async fn handle_shard_websocket_connection<S>(
where
S: futures::Sink<FromShardWebsocket, Error = anyhow::Error> + Unpin + Send + 'static,
{
let (tx_to_shard_conn, mut rx_from_aggregator) = mpsc::unbounded();
let (tx_to_shard_conn, rx_from_aggregator) = flume::unbounded();
let mut rx_from_aggregator = rx_from_aggregator.into_stream();
// Tell the aggregator about this new connection, and give it a way to send messages to us:
let init_msg = FromShardWebsocket::Initialize {
@@ -343,8 +344,8 @@ where
S: futures::Sink<FromFeedWebsocket, Error = anyhow::Error> + Unpin + Send + 'static,
{
// unbounded channel so that slow feeds don't block aggregator progress:
let (tx_to_feed_conn, rx_from_aggregator) = mpsc::unbounded();
let mut rx_from_aggregator_chunks = ReadyChunksAll::new(rx_from_aggregator);
let (tx_to_feed_conn, rx_from_aggregator) = flume::unbounded();
let mut rx_from_aggregator_chunks = ReadyChunksAll::new(rx_from_aggregator.into_stream());
// Tell the aggregator about this new connection, and give it a way to send messages to us:
let init_msg = FromFeedWebsocket::Initialize {