mirror of
https://github.com/pezkuwichain/pezkuwi-telemetry.git
synced 2026-07-11 19:55:45 +00:00
Squashed diff from mh-backend-shard
This commit is contained in:
Maciej Hirsz
@@ -0,0 +1,239 @@
|
||||
use std::net::Ipv4Addr;
|
||||
use std::fmt;
|
||||
// use std::sync::mpsc::{self, Sender};
|
||||
|
||||
use actix::prelude::*;
|
||||
use actix_http::http::Uri;
|
||||
use bincode::Options;
|
||||
use rustc_hash::FxHashMap;
|
||||
use shared::util::{Hash, DenseMap};
|
||||
use shared::types::{ConnId, NodeDetails, NodeId};
|
||||
use shared::node::Payload;
|
||||
use shared::shard::{ShardConnId, ShardMessage, BackendMessage};
|
||||
use soketto::handshake::{Client, ServerResponse};
|
||||
use crate::node::{NodeConnector, Initialize};
|
||||
use tokio::net::TcpStream;
|
||||
use tokio::sync::mpsc::{self, UnboundedSender};
|
||||
use tokio_util::compat::{Compat, TokioAsyncReadCompatExt};
|
||||
|
||||
type WsSender = soketto::Sender<Compat<TcpStream>>;
|
||||
type WsReceiver = soketto::Receiver<Compat<TcpStream>>;
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct Aggregator {
|
||||
url: Uri,
|
||||
chains: FxHashMap<Hash, UnboundedSender<ChainMessage>>,
|
||||
}
|
||||
|
||||
impl Actor for Aggregator {
|
||||
type Context = Context<Self>;
|
||||
}
|
||||
|
||||
impl Aggregator {
|
||||
pub fn new(url: Uri) -> Self {
|
||||
Aggregator {
|
||||
url,
|
||||
chains: Default::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct Chain {
|
||||
/// Base URL of Backend Core
|
||||
url: Uri,
|
||||
/// Genesis hash of the chain, required to construct the URL to connect to the Backend Core
|
||||
genesis_hash: Hash,
|
||||
/// Dense mapping of SharedConnId -> Addr<NodeConnector> + multiplexing ConnId sent from the node.
|
||||
nodes: DenseMap<(Addr<NodeConnector>, ConnId)>,
|
||||
}
|
||||
|
||||
impl Chain {
|
||||
pub fn new(url: Uri, genesis_hash: Hash) -> Self {
|
||||
Chain {
|
||||
url,
|
||||
genesis_hash,
|
||||
nodes: DenseMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn spawn(mut self) -> UnboundedSender<ChainMessage> {
|
||||
let (tx_ret, mut rx) = mpsc::unbounded_channel();
|
||||
|
||||
let tx = tx_ret.clone();
|
||||
|
||||
tokio::task::spawn(async move {
|
||||
let mut sender = match self.connect(tx.clone()).await {
|
||||
Ok(pair) => pair,
|
||||
Err(err) => {
|
||||
log::error!("Failed to connect to Backend Core: {:?}", err);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// tokio::task::spawn(async move {
|
||||
|
||||
// });
|
||||
|
||||
loop {
|
||||
match rx.recv().await {
|
||||
Some(ChainMessage::AddNode(msg)) => {
|
||||
println!("Add node {:?}", msg);
|
||||
|
||||
let AddNode { node, ip, conn_id, node_connector, .. } = msg;
|
||||
let sid = self.nodes.add((node_connector, conn_id)) as ShardConnId;
|
||||
|
||||
let bytes = bincode::options().serialize(&ShardMessage::AddNode {
|
||||
ip,
|
||||
node,
|
||||
sid,
|
||||
}).unwrap();
|
||||
|
||||
println!("Sending {} bytes", bytes.len());
|
||||
|
||||
let _ = sender.send_binary_mut(bytes).await;
|
||||
let _ = sender.flush().await;
|
||||
},
|
||||
Some(ChainMessage::UpdateNode(nid, payload)) => {
|
||||
let msg = ShardMessage::UpdateNode {
|
||||
nid,
|
||||
payload,
|
||||
};
|
||||
|
||||
println!("Serialize {:?}", msg);
|
||||
|
||||
let bytes = bincode::options().serialize(&msg).unwrap();
|
||||
|
||||
println!("Sending update: {} bytes", bytes.len());
|
||||
|
||||
let _ = sender.send_binary_mut(bytes).await;
|
||||
let _ = sender.flush().await;
|
||||
},
|
||||
Some(ChainMessage::Backend(BackendMessage::Initialize { sid, nid })) => {
|
||||
if let Some((addr, conn_id)) = self.nodes.get(sid as usize) {
|
||||
addr.do_send(Initialize {
|
||||
nid,
|
||||
conn_id: *conn_id,
|
||||
chain: tx.clone(),
|
||||
})
|
||||
}
|
||||
},
|
||||
Some(ChainMessage::Backend(BackendMessage::Mute { sid, reason })) => {
|
||||
// TODO
|
||||
},
|
||||
None => (),
|
||||
}
|
||||
}
|
||||
// let mut client = Client::new(socket.compat(), host, &path);
|
||||
|
||||
// let (mut sender, mut receiver) = match client.handshake().await? {
|
||||
// ServerResponse::Accepted { .. } => client.into_builder().finish(),
|
||||
// ServerResponse::Redirect { status_code, location } => unimplemented!("follow location URL"),
|
||||
// ServerResponse::Rejected { status_code } => unimplemented!("handle failure")
|
||||
// };
|
||||
});
|
||||
|
||||
tx_ret
|
||||
}
|
||||
|
||||
pub async fn connect(&self, tx: UnboundedSender<ChainMessage>) -> anyhow::Result<WsSender> {
|
||||
let host = self.url.host().unwrap_or("127.0.0.1");
|
||||
let port = self.url.port_u16().unwrap_or(8000);
|
||||
let path = format!("{}{}", self.url.path(), self.genesis_hash);
|
||||
|
||||
let socket = TcpStream::connect((host, port)).await?;
|
||||
|
||||
socket.set_nodelay(true).unwrap();
|
||||
|
||||
let mut client = Client::new(socket.compat(), host, &path);
|
||||
|
||||
let (sender, receiver) = match client.handshake().await? {
|
||||
ServerResponse::Accepted { .. } => client.into_builder().finish(),
|
||||
ServerResponse::Redirect { status_code, .. } |
|
||||
ServerResponse::Rejected { status_code } => {
|
||||
return Err(anyhow::anyhow!("Failed to connect, status code: {}", status_code));
|
||||
}
|
||||
};
|
||||
|
||||
async fn read(tx: UnboundedSender<ChainMessage>, mut receiver: WsReceiver) -> anyhow::Result<()> {
|
||||
let mut data = Vec::with_capacity(128);
|
||||
|
||||
loop {
|
||||
data.clear();
|
||||
|
||||
receiver.receive_data(&mut data).await?;
|
||||
|
||||
println!("Received {} bytes from Backend Core", data.len());
|
||||
|
||||
match bincode::options().deserialize(&data) {
|
||||
Ok(msg) => tx.send(ChainMessage::Backend(msg))?,
|
||||
Err(err) => {
|
||||
log::error!("Failed to read message from Backend Core: {:?}", err);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
tokio::task::spawn(read(tx, receiver));
|
||||
|
||||
Ok(sender)
|
||||
}
|
||||
}
|
||||
|
||||
impl Actor for Chain {
|
||||
type Context = Context<Self>;
|
||||
}
|
||||
|
||||
#[derive(Message)]
|
||||
#[rtype(result = "()")]
|
||||
pub struct AddNode {
|
||||
pub ip: Option<Ipv4Addr>,
|
||||
pub genesis_hash: Hash,
|
||||
pub node: NodeDetails,
|
||||
pub conn_id: ConnId,
|
||||
pub node_connector: Addr<NodeConnector>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum ChainMessage {
|
||||
AddNode(AddNode),
|
||||
UpdateNode(NodeId, Payload),
|
||||
Backend(BackendMessage),
|
||||
}
|
||||
|
||||
impl fmt::Debug for AddNode {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
f.write_str("AddNode")
|
||||
}
|
||||
}
|
||||
|
||||
impl Handler<AddNode> for Aggregator {
|
||||
type Result = ();
|
||||
|
||||
fn handle(&mut self, msg: AddNode, ctx: &mut Self::Context) {
|
||||
let AddNode { genesis_hash, .. } = msg;
|
||||
|
||||
let url = &self.url;
|
||||
let chain = self
|
||||
.chains
|
||||
.entry(genesis_hash)
|
||||
.or_insert_with(move || Chain::new(url.clone(), genesis_hash).spawn());
|
||||
|
||||
if let Err(err) = chain.send(ChainMessage::AddNode(msg)) {
|
||||
let msg = err.0;
|
||||
log::error!("Failed to add node to chain, shutting down chain");
|
||||
self.chains.remove(&genesis_hash);
|
||||
// TODO: Send a message back to clean up node connections
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Handler<AddNode> for Chain {
|
||||
type Result = ();
|
||||
|
||||
fn handle(&mut self, msg: AddNode, ctx: &mut Self::Context) {
|
||||
let AddNode { ip, node_connector, .. } = msg;
|
||||
|
||||
println!("Node connected to {}: {:?}", self.genesis_hash, ip);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,125 @@
|
||||
use std::net::Ipv4Addr;
|
||||
|
||||
use actix::prelude::*;
|
||||
use actix_http::ws::Codec;
|
||||
use actix_http::http::Uri;
|
||||
use actix_web::{get, middleware, web, App, Error, HttpRequest, HttpResponse, HttpServer};
|
||||
use actix_web_actors::ws;
|
||||
use clap::Clap;
|
||||
use simple_logger::SimpleLogger;
|
||||
|
||||
mod aggregator;
|
||||
mod node;
|
||||
|
||||
use aggregator::Aggregator;
|
||||
use node::NodeConnector;
|
||||
|
||||
const VERSION: &str = env!("CARGO_PKG_VERSION");
|
||||
const AUTHORS: &str = env!("CARGO_PKG_AUTHORS");
|
||||
const NAME: &str = "Substrate Telemetry Backend Shard";
|
||||
const ABOUT: &str = "This is the Telemetry Backend Shard that forwards the data sent by Substrate/Polkadot nodes to the Backend Core";
|
||||
|
||||
#[derive(Clap, Debug)]
|
||||
#[clap(name = NAME, version = VERSION, author = AUTHORS, about = ABOUT)]
|
||||
struct Opts {
|
||||
#[clap(
|
||||
short = 'l',
|
||||
long = "listen",
|
||||
default_value = "127.0.0.1:8001",
|
||||
about = "This is the socket address Telemetry is 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'"
|
||||
)]
|
||||
socket: std::net::SocketAddr,
|
||||
#[clap(
|
||||
arg_enum,
|
||||
required = false,
|
||||
long = "log",
|
||||
default_value = "info",
|
||||
about = "Log level."
|
||||
)]
|
||||
log_level: LogLevel,
|
||||
#[clap(
|
||||
short = 'c',
|
||||
long = "core",
|
||||
default_value = "ws://127.0.0.1:8000/shard_submit/",
|
||||
about = "Url to the Backend Core endpoint accepting shard connections"
|
||||
)]
|
||||
core_url: Uri,
|
||||
}
|
||||
|
||||
#[derive(Clap, Debug, PartialEq)]
|
||||
enum LogLevel {
|
||||
Error,
|
||||
Warn,
|
||||
Info,
|
||||
Debug,
|
||||
Trace,
|
||||
}
|
||||
|
||||
impl From<&LogLevel> for log::LevelFilter {
|
||||
fn from(log_level: &LogLevel) -> Self {
|
||||
match log_level {
|
||||
LogLevel::Error => log::LevelFilter::Error,
|
||||
LogLevel::Warn => log::LevelFilter::Warn,
|
||||
LogLevel::Info => log::LevelFilter::Info,
|
||||
LogLevel::Debug => log::LevelFilter::Debug,
|
||||
LogLevel::Trace => log::LevelFilter::Trace,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Entry point for connecting nodes
|
||||
#[get("/submit")]
|
||||
async fn node_route(
|
||||
req: HttpRequest,
|
||||
stream: web::Payload,
|
||||
aggregator: web::Data<Addr<Aggregator>>,
|
||||
) -> Result<HttpResponse, Error> {
|
||||
let ip = req
|
||||
.connection_info()
|
||||
.realip_remote_addr()
|
||||
.and_then(|mut addr| {
|
||||
if let Some(port_idx) = addr.find(':') {
|
||||
addr = &addr[..port_idx];
|
||||
}
|
||||
addr.parse::<Ipv4Addr>().ok()
|
||||
});
|
||||
|
||||
let mut res = ws::handshake(&req)?;
|
||||
let aggregator = aggregator.get_ref().clone();
|
||||
|
||||
Ok(res.streaming(ws::WebsocketContext::with_codec(
|
||||
NodeConnector::new(aggregator, ip),
|
||||
stream,
|
||||
Codec::new().max_size(10 * 1024 * 1024), // 10mb frame limit
|
||||
)))
|
||||
}
|
||||
|
||||
/// Telemetry entry point. Listening by default on 127.0.0.1:8000.
|
||||
/// This can be changed using the `PORT` and `BIND` ENV variables.
|
||||
#[actix_web::main]
|
||||
async fn main() -> std::io::Result<()> {
|
||||
let opts = Opts::parse();
|
||||
let log_level = &opts.log_level;
|
||||
SimpleLogger::new()
|
||||
.with_level(log_level.into())
|
||||
.init()
|
||||
.expect("Must be able to start a logger");
|
||||
|
||||
println!("URL? {:?} {:?}", opts.core_url.host(), opts.core_url.port_u16());
|
||||
|
||||
let aggregator = Aggregator::new(opts.core_url).start();
|
||||
|
||||
log::info!(
|
||||
"Starting Telemetry Shard version: {}",
|
||||
env!("CARGO_PKG_VERSION")
|
||||
);
|
||||
HttpServer::new(move || {
|
||||
App::new()
|
||||
.wrap(middleware::NormalizePath::default())
|
||||
.data(aggregator.clone())
|
||||
.service(node_route)
|
||||
})
|
||||
.bind(opts.socket)?
|
||||
.run()
|
||||
.await
|
||||
}
|
||||
@@ -0,0 +1,201 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::net::Ipv4Addr;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use crate::aggregator::{AddNode, Aggregator, ChainMessage};
|
||||
// use crate::chain::{Chain, RemoveNode, UpdateNode};
|
||||
use actix::prelude::*;
|
||||
use actix_web_actors::ws::{self, CloseReason};
|
||||
use shared::node::{NodeMessage, Payload};
|
||||
use shared::types::{ConnId, NodeId};
|
||||
use shared::ws::{MultipartHandler, WsMessage};
|
||||
use tokio::sync::mpsc::UnboundedSender;
|
||||
|
||||
/// How often heartbeat pings are sent
|
||||
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(20);
|
||||
/// How long before lack of client response causes a timeout
|
||||
const CLIENT_TIMEOUT: Duration = Duration::from_secs(60);
|
||||
|
||||
pub struct NodeConnector {
|
||||
/// Multiplexing connections by id
|
||||
multiplex: BTreeMap<ConnId, ConnMultiplex>,
|
||||
/// Client must send ping at least once every 60 seconds (CLIENT_TIMEOUT),
|
||||
hb: Instant,
|
||||
/// Aggregator actor address
|
||||
aggregator: Addr<Aggregator>,
|
||||
/// IP address of the node this connector is responsible for
|
||||
ip: Option<Ipv4Addr>,
|
||||
/// Helper for handling continuation messages
|
||||
multipart: MultipartHandler,
|
||||
}
|
||||
|
||||
enum ConnMultiplex {
|
||||
Connected {
|
||||
/// Id of the node this multiplex connector is responsible for handling
|
||||
nid: NodeId,
|
||||
/// Chain address to which this multiplex connector is delegating messages
|
||||
chain: UnboundedSender<ChainMessage>,
|
||||
},
|
||||
Waiting {
|
||||
/// Backlog of messages to be sent once we get a recipient handle to the chain
|
||||
backlog: Vec<Payload>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Default for ConnMultiplex {
|
||||
fn default() -> Self {
|
||||
ConnMultiplex::Waiting {
|
||||
backlog: Vec::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Actor for NodeConnector {
|
||||
type Context = ws::WebsocketContext<Self>;
|
||||
|
||||
fn started(&mut self, ctx: &mut Self::Context) {
|
||||
self.heartbeat(ctx);
|
||||
}
|
||||
|
||||
fn stopped(&mut self, _: &mut Self::Context) {
|
||||
// for mx in self.multiplex.values() {
|
||||
// if let ConnMultiplex::Connected { chain, nid } = mx {
|
||||
// chain.do_send(RemoveNode(*nid));
|
||||
// }
|
||||
// }
|
||||
}
|
||||
}
|
||||
|
||||
impl NodeConnector {
|
||||
pub fn new(aggregator: Addr<Aggregator>, ip: Option<Ipv4Addr>) -> Self {
|
||||
Self {
|
||||
multiplex: BTreeMap::new(),
|
||||
hb: Instant::now(),
|
||||
aggregator,
|
||||
ip,
|
||||
multipart: MultipartHandler::default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn heartbeat(&self, ctx: &mut <Self as Actor>::Context) {
|
||||
ctx.run_interval(HEARTBEAT_INTERVAL, |act, ctx| {
|
||||
// check client heartbeats
|
||||
if Instant::now().duration_since(act.hb) > CLIENT_TIMEOUT {
|
||||
// stop actor
|
||||
ctx.close(Some(CloseReason {
|
||||
code: ws::CloseCode::Abnormal,
|
||||
description: Some("Missed heartbeat".into()),
|
||||
}));
|
||||
ctx.stop();
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn handle_message(
|
||||
&mut self,
|
||||
msg: NodeMessage,
|
||||
ctx: &mut <Self as Actor>::Context,
|
||||
) {
|
||||
let conn_id = msg.id();
|
||||
let payload = msg.into();
|
||||
|
||||
match self.multiplex.entry(conn_id).or_default() {
|
||||
ConnMultiplex::Connected { nid, chain } => {
|
||||
// TODO: error handle
|
||||
let _ = chain.send(ChainMessage::UpdateNode(*nid, payload));
|
||||
}
|
||||
ConnMultiplex::Waiting { backlog } => {
|
||||
if let Payload::SystemConnected(connected) = payload {
|
||||
println!("Node connected {:?}", connected.node);
|
||||
self.aggregator.do_send(AddNode {
|
||||
genesis_hash: connected.genesis_hash,
|
||||
ip: self.ip,
|
||||
node: connected.node,
|
||||
conn_id,
|
||||
node_connector: ctx.address(),
|
||||
});
|
||||
} else {
|
||||
if backlog.len() >= 10 {
|
||||
backlog.remove(0);
|
||||
}
|
||||
|
||||
backlog.push(payload);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Message)]
|
||||
#[rtype(result = "()")]
|
||||
pub struct Initialize {
|
||||
pub nid: NodeId,
|
||||
pub conn_id: ConnId,
|
||||
pub chain: UnboundedSender<ChainMessage>,
|
||||
}
|
||||
|
||||
impl Handler<Initialize> for NodeConnector {
|
||||
type Result = ();
|
||||
|
||||
fn handle(&mut self, msg: Initialize, _: &mut Self::Context) {
|
||||
let Initialize {
|
||||
nid,
|
||||
conn_id,
|
||||
chain,
|
||||
} = msg;
|
||||
log::trace!(target: "NodeConnector::Initialize", "Initializing a node, nid={}, on conn_id={}", nid, conn_id);
|
||||
let mx = self.multiplex.entry(conn_id).or_default();
|
||||
|
||||
if let ConnMultiplex::Waiting { backlog } = mx {
|
||||
for payload in backlog.drain(..) {
|
||||
// TODO: error handle.
|
||||
let _ = chain.send(ChainMessage::UpdateNode(nid, payload));
|
||||
}
|
||||
|
||||
*mx = ConnMultiplex::Connected {
|
||||
nid,
|
||||
chain,
|
||||
};
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for NodeConnector {
|
||||
fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
|
||||
self.hb = Instant::now();
|
||||
|
||||
let data = match msg.map(|msg| self.multipart.handle(msg)) {
|
||||
Ok(WsMessage::Nop) => return,
|
||||
Ok(WsMessage::Ping(msg)) => {
|
||||
ctx.pong(&msg);
|
||||
return;
|
||||
}
|
||||
Ok(WsMessage::Data(data)) => data,
|
||||
Ok(WsMessage::Close(reason)) => {
|
||||
ctx.close(reason);
|
||||
ctx.stop();
|
||||
return;
|
||||
}
|
||||
Err(error) => {
|
||||
log::error!("{:?}", error);
|
||||
ctx.stop();
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match serde_json::from_slice(&data) {
|
||||
Ok(msg) => self.handle_message(msg, ctx),
|
||||
#[cfg(debug)]
|
||||
Err(err) => {
|
||||
let data: &[u8] = data.get(..512).unwrap_or_else(|| &data);
|
||||
log::warn!(
|
||||
"Failed to parse node message: {} {}",
|
||||
err,
|
||||
std::str::from_utf8(data).unwrap_or_else(|_| "INVALID UTF8")
|
||||
)
|
||||
}
|
||||
#[cfg(not(debug))]
|
||||
Err(_) => (),
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user