use std::collections::HashMap;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::{RwLock, mpsc};
use tokio_util::sync::CancellationToken;
use crate::instrument;
use super::certs;
use super::config::{NodeClass, NodeIdentity, TransportTuning};
use super::error::ClusterError;
use super::framing::{self, ControlMessage, FrameCodec, HandshakePayload, PROTOCOL_VERSION};
#[derive(Debug, Clone)]
pub enum ConnectionEvent {
Connected(String),
Disconnected(String),
}
pub struct IncomingConnection {
pub remote_identity: NodeIdentity,
pub connection: quinn::Connection,
pub control_tx: mpsc::UnboundedSender<ControlMessage>,
pub control_recv: quinn::RecvStream,
pub node_class: NodeClass,
pub node_metadata: HashMap<String, String>,
pub is_edge_client: bool,
}
pub struct NodeConnection {
pub connection: quinn::Connection,
pub remote_identity: NodeIdentity,
pub control_tx: mpsc::UnboundedSender<ControlMessage>,
}
pub struct Transport {
endpoint: quinn::Endpoint,
identity: NodeIdentity,
cookie: String,
type_manifest: Vec<String>,
node_class: NodeClass,
node_metadata: HashMap<String, String>,
connections: Arc<RwLock<HashMap<String, NodeConnection>>>,
client_config: quinn::ClientConfig,
shutdown: CancellationToken,
connection_events_tx: mpsc::UnboundedSender<ConnectionEvent>,
is_edge_client: bool,
}
impl Transport {
pub async fn bind(
identity: NodeIdentity,
cookie: String,
type_manifest: Vec<String>,
node_class: NodeClass,
node_metadata: HashMap<String, String>,
tuning: TransportTuning,
shutdown: CancellationToken,
) -> Result<
(
Arc<Self>,
mpsc::UnboundedReceiver<IncomingConnection>,
mpsc::UnboundedReceiver<ConnectionEvent>,
),
ClusterError,
> {
let (cert_chain, private_key) = certs::generate_self_signed_cert(&identity)?;
let mut server_config = certs::create_server_config(cert_chain, private_key)?;
let mut transport_config = quinn::TransportConfig::default();
transport_config.initial_rtt(Duration::from_millis(tuning.initial_rtt_ms));
transport_config.max_idle_timeout(Some(
quinn::IdleTimeout::try_from(Duration::from_secs(tuning.max_idle_timeout_secs))
.expect("valid idle timeout"),
));
if let Some(interval) = tuning.keep_alive_interval_secs {
transport_config.keep_alive_interval(Some(Duration::from_secs(interval)));
}
transport_config.max_concurrent_bidi_streams(quinn::VarInt::from_u32(
tuning.max_concurrent_bidi_streams,
));
transport_config
.stream_receive_window(quinn::VarInt::from_u32(tuning.stream_receive_window));
transport_config.receive_window(quinn::VarInt::from_u32(tuning.receive_window));
transport_config.send_window(tuning.send_window);
transport_config.initial_mtu(tuning.initial_mtu);
let transport_config = Arc::new(transport_config);
server_config.transport_config(Arc::clone(&transport_config));
let mut client_config = certs::create_client_config()?;
client_config.transport_config(transport_config);
let endpoint = quinn::Endpoint::server(server_config, identity.socket_addr())?;
tracing::info!("Transport bound on {}", identity.socket_addr());
let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();
let (connection_events_tx, connection_events_rx) = mpsc::unbounded_channel();
let transport = Arc::new(Self {
endpoint,
identity: identity.clone(),
cookie: cookie.clone(),
type_manifest: type_manifest.clone(),
node_class,
node_metadata,
connections: Arc::new(RwLock::new(HashMap::new())),
client_config,
shutdown: shutdown.clone(),
connection_events_tx,
is_edge_client: false,
});
let transport_clone = Arc::clone(&transport);
tokio::spawn(async move {
transport_clone.accept_loop(incoming_tx).await;
});
Ok((transport, incoming_rx, connection_events_rx))
}
pub async fn connect_only(
cookie: String,
node_class: NodeClass,
node_metadata: HashMap<String, String>,
tuning: TransportTuning,
shutdown: CancellationToken,
) -> Result<(Arc<Self>, mpsc::UnboundedReceiver<ConnectionEvent>), ClusterError> {
let mut transport_config = quinn::TransportConfig::default();
transport_config.initial_rtt(Duration::from_millis(tuning.initial_rtt_ms));
transport_config.max_idle_timeout(Some(
quinn::IdleTimeout::try_from(Duration::from_secs(tuning.max_idle_timeout_secs))
.expect("valid idle timeout"),
));
if let Some(interval) = tuning.keep_alive_interval_secs {
transport_config.keep_alive_interval(Some(Duration::from_secs(interval)));
}
transport_config.max_concurrent_bidi_streams(quinn::VarInt::from_u32(
tuning.max_concurrent_bidi_streams,
));
transport_config
.stream_receive_window(quinn::VarInt::from_u32(tuning.stream_receive_window));
transport_config.receive_window(quinn::VarInt::from_u32(tuning.receive_window));
transport_config.send_window(tuning.send_window);
transport_config.initial_mtu(tuning.initial_mtu);
let transport_config = Arc::new(transport_config);
let mut client_config = certs::create_client_config()?;
client_config.transport_config(transport_config);
let endpoint = quinn::Endpoint::client("0.0.0.0:0".parse().unwrap())
.map_err(|e| ClusterError::Transport(e.to_string()))?;
let local_addr = endpoint
.local_addr()
.map_err(|e| ClusterError::Transport(e.to_string()))?;
let identity = NodeIdentity {
name: format!("edge-{}", rand::random::<u32>()),
host: local_addr.ip().to_string(),
port: local_addr.port(),
incarnation: 0,
};
let (connection_events_tx, connection_events_rx) = mpsc::unbounded_channel();
let transport = Arc::new(Self {
endpoint,
identity,
cookie,
type_manifest: Vec::new(),
node_class,
node_metadata,
connections: Arc::new(RwLock::new(HashMap::new())),
client_config,
shutdown,
connection_events_tx,
is_edge_client: true,
});
Ok((transport, connection_events_rx))
}
pub fn local_addr(&self) -> SocketAddr {
self.endpoint.local_addr().unwrap()
}
pub async fn connect(
self: &Arc<Self>,
addr: SocketAddr,
) -> Result<IncomingConnection, ClusterError> {
let mut endpoint = self.endpoint.clone();
endpoint.set_default_client_config(self.client_config.clone());
let connection = endpoint
.connect(addr, "localhost")
.map_err(|e| ClusterError::Transport(e.to_string()))?
.await?;
let (mut send, recv) = connection
.open_bi()
.await
.map_err(|e| ClusterError::Transport(e.to_string()))?;
let our_handshake = HandshakePayload {
identity: self.identity.clone(),
cookie: self.cookie.clone(),
type_manifest: self.type_manifest.clone(),
protocol_version: PROTOCOL_VERSION,
node_class: self.node_class.clone(),
node_metadata: self.node_metadata.clone(),
is_edge_client: self.is_edge_client,
};
let frame = framing::encode_message(&ControlMessage::Handshake(our_handshake))
.map_err(ClusterError::Serialization)?;
send.write_all(&frame)
.await
.map_err(|e| ClusterError::Transport(e.to_string()))?;
let (peer_handshake, recv) = read_handshake(recv).await?;
if peer_handshake.cookie != self.cookie {
connection.close(quinn::VarInt::from_u32(1), b"cookie mismatch");
return Err(ClusterError::CookieMismatch(addr));
}
if peer_handshake.protocol_version != PROTOCOL_VERSION {
connection.close(quinn::VarInt::from_u32(2), b"protocol version mismatch");
return Err(ClusterError::ProtocolMismatch {
local: PROTOCOL_VERSION,
remote: peer_handshake.protocol_version,
addr,
});
}
let remote_identity = peer_handshake.identity.clone();
let node_key = remote_identity.node_id_string();
let (control_out_tx, control_out_rx) = mpsc::unbounded_channel();
tokio::spawn(run_control_stream_writer(
send,
control_out_rx,
self.shutdown.clone(),
));
{
let mut conns = self.connections.write().await;
let stale_key = conns
.iter()
.find(|(_, nc)| nc.remote_identity.socket_addr() == addr)
.map(|(key, nc)| (key.clone(), nc.remote_identity.incarnation));
if let Some((existing_key, existing_nid)) = stale_key {
if existing_nid == remote_identity.incarnation {
connection.close(quinn::VarInt::from_u32(0), b"duplicate connection");
return Err(ClusterError::Transport(format!(
"already connected to {addr} (same incarnation)"
)));
}
if let Some(old) = conns.remove(&existing_key) {
old.connection
.close(quinn::VarInt::from_u32(0), b"stale incarnation");
tracing::info!(
"Replaced stale connection {existing_key} (nid {existing_nid}) \
with new incarnation (nid {})",
remote_identity.incarnation
);
instrument::connection_closed();
let _ = self
.connection_events_tx
.send(ConnectionEvent::Disconnected(existing_key));
}
}
conns.insert(
node_key.clone(),
NodeConnection {
connection: connection.clone(),
remote_identity: remote_identity.clone(),
control_tx: control_out_tx.clone(),
},
);
}
instrument::connection_opened();
let _ = self
.connection_events_tx
.send(ConnectionEvent::Connected(node_key));
Ok(IncomingConnection {
remote_identity,
connection,
control_tx: control_out_tx,
control_recv: recv,
node_class: peer_handshake.node_class,
node_metadata: peer_handshake.node_metadata,
is_edge_client: peer_handshake.is_edge_client,
})
}
pub async fn get_connection(&self, node_id: &str) -> Option<quinn::Connection> {
let conns = self.connections.read().await;
conns.get(node_id).map(|nc| nc.connection.clone())
}
pub async fn send_control(
&self,
node_id: &str,
msg: ControlMessage,
) -> Result<(), ClusterError> {
let conns = self.connections.read().await;
let nc = conns
.get(node_id)
.ok_or_else(|| ClusterError::NodeNotFound(node_id.to_string()))?;
nc.control_tx
.send(msg)
.map_err(|_| ClusterError::Transport("control channel closed".into()))?;
Ok(())
}
pub async fn broadcast_control(&self, msg: &ControlMessage) {
let conns = self.connections.read().await;
for (node_id, nc) in conns.iter() {
if nc.control_tx.send(msg.clone()).is_err() {
tracing::warn!("Failed to send control message to {node_id}");
}
}
}
pub async fn connected_nodes(&self) -> Vec<String> {
let conns = self.connections.read().await;
conns.keys().cloned().collect()
}
pub async fn remove_connection(&self, node_id: &str) {
let mut conns = self.connections.write().await;
if let Some(nc) = conns.remove(node_id) {
nc.connection
.close(quinn::VarInt::from_u32(0), b"node departed");
instrument::connection_closed();
let _ = self
.connection_events_tx
.send(ConnectionEvent::Disconnected(node_id.to_string()));
}
}
pub async fn open_actor_stream(
&self,
node_id: &str,
) -> Result<(quinn::SendStream, quinn::RecvStream), ClusterError> {
let conns = self.connections.read().await;
let nc = conns
.get(node_id)
.ok_or_else(|| ClusterError::NodeNotFound(node_id.to_string()))?;
let (send, recv) = nc
.connection
.open_bi()
.await
.map_err(|e| ClusterError::Transport(e.to_string()))?;
Ok((send, recv))
}
async fn accept_loop(self: Arc<Self>, incoming_tx: mpsc::UnboundedSender<IncomingConnection>) {
loop {
tokio::select! {
incoming = self.endpoint.accept() => {
let Some(incoming) = incoming else {
tracing::info!("QUIC endpoint closed");
break;
};
let transport = Arc::clone(&self);
let tx = incoming_tx.clone();
tokio::spawn(async move {
match transport.handle_incoming(incoming).await {
Ok(ic) => { let _ = tx.send(ic); }
Err(e) => tracing::warn!("Failed to accept connection: {e}"),
}
});
}
_ = self.shutdown.cancelled() => {
tracing::info!("Transport shutting down");
break;
}
}
}
}
async fn handle_incoming(
self: &Arc<Self>,
incoming: quinn::Incoming,
) -> Result<IncomingConnection, ClusterError> {
let connection = incoming.await?;
let remote_addr = connection.remote_address();
tracing::debug!("Incoming connection from {remote_addr}");
let (send, recv) = connection
.accept_bi()
.await
.map_err(|e| ClusterError::Transport(e.to_string()))?;
let (peer_handshake, recv) = read_handshake(recv).await?;
if peer_handshake.cookie != self.cookie {
connection.close(quinn::VarInt::from_u32(1), b"cookie mismatch");
return Err(ClusterError::CookieMismatch(remote_addr));
}
if peer_handshake.protocol_version != PROTOCOL_VERSION {
connection.close(quinn::VarInt::from_u32(2), b"protocol version mismatch");
return Err(ClusterError::ProtocolMismatch {
local: PROTOCOL_VERSION,
remote: peer_handshake.protocol_version,
addr: remote_addr,
});
}
let our_handshake = HandshakePayload {
identity: self.identity.clone(),
cookie: self.cookie.clone(),
type_manifest: self.type_manifest.clone(),
protocol_version: PROTOCOL_VERSION,
node_class: self.node_class.clone(),
node_metadata: self.node_metadata.clone(),
is_edge_client: self.is_edge_client,
};
let frame = framing::encode_message(&ControlMessage::Handshake(our_handshake))
.map_err(ClusterError::Serialization)?;
let mut send = send;
send.write_all(&frame)
.await
.map_err(|e| ClusterError::Transport(e.to_string()))?;
let remote_identity = peer_handshake.identity.clone();
let node_key = remote_identity.node_id_string();
let (control_out_tx, control_out_rx) = mpsc::unbounded_channel();
tokio::spawn(run_control_stream_writer(
send,
control_out_rx,
self.shutdown.clone(),
));
{
let mut conns = self.connections.write().await;
let stale_key = conns
.iter()
.find(|(_, nc)| nc.remote_identity.socket_addr() == remote_addr)
.map(|(key, nc)| (key.clone(), nc.remote_identity.incarnation));
if let Some((existing_key, existing_nid)) = stale_key {
if existing_nid == remote_identity.incarnation {
connection.close(quinn::VarInt::from_u32(0), b"duplicate connection");
return Err(ClusterError::Transport(format!(
"already connected to {remote_addr} (same incarnation)"
)));
}
if let Some(old) = conns.remove(&existing_key) {
old.connection
.close(quinn::VarInt::from_u32(0), b"stale incarnation");
tracing::info!(
"Replaced stale incoming connection {existing_key} (nid {existing_nid}) \
with new incarnation (nid {})",
remote_identity.incarnation
);
instrument::connection_closed();
let _ = self
.connection_events_tx
.send(ConnectionEvent::Disconnected(existing_key));
}
}
conns.insert(
node_key.clone(),
NodeConnection {
connection: connection.clone(),
remote_identity: remote_identity.clone(),
control_tx: control_out_tx.clone(),
},
);
}
instrument::connection_opened();
let _ = self
.connection_events_tx
.send(ConnectionEvent::Connected(node_key));
tracing::info!("Accepted connection from {remote_identity}");
Ok(IncomingConnection {
remote_identity,
connection,
control_tx: control_out_tx,
control_recv: recv,
node_class: peer_handshake.node_class,
node_metadata: peer_handshake.node_metadata,
is_edge_client: peer_handshake.is_edge_client,
})
}
}
impl Drop for Transport {
fn drop(&mut self) {
self.endpoint.close(quinn::VarInt::from_u32(0), b"shutdown");
}
}
async fn read_handshake(
mut recv: quinn::RecvStream,
) -> Result<(HandshakePayload, quinn::RecvStream), ClusterError> {
let mut codec = FrameCodec::new();
let mut buf = vec![0u8; 8192];
let frame = loop {
match recv
.read(&mut buf)
.await
.map_err(|e| ClusterError::HandshakeFailed(e.to_string()))?
{
Some(n) => {
codec.push_data(&buf[..n]);
if let Some(frame) = codec
.next_frame()
.map_err(|e| ClusterError::HandshakeFailed(e.to_string()))?
{
break frame;
}
}
None => {
return Err(ClusterError::HandshakeFailed(
"stream closed before handshake complete".into(),
));
}
}
};
let msg: ControlMessage =
framing::decode_message(&frame).map_err(ClusterError::Deserialization)?;
match msg {
ControlMessage::Handshake(payload) => Ok((payload, recv)),
other => Err(ClusterError::HandshakeFailed(format!(
"expected Handshake, got {other:?}"
))),
}
}
async fn run_control_stream_writer(
mut send: quinn::SendStream,
mut rx: mpsc::UnboundedReceiver<ControlMessage>,
shutdown: CancellationToken,
) {
loop {
tokio::select! {
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;
}
}
_ = shutdown.cancelled() => break,
}
}
let _ = send.finish();
}
pub async fn run_control_stream_reader(
mut recv: quinn::RecvStream,
tx: mpsc::UnboundedSender<(String, ControlMessage)>,
node_id: String,
shutdown: CancellationToken,
) {
let mut codec = FrameCodec::new();
let mut buf = vec![0u8; 8192];
loop {
tokio::select! {
result = recv.read(&mut buf) => {
match result {
Ok(Some(n)) => {
codec.push_data(&buf[..n]);
while let Ok(Some(frame)) = codec.next_frame() {
match framing::decode_message::<ControlMessage>(&frame) {
Ok(msg) => {
if tx.send((node_id.clone(), msg)).is_err() {
return;
}
}
Err(e) => {
tracing::warn!("Failed to decode control message from {node_id}: {e}");
}
}
}
}
Ok(None) => {
tracing::debug!("Control stream from {node_id} closed");
break;
}
Err(e) => {
tracing::warn!("Control stream read error from {node_id}: {e}");
break;
}
}
}
_ = shutdown.cancelled() => break,
}
}
}