pub mod iroh;
#[cfg(feature = "sim")]
pub mod sim;
#[cfg(feature = "sim")]
pub mod sim_cluster;
use std::collections::HashMap;
use std::net::SocketAddr;
use std::ops::ControlFlow;
use async_trait::async_trait;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use super::config::{NodeClass, NodeIdentity};
use super::error::ClusterError;
use super::framing::{self, ControlMessage, FrameCodec};
#[derive(Clone, PartialEq, Eq, Hash, Debug, serde::Serialize, serde::Deserialize)]
pub struct NodeId(pub String);
impl std::fmt::Display for NodeId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{}", self.0)
}
}
#[derive(Clone, Debug)]
pub struct PeerAddr {
pub id: NodeId,
pub hint: Vec<SocketAddr>,
}
#[derive(Clone, Debug)]
pub enum ConnectionEvent {
Connected(String),
Disconnected(String),
}
pub struct IncomingConnection {
pub remote_identity: NodeIdentity,
pub connection: Box<dyn Connection>,
pub control_tx: mpsc::UnboundedSender<ControlMessage>,
pub control_recv: Box<dyn RecvHalf>,
pub node_class: NodeClass,
pub node_metadata: HashMap<String, String>,
pub is_edge_client: bool,
}
#[async_trait]
pub trait Net: Send + Sync + 'static {
fn node_id(&self) -> NodeId;
fn local_addr(&self) -> SocketAddr;
async fn connect(&self, addr: PeerAddr) -> Result<IncomingConnection, ClusterError>;
async fn send_control(&self, node_id: &str, msg: ControlMessage) -> Result<(), ClusterError>;
async fn broadcast_control(&self, msg: &ControlMessage);
async fn open_actor_stream(
&self,
node_id: &str,
) -> Result<(Box<dyn SendHalf>, Box<dyn RecvHalf>), ClusterError>;
async fn connected_nodes(&self) -> Vec<String>;
async fn remove_connection(&self, node_id: &str);
}
#[async_trait]
pub trait Connection: Send + Sync {
fn remote_id(&self) -> NodeId;
async fn open_bi(&self) -> Result<(Box<dyn SendHalf>, Box<dyn RecvHalf>), ClusterError>;
async fn accept_bi(&self) -> Option<(Box<dyn SendHalf>, Box<dyn RecvHalf>)>;
fn close(&self, code: u32, reason: &[u8]);
}
#[async_trait]
pub trait SendHalf: Send {
async fn write_all(&mut self, buf: &[u8]) -> Result<(), ClusterError>;
fn finish(&mut self) -> Result<(), ClusterError>;
}
#[async_trait]
pub trait RecvHalf: Send {
async fn read(&mut self, buf: &mut [u8]) -> Result<Option<usize>, ClusterError>;
}
pub async fn run_control_stream_writer(
mut send: Box<dyn SendHalf>,
mut rx: mpsc::UnboundedReceiver<ControlMessage>,
shutdown: CancellationToken,
) {
loop {
tokio::select! {
biased;
_ = shutdown.cancelled() => break,
msg = rx.recv() => {
let Some(msg) = msg else { break };
let frame = match framing::encode_message(&msg) {
Ok(f) => f,
Err(e) => {
tracing::error!("Failed to encode control message: {e}");
continue;
}
};
if let Err(e) = send.write_all(&frame).await {
tracing::warn!("Control stream write failed: {e}");
break;
}
}
}
}
let _ = send.finish();
}
pub async fn pump_frames(
recv: &mut dyn RecvHalf,
codec: &mut FrameCodec,
mut on_read: impl FnMut(usize),
mut on_frame: impl FnMut(Vec<u8>) -> ControlFlow<()>,
) -> Result<(), ClusterError> {
let mut buf = vec![0u8; framing::FRAME_READ_BUF];
loop {
match recv.read(&mut buf).await? {
Some(n) => {
on_read(n);
codec.push_data(&buf[..n]);
while let Ok(Some(frame)) = codec.next_frame() {
if on_frame(frame).is_break() {
return Ok(());
}
}
}
None => return Ok(()),
}
}
}
pub async fn run_control_stream_reader(
mut recv: Box<dyn RecvHalf>,
tx: mpsc::UnboundedSender<(String, ControlMessage)>,
node_id: String,
shutdown: CancellationToken,
) {
let mut codec = FrameCodec::new();
let pump = pump_frames(
recv.as_mut(),
&mut codec,
|_| {},
|frame| match framing::decode_message::<ControlMessage>(&frame) {
Ok(msg) => {
if tx.send((node_id.clone(), msg)).is_ok() {
ControlFlow::Continue(())
} else {
ControlFlow::Break(()) }
}
Err(e) => {
tracing::warn!("Failed to decode control message from {node_id}: {e}");
ControlFlow::Continue(())
}
},
);
tokio::select! {
biased;
_ = shutdown.cancelled() => {}
result = pump => match result {
Ok(()) => tracing::debug!("Control stream from {node_id} closed"),
Err(e) => tracing::warn!("Control stream read error from {node_id}: {e}"),
},
}
}