use std::{
io,
net::SocketAddr,
sync::{OnceLock, atomic::Ordering},
};
use tokio::sync::{mpsc, oneshot};
use crate::{
connections::{Connection, DisconnectOrigin},
node::Node,
protocols::{on_connect::OnConnectBundle, on_disconnect::OnDisconnectBundle},
};
mod handshake;
mod on_connect;
mod on_disconnect;
mod reading;
mod writing;
pub use handshake::Handshake;
pub use on_connect::OnConnect;
pub use on_disconnect::OnDisconnect;
pub use reading::Reading;
pub use writing::Writing;
#[derive(Default)]
pub(crate) struct Protocols {
pub(crate) handshake: OnceLock<ProtocolHandler<Connection, io::Result<Connection>>>,
pub(crate) reading: OnceLock<ProtocolHandler<Connection, io::Result<Connection>>>,
pub(crate) writing: OnceLock<writing::WritingHandler>,
pub(crate) on_connect: OnceLock<ProtocolHandler<SocketAddr, OnConnectBundle>>,
pub(crate) on_disconnect:
OnceLock<ProtocolHandler<(SocketAddr, DisconnectOrigin), OnDisconnectBundle>>,
}
pub(crate) type ReturnableItem<T, U> = (T, oneshot::Sender<U>);
pub(crate) type ReturnableConnection = ReturnableItem<Connection, io::Result<Connection>>;
pub(crate) struct ProtocolHandler<T, U>(mpsc::Sender<ReturnableItem<T, U>>);
pub(crate) trait Protocol<T, U> {
async fn trigger(&self, item: ReturnableItem<T, U>);
}
impl<T, U> Protocol<T, U> for ProtocolHandler<T, U> {
async fn trigger(&self, item: ReturnableItem<T, U>) {
let _ = self.0.send(item).await;
}
}
pub(crate) struct DisconnectOnDrop {
pub(crate) node: Option<Node>,
pub(crate) addr: SocketAddr,
pub(crate) origin: DisconnectOrigin,
}
impl DisconnectOnDrop {
pub(crate) fn new(node: Node, addr: SocketAddr, origin: DisconnectOrigin) -> Self {
Self {
node: Some(node),
addr,
origin,
}
}
}
impl Drop for DisconnectOnDrop {
fn drop(&mut self) {
if let Some(node) = self.node.take() {
let (addr, origin) = (self.addr, self.origin);
let needs_recovery = node
.connections
.active
.read()
.get(&addr)
.is_some_and(|c| !c.disconnecting.load(Ordering::Acquire));
if needs_recovery && let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move { node.disconnect_w_origin(addr, origin).await });
}
}
}
}
pub(crate) fn log_setup_join(
span: &tracing::Span,
protocol: &'static str,
res: Option<Result<(), tokio::task::JoinError>>,
) {
if let Some(Err(e)) = res
&& e.is_panic()
{
tracing::error!(parent: span, "a {protocol} setup task panicked: {e}");
}
}
pub(crate) fn panic_message(payload: &(dyn std::any::Any + Send)) -> &str {
payload
.downcast_ref::<&'static str>()
.copied()
.or_else(|| payload.downcast_ref::<String>().map(String::as_str))
.unwrap_or("<non-string panic payload>")
}