pezkuwichain/pezkuwi-telemetry
1
0
Fork 0
mirror of https://github.com/pezkuwichain/pezkuwi-telemetry.git synced 2026-06-14 20:51:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Migrate telemetry_core to Hyper+Soketto

Browse Source
This commit is contained in:
James Wilson
2021-07-23 15:57:50 +01:00
parent 649fb966d2
commit 0eff32d10e
7 changed files with 389 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/Cargo.lock
Generated
+9 -4
View File
@@ -20,9 +20,9 @@ dependencies = [
[[package]] [[package]]
name = "anyhow" name = "anyhow"
version = "1.0.41" version = "1.0.42"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "15af2628f6890fe2609a3b91bef4c83450512802e59489f9c1cb1fa5df064a61" checksum = "595d3cfa7a60d4555cb5067b99f07142a08ea778de5cf993f7b75c7d8fabc486"
[[package]] [[package]]
name = "arrayvec" name = "arrayvec"
@@ -205,6 +205,8 @@ dependencies = [
name = "common" name = "common"
version = "0.1.0" version = "0.1.0"
dependencies = [ dependencies = [
"anyhow",
"base64",
"bimap", "bimap",
"bincode", "bincode",
"bytes", "bytes",
@@ -212,6 +214,7 @@ dependencies = [
"futures", "futures",
"hex", "hex",
"http", "http",
"hyper",
"log", "log",
"num-traits", "num-traits",
"pin-project-lite", "pin-project-lite",
@@ -219,6 +222,7 @@ dependencies = [
"rustc-hash", "rustc-hash",
"serde", "serde",
"serde_json", "serde_json",
"sha-1",
"soketto", "soketto",
"thiserror", "thiserror",
"tokio", "tokio",
@@ -675,9 +679,9 @@ checksum = "6456b8a6c8f33fee7d958fcd1b60d55b11940a79e63ae87013e6d22e26034440"
[[package]] [[package]]
name = "hyper" name = "hyper"
version = "0.14.9" version = "0.14.11"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "07d6baa1b441335f3ce5098ac421fb6547c46dda735ca1bc6d0153c838f9dd83" checksum = "0b61cf2d1aebcf6e6352c97b81dc2244ca29194be1b276f5d8ad5c6330fffb11"
dependencies = [ dependencies = [
"bytes", "bytes",
"futures-channel", "futures-channel",
@@ -1689,6 +1693,7 @@ dependencies = [
"futures", "futures",
"hex", "hex",
"http", "http",
"hyper",
"log", "log",
"once_cell", "once_cell",
"parking_lot", "parking_lot",
backend/common/Cargo.toml
+4
View File
@@ -6,12 +6,15 @@ edition = "2018"
license = "GPL-3.0" license = "GPL-3.0"
[dependencies] [dependencies]
anyhow = "1.0.42"
base64 = { default-features = false, features = ["alloc"], version = "0.13" }
bimap = "0.6.1" bimap = "0.6.1"
bytes = "1.0.1" bytes = "1.0.1"
fnv = "1.0.7" fnv = "1.0.7"
futures = "0.3.15" futures = "0.3.15"
hex = "0.4.3" hex = "0.4.3"
http = "0.2.4" http = "0.2.4"
hyper = { version = "0.14.11", features = ["full"] }
log = "0.4" log = "0.4"
num-traits = "0.2" num-traits = "0.2"
pin-project-lite = "0.2.7" pin-project-lite = "0.2.7"
@@ -19,6 +22,7 @@ primitive-types = { version = "0.9.0", features = ["serde"] }
rustc-hash = "1.1.0" rustc-hash = "1.1.0"
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_json = { version = "1.0", features = ["raw_value"] } serde_json = { version = "1.0", features = ["raw_value"] }
sha-1 = { default-features = false, version = "0.9" }
soketto = "0.6.0" soketto = "0.6.0"
thiserror = "1.0.24" thiserror = "1.0.24"
tokio = { version = "1.8.2", features = ["full"] } tokio = { version = "1.8.2", features = ["full"] }
backend/common/src/http_utils.rs
+139
View File
@@ -0,0 +1,139 @@
use std::net::SocketAddr;
use hyper::{ Server, Request, Response, Body };
use std::future::Future;
use tokio_util::compat::{Compat,TokioAsyncReadCompatExt};
use futures::io::{BufReader, BufWriter};
use hyper::server::conn::AddrStream;
/// A convenience function to start up a Hyper server and handle requests.
pub async fn start_server<H, F>(addr: SocketAddr, handler: H) -> Result<(), anyhow::Error>
where
H: Clone + Send + Sync + 'static + FnMut(SocketAddr, Request<Body>) -> F,
F: Send + 'static + Future<Output = Result<Response<Body>, anyhow::Error>>
{
let service =
hyper::service::make_service_fn(move |addr: &AddrStream| {
let mut handler = handler.clone();
let addr = addr.remote_addr();
async move { Ok::<_, hyper::Error>(hyper::service::service_fn(move |r| handler(addr, r))) }
});
let server = Server::bind(&addr).serve(service);
log::info!("listening on http://{}", server.local_addr());
server.await?;
Ok(())
}
type WsStream = BufReader<BufWriter<Compat<hyper::upgrade::Upgraded>>>;
pub type WsSender = soketto::connection::Sender<WsStream>;
pub type WsReceiver = soketto::connection::Receiver<WsStream>;
/// A convenience function to upgrade a Hyper request into a Soketto Websocket.
pub fn upgrade_to_websocket<H,F>(req: Request<Body>, on_upgrade: H) -> hyper::Response<Body>
where
H: 'static + Send + FnOnce(WsSender, WsReceiver) -> F,
F: Send + Future<Output = ()>
{
if !is_upgrade_request(&req) {
return basic_response(400, "Expecting WebSocket upgrade headers");
}
let key = match req.headers().get("Sec-WebSocket-Key") {
Some(key) => key,
None => return basic_response(400, "Upgrade to websocket connection failed; Sec-WebSocket-Key header not provided")
};
if req.headers().get("Sec-WebSocket-Version").map(|v| v.as_bytes()) != Some(b"13") {
return basic_response(400, "Sec-WebSocket-Version header should have a value of 13");
}
// Just a little ceremony we need to go to to return the correct response key:
let mut accept_key_buf = [0; 32];
let accept_key = generate_websocket_accept_key(key.as_bytes(), &mut accept_key_buf);
// Tell the client that we accept the upgrade-to-WS request:
let response = Response::builder()
.status(hyper::StatusCode::SWITCHING_PROTOCOLS)
.header(hyper::header::CONNECTION, "upgrade")
.header(hyper::header::UPGRADE, "websocket")
.header("Sec-WebSocket-Accept", accept_key)
.body(Body::empty())
.expect("bug: failed to build response");
// Spawn our handler to work with the WS connection:
tokio::spawn(async move {
// Get our underlying TCP stream:
let stream = match hyper::upgrade::on(req).await {
Ok(stream) => stream,
Err(e) => {
log::error!("Error upgrading connection to websocket: {}", e);
return;
}
};
// Start a Soketto server with it:
let server = soketto::handshake::Server::new(
BufReader::new(BufWriter::new(stream.compat()))
);
// Get hold of a way to send and receive messages:
let (sender, receiver)
= server.into_builder().finish();
// Pass these to our when-upgraded handler:
on_upgrade(sender, receiver).await;
});
response
}
/// A helper to return a basic HTTP response with a code and text body.
fn basic_response(code: u16, msg: impl AsRef<str>) -> Response<Body> {
Response::builder()
.status(code)
.body(Body::from(msg.as_ref().to_owned()))
.expect("bug: failed to build response body")
}
/// Defined in RFC 6455. this is how we convert the Sec-WebSocket-Key in a request into a
/// Sec-WebSocket-Accept that we return in the response.
fn generate_websocket_accept_key<'a>(key: &[u8], buf: &'a mut [u8; 32]) -> &'a [u8] {
// Defined in RFC 6455, we append this to the key to generate the response:
const KEY: &[u8] = b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
use sha1::{Digest, Sha1};
let mut digest = Sha1::new();
digest.update(key);
digest.update(KEY);
let d = digest.finalize();
let n = base64::encode_config_slice(&d, base64::STANDARD, buf);
&buf[..n]
}
/// Check if a request is a websocket upgrade request.
fn is_upgrade_request<B>(request: &hyper::Request<B>) -> bool {
header_contains_value(request.headers(), hyper::header::CONNECTION, b"upgrade")
&& header_contains_value(request.headers(), hyper::header::UPGRADE, b"websocket")
}
/// Check if there is a header of the given name containing the wanted value.
fn header_contains_value(headers: &hyper::HeaderMap, header: hyper::header::HeaderName, value: &[u8]) -> bool {
pub fn trim(x: &[u8]) -> &[u8] {
let from = match x.iter().position(|x| !x.is_ascii_whitespace()) {
Some(i) => i,
None => return &x[0..0],
};
let to = x.iter().rposition(|x| !x.is_ascii_whitespace()).unwrap();
&x[from..=to]
}
for header in headers.get_all(header) {
if header.as_bytes().split(|&c| c == b',').any(|x| trim(x).eq_ignore_ascii_case(value)) {
return true;
}
}
false
}
backend/common/src/lib.rs
+1
View File
@@ -5,6 +5,7 @@ pub mod node_types;
pub mod time; pub mod time;
pub mod ws_client; pub mod ws_client;
pub mod ready_chunks_all; pub mod ready_chunks_all;
pub mod http_utils;
mod assign_id; mod assign_id;
mod dense_map; mod dense_map;
backend/telemetry_core/Cargo.toml
+1
View File
@@ -14,6 +14,7 @@ common = { path = "../common" }
futures = "0.3.15" futures = "0.3.15"
hex = "0.4.3" hex = "0.4.3"
http = "0.2.4" http = "0.2.4"
hyper = "0.14.11"
log = "0.4.14" log = "0.4.14"
once_cell = "1.8.0" once_cell = "1.8.0"
parking_lot = "0.11.1" parking_lot = "0.11.1"
backend/telemetry_core/src/aggregator/aggregator.rs
+2 -2
View File
@@ -74,7 +74,7 @@ impl Aggregator {
/// Return a sink that a shard can send messages into to be handled by the aggregator. /// Return a sink that a shard can send messages into to be handled by the aggregator.
pub fn subscribe_shard( pub fn subscribe_shard(
&self, &self,
) -> impl Sink<inner_loop::FromShardWebsocket, Error = anyhow::Error> + Unpin { ) -> impl Sink<inner_loop::FromShardWebsocket, Error = anyhow::Error> + Send + Sync + Unpin + 'static {
// Assign a unique aggregator-local ID to each connection that subscribes, and pass // Assign a unique aggregator-local ID to each connection that subscribes, and pass
// that along with every message to the aggregator loop: // that along with every message to the aggregator loop:
let shard_conn_id = self let shard_conn_id = self
@@ -96,7 +96,7 @@ impl Aggregator {
/// Return a sink that a feed can send messages into to be handled by the aggregator. /// Return a sink that a feed can send messages into to be handled by the aggregator.
pub fn subscribe_feed( pub fn subscribe_feed(
&self, &self,
) -> impl Sink<inner_loop::FromFeedWebsocket, Error = anyhow::Error> + Unpin { ) -> impl Sink<inner_loop::FromFeedWebsocket, Error = anyhow::Error> + Send + Sync + Unpin + 'static {
// Assign a unique aggregator-local ID to each connection that subscribes, and pass // Assign a unique aggregator-local ID to each connection that subscribes, and pass
// that along with every message to the aggregator loop: // that along with every message to the aggregator loop:
let feed_conn_id = self let feed_conn_id = self
backend/telemetry_core/src/main.rs
+233 -218
View File
@@ -2,8 +2,6 @@ mod aggregator;
mod feed_message; mod feed_message;
mod find_location; mod find_location;
mod state; mod state;
use std::net::SocketAddr;
use std::str::FromStr; use std::str::FromStr;
use aggregator::{ use aggregator::{
@@ -15,8 +13,8 @@ use common::ready_chunks_all::ReadyChunksAll;
use futures::{channel::mpsc, SinkExt, StreamExt}; use futures::{channel::mpsc, SinkExt, StreamExt};
use simple_logger::SimpleLogger; use simple_logger::SimpleLogger;
use structopt::StructOpt; use structopt::StructOpt;
use warp::filters::ws; use hyper::{ Response, Method };
use warp::Filter; use common::http_utils;
const VERSION: &str = env!("CARGO_PKG_VERSION"); const VERSION: &str = env!("CARGO_PKG_VERSION");
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS"); const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
@@ -60,65 +58,63 @@ async fn main() {
/// Declare our routes and start the server. /// Declare our routes and start the server.
async fn start_server(opts: Opts) -> anyhow::Result<()> { async fn start_server(opts: Opts) -> anyhow::Result<()> {
let shard_aggregator = Aggregator::spawn(opts.denylist).await?; let aggregator = Aggregator::spawn(opts.denylist).await?;
let feed_aggregator = shard_aggregator.clone(); let server = http_utils::start_server(opts.socket, move |addr, req| {
let aggregator = aggregator.clone();
println!("REQUEST: {:?}", (req.method(), req.uri().path()));
async move {
match (req.method(), req.uri().path().trim_end_matches('/')) {
// Check that the server is up and running:
(&Method::GET, "/health") => {
Ok(Response::new("OK".into()))
},
// Subscribe to feed messages:
(&Method::GET, "/feed") => {
Ok(http_utils::upgrade_to_websocket(req, move |ws_send, ws_recv| async move {
let tx_to_aggregator = aggregator.subscribe_feed();
let (mut tx_to_aggregator, mut ws_send)
= handle_feed_websocket_connection(ws_send, ws_recv, 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 _ = ws_send.close().await;
}))
},
// Subscribe to shard messages:
(&Method::GET, "/shard_submit") => {
Ok(http_utils::upgrade_to_websocket(req, move |ws_send, ws_recv| async move {
let tx_to_aggregator = aggregator.subscribe_shard();
let (mut tx_to_aggregator, mut ws_send)
= handle_shard_websocket_connection(ws_send, ws_recv, 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 _ = ws_send.close().await;
}))
},
// 404 for anything else:
_ => {
Ok(Response::builder()
.status(404)
.body("Not found".into())
.unwrap())
}
}
}
});
// Handle requests to /health by returning OK. server.await?;
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(()) Ok(())
} }
/// This handles messages coming to/from a shard connection /// This handles messages coming to/from a shard connection
async fn handle_shard_websocket_connection<S>( async fn handle_shard_websocket_connection<S>(
mut websocket: ws::WebSocket, mut ws_send: http_utils::WsSender,
mut ws_recv: http_utils::WsReceiver,
mut tx_to_aggregator: S, mut tx_to_aggregator: S,
) -> (S, ws::WebSocket) ) -> (S, http_utils::WsSender)
where where
S: futures::Sink<FromShardWebsocket, Error = anyhow::Error> + Unpin, 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, mut rx_from_aggregator) = mpsc::unbounded();
@@ -128,102 +124,119 @@ where
}; };
if let Err(e) = tx_to_aggregator.send(init_msg).await { if let Err(e) = tx_to_aggregator.send(init_msg).await {
log::error!("Error sending message to aggregator: {}", e); log::error!("Error sending message to aggregator: {}", e);
return (tx_to_aggregator, websocket); return (tx_to_aggregator, ws_send);
} }
// Loop, handling new messages from the shard or from the aggregator: // Channels to notify each loop if the other closes:
loop { let (recv_closer_tx, mut recv_closer_rx) = tokio::sync::oneshot::channel::<()>();
tokio::select! { let (send_closer_tx, mut send_closer_rx) = tokio::sync::oneshot::channel::<()>();
// 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 { // Receive messages from a shard:
ToShardWebsocket::Mute { local_id, reason } => { let recv_handle = tokio::spawn(async move {
internal_messages::FromTelemetryCore::Mute { local_id, reason } loop {
} let mut bytes = Vec::new();
};
let bytes = bincode::options() // Receive a message, or bail if closer called. We don't care about cancel safety;
.serialize(&internal_msg) // if we're halfway through receiving a message, no biggie since we're closing the
.expect("message to shard should serialize"); // connection anyway.
let msg_info = tokio::select! {
msg_info = ws_recv.receive_data(&mut bytes) => msg_info,
_ = &mut recv_closer_rx => { break }
};
if let Err(e) = websocket.send(ws::Message::binary(bytes)).await { // Handle the socket closing, or errors receiving the message.
log::error!("Error sending message to shard; booting it: {}", e); if let Err(soketto::connection::Error::Closed) = msg_info {
break break;
} }
if let Err(e) = msg_info {
log::error!("Shutting down websocket connection: Failed to receive data: {}", 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 { let msg: internal_messages::FromShardAggregator = match bincode::options().deserialize(&bytes) {
Err(e) => { Ok(msg) => msg,
log::error!("Error receiving message from shard; booting it: {}", e); Err(e) => {
break; log::error!("Failed to deserialize message from shard; booting it: {}", e);
},
Ok(msg) => msg
};
// Close message? Break and allow connection to be dropped.
if msg.is_close() {
break; break;
} }
};
// If the message isn't something we want to handle, just ignore it. // Convert and send to the aggregator:
// This includes system messages like "pings" and such, so don't log anything. let aggregator_msg = match msg {
if !msg.is_binary() && !msg.is_text() { internal_messages::FromShardAggregator::AddNode { ip, node, local_id, genesis_hash } => {
continue; 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 }
},
};
let bytes = msg.as_bytes(); if let Err(e) = tx_to_aggregator.send(aggregator_msg).await {
let msg: internal_messages::FromShardAggregator = match bincode::options().deserialize(bytes) { log::error!("Failed to send message to aggregator; closing shard: {}", e);
Ok(msg) => msg, break;
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;
}
} }
} }
}
drop(send_closer_tx); // Kill the send task if this recv task ends
tx_to_aggregator
});
// Send messages to the shard:
let send_handle = tokio::spawn(async move {
loop {
let msg = tokio::select! {
msg = rx_from_aggregator.next() => msg,
_ = &mut send_closer_rx => { break }
};
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) = ws_send.send_binary(bytes).await {
log::error!("Failed to send message to aggregator; closing shard: {}", e)
}
if let Err(e) = ws_send.flush().await {
log::error!("Failed to flush message to aggregator; closing shard: {}", e)
}
}
drop(recv_closer_tx); // Kill the recv task if this send task ends
ws_send
});
// If our send/recv tasks are stopped (if one of them dies, they both will),
// collect the bits we need to hand back from them:
let ws_send = send_handle.await.unwrap();
let tx_to_aggregator = recv_handle.await.unwrap();
// loop ended; give socket back to parent: // loop ended; give socket back to parent:
(tx_to_aggregator, websocket) (tx_to_aggregator, ws_send)
} }
/// This handles messages coming from a feed connection /// This handles messages coming from a feed connection
async fn handle_feed_websocket_connection<S>( async fn handle_feed_websocket_connection<S>(
mut websocket: ws::WebSocket, mut ws_send: http_utils::WsSender,
mut ws_recv: http_utils::WsReceiver,
mut tx_to_aggregator: S, mut tx_to_aggregator: S,
) -> (S, ws::WebSocket) ) -> (S, http_utils::WsSender)
where where
S: futures::Sink<FromFeedWebsocket, Error = anyhow::Error> + Unpin, S: futures::Sink<FromFeedWebsocket, Error = anyhow::Error> + Unpin + Send + 'static,
{ {
// unbounded channel so that slow feeds don't block aggregator progress: // unbounded channel so that slow feeds don't block aggregator progress:
let (tx_to_feed_conn, rx_from_aggregator) = mpsc::unbounded(); let (tx_to_feed_conn, rx_from_aggregator) = mpsc::unbounded();
@@ -235,106 +248,108 @@ where
}; };
if let Err(e) = tx_to_aggregator.send(init_msg).await { if let Err(e) = tx_to_aggregator.send(init_msg).await {
log::error!("Error sending message to aggregator: {}", e); log::error!("Error sending message to aggregator: {}", e);
return (tx_to_aggregator, websocket); return (tx_to_aggregator, ws_send);
} }
// Loop, handling new messages from the shard or from the aggregator: // Channels to notify each loop if the other closes:
loop { let (recv_closer_tx, mut recv_closer_rx) = tokio::sync::oneshot::channel::<()>();
// Without any special handling, if messages come in every ~2.5ms to each feed, the select! loop let (send_closer_tx, mut send_closer_rx) = tokio::sync::oneshot::channel::<()>();
// has to wake up 400 times a second to poll things. If we have 1000 feeds, that's 400,000 wakeups
// per second. Even without any work in the loop, that uses a bunch of CPU. As an example, try
// replacing the loop with this:
//
// ```
// let s = tokio::time::sleep(tokio::time::Duration::from_micros(2500));
// tokio::select! {
// _ = s => {},
// _ = websocket.next() => {}
// }
// continue;
// ```
//
// To combat this, we add a small wait to reduce how often the select loop will be woken up under high load. We
// buffer messages to feeds so that we do as much work as possible during each wakeup, and if the
// wakeup lasts longer than 75ms we don't wait before polling again. This knocks ~80% of a CPU worth of usage
// off on my machine running a soak test with 500 feeds, 4 shards and 100 nodes, doesn't seem to impact
// memory usage much, and still ensures that messages are delivered in a timely fashion.
//
// Increasing the wait to 100ms or more doesn't seem to have much more of a positive impact anyway.
let debounce = tokio::time::sleep_until(tokio::time::Instant::now() + std::time::Duration::from_millis(75));
tokio::select! {biased; // Receive messages from the feed:
let recv_handle = tokio::spawn(async move {
loop {
let mut bytes = Vec::new();
// FRONTEND -> AGGREGATOR (relay messages to the aggregator). Biased, so messages // Receive a message, or bail if closer called. We don't care about cancel safety;
// from the UI will have priority (especially important with our debounce delay). // if we're halfway through receiving a message, no biggie since we're closing the
msg = websocket.next() => { // connection anyway.
// End the loop when connection from feed ends: let msg_info = tokio::select! {
let msg = match msg { msg_info = ws_recv.receive_data(&mut bytes) => msg_info,
Some(msg) => msg, _ = &mut recv_closer_rx => { break }
None => break };
};
// If we see any errors, log them and end our loop: // Handle the socket closing, or errors receiving the message.
let msg = match msg { if let Err(soketto::connection::Error::Closed) = msg_info {
Err(e) => { break;
log::error!("Error in node websocket connection: {}", e); }
break; if let Err(e) = msg_info {
}, log::error!("Shutting down websocket connection: Failed to receive data: {}", e);
Ok(msg) => msg break;
};
// 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;
}
} }
// AGGREGATOR -> FRONTEND (buffer messages to the UI) // We ignore all but valid UTF8 text messages from the frontend:
msgs = rx_from_aggregator_chunks.next() => { let text = match String::from_utf8(bytes) {
// End the loop when connection from aggregator ends: Ok(s) => s,
let msgs = match msgs { Err(_) => continue
Some(msgs) => msgs, };
None => break
};
// There is only one message type at the mo; bytes to send // Parse the message into a command we understand and send it to the aggregator:
// to the websocket. collect them all up to dispatch in one shot. let cmd = match FromFeedWebsocket::from_str(&text) {
let all_ws_msgs = msgs.into_iter().map(|msg| { Ok(cmd) => cmd,
let bytes = match msg { Err(e) => {
ToFeedWebsocket::Bytes(bytes) => bytes log::warn!("Ignoring invalid command '{}' from the frontend: {}", text, e);
}; continue
Ok(ws::Message::binary(&*bytes))
});
if let Err(e) = websocket.send_all(&mut futures::stream::iter(all_ws_msgs)).await {
log::warn!("Closing feed websocket due to error: {}", e);
break;
} }
};
if let Err(e) = tx_to_aggregator.send(cmd).await {
log::error!("Failed to send message to aggregator; closing feed: {}", e);
break;
} }
} }
debounce.await; drop(send_closer_tx); // Kill the send task if this recv task ends
} tx_to_aggregator
});
// Send messages to the feed:
let send_handle = tokio::spawn(async move {
loop {
let debounce = tokio::time::sleep_until(tokio::time::Instant::now() + std::time::Duration::from_millis(75));
let msgs = tokio::select! {
msgs = rx_from_aggregator_chunks.next() => msgs,
_ = &mut send_closer_rx => { break }
};
// End the loop when connection from aggregator ends:
let msgs = match msgs {
Some(msgs) => msgs,
None => break
};
// There is only one message type at the mo; bytes to send
// to the websocket. collect them all up to dispatch in one shot.
let all_msg_bytes = msgs.into_iter().map(|msg| {
match msg {
ToFeedWebsocket::Bytes(bytes) => bytes
}
});
for bytes in all_msg_bytes {
if let Err(e) = ws_send.send_binary(&bytes).await {
log::warn!("Closing feed websocket due to error sending data: {}", e);
break;
}
}
if let Err(e) = ws_send.flush().await {
log::warn!("Closing feed websocket due to error flushing data: {}", e);
break;
}
debounce.await;
}
drop(recv_closer_tx); // Kill the recv task if this send task ends
ws_send
});
// If our send/recv tasks are stopped (if one of them dies, they both will),
// collect the bits we need to hand back from them:
let ws_send = send_handle.await.unwrap();
let tx_to_aggregator = recv_handle.await.unwrap();
// loop ended; give socket back to parent: // loop ended; give socket back to parent:
(tx_to_aggregator, websocket) (tx_to_aggregator, ws_send)
} }