use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::Arc;
use tokio::task::JoinSet;
use unb_core::validate_node_identifier;
use unb_runtime::Pipe;
use crate::host::HostConfig;
use crate::{
ConnectionStatus, Endpoint, EndpointSet, Hosting, Node, PeerConnection, TransportKind,
};
use unb_runtime::WsError;
pub async fn connect_unix(path: impl AsRef<std::path::Path>) -> Result<Pipe, WsError> {
Ok(Pipe::Piped {
pipe: unb_transport::unix::connect(path).await?,
initiator: true,
})
}
pub async fn accept_unix(listener: &unb_transport::unix::UnixListener) -> Result<Pipe, WsError> {
Ok(Pipe::Piped {
pipe: listener.accept().await?,
initiator: false,
})
}
#[derive(Clone, Debug)]
pub struct ParentLink {
pub node: String,
pub path: PathBuf,
}
impl ParentLink {
pub fn unix(node: impl Into<String>, path: impl Into<PathBuf>) -> ParentLink {
ParentLink {
node: node.into(),
path: path.into(),
}
}
fn validate(&self) -> Result<(), WsError> {
validate_node_identifier(&self.node)
.map_err(|error| WsError::Connect(error.to_string()))?;
if self.path.as_os_str().is_empty() {
return Err(WsError::Connect("parent Unix path is empty".into()));
}
Ok(())
}
}
#[derive(Default)]
pub struct TopologyConfig {
pub host: Option<HostConfig>,
pub parent: Option<ParentLink>,
}
impl TopologyConfig {
pub fn new() -> TopologyConfig {
TopologyConfig::default()
}
pub fn host(host: HostConfig) -> TopologyConfig {
TopologyConfig {
host: Some(host),
parent: None,
}
}
pub fn parent(parent: ParentLink) -> TopologyConfig {
TopologyConfig {
host: None,
parent: Some(parent),
}
}
pub fn with_parent(mut self, parent: ParentLink) -> TopologyConfig {
self.parent = Some(parent);
self
}
pub fn validate(&self) -> Result<(), WsError> {
if self.host.is_none() && self.parent.is_none() {
return Err(WsError::Connect(
"topology requires a host or parent".into(),
));
}
if let Some(host) = &self.host {
host.validate()
.map_err(|error| WsError::Connect(error.to_string()))?;
}
if let Some(parent) = &self.parent {
parent.validate()?;
}
Ok(())
}
}
pub struct UnbTopology {
hosting: Option<Hosting>,
parent: Option<PeerConnection>,
}
impl UnbTopology {
pub fn is_finished(&self) -> bool {
self.hosting.as_ref().is_some_and(Hosting::is_finished)
|| self.parent.as_ref().is_some_and(|parent| {
matches!(parent.status(), ConnectionStatus::Disconnected { .. })
})
}
pub fn health(&self) -> crate::host::HealthStatus {
let hosting = self.hosting.as_ref().map(Hosting::health);
let parent_link_ready = match &self.parent {
None => true,
Some(parent) => parent.status() == ConnectionStatus::Connected,
};
crate::host::HealthStatus {
process_alive: true,
websocket_bound: hosting.is_some_and(|health| health.websocket_bound),
websocket_addr: hosting.and_then(|health| health.websocket_addr),
webtransport_bound: hosting.is_some_and(|health| health.webtransport_bound),
webtransport_addr: hosting.and_then(|health| health.webtransport_addr),
listeners_running: hosting.is_some_and(|health| health.listeners_running),
parent_link_ready,
child_link_ready: true,
}
}
pub fn websocket_addr(&self) -> Option<SocketAddr> {
self.hosting.as_ref().and_then(Hosting::websocket_addr)
}
pub fn webtransport_addr(&self) -> Option<SocketAddr> {
self.hosting.as_ref().and_then(Hosting::webtransport_addr)
}
pub async fn shutdown(self) -> Result<(), WsError> {
if let Some(parent) = &self.parent {
parent.disconnect();
}
if let Some(hosting) = &self.hosting {
hosting.cancel();
}
let hosting = async {
match self.hosting {
Some(hosting) => hosting
.shutdown()
.await
.map_err(|error| WsError::Connect(error.to_string())),
None => Ok(()),
}
};
hosting.await
}
pub async fn wait(&mut self) -> Result<(), WsError> {
let host_failure = |error: crate::host::HostError| WsError::Connect(error.to_string());
match (&mut self.hosting, &mut self.parent) {
(Some(hosting), Some(parent)) => tokio::select! {
biased;
result = hosting.wait() => result.map_err(host_failure),
() = wait_for_parent_termination(parent.clone()) => Ok(()),
},
(Some(hosting), None) => hosting.wait().await.map_err(host_failure),
(None, Some(parent)) => {
wait_for_parent_termination(parent.clone()).await;
Ok(())
}
(None, None) => Ok(()),
}
}
}
async fn wait_for_parent_termination(parent: PeerConnection) {
loop {
let changed = parent.changed();
if matches!(parent.status(), ConnectionStatus::Disconnected { .. }) {
return;
}
let _ = changed.await;
}
}
impl Node {
pub async fn start_topology(
self: &Arc<Self>,
config: TopologyConfig,
) -> Result<UnbTopology, WsError> {
config.validate()?;
let hosting = match config.host {
Some(host) => Some(
host.start(self)
.await
.map_err(|error| WsError::Connect(error.to_string()))?,
),
None => None,
};
let parent = match config.parent {
Some(parent) => {
let endpoint = Endpoint {
kind: TransportKind::Unix,
address: parent.path.to_string_lossy().into_owned(),
cert_hash: None,
};
match self
.connect_expected(EndpointSet::from(endpoint), &parent.node)
.await
{
Ok(connection) if connection.peer() == parent.node => Some(connection),
Ok(connection) => {
let actual = connection.peer().to_owned();
connection.disconnect();
if let Some(hosting) = hosting {
let _ = hosting.shutdown().await;
}
return Err(WsError::Connect(format!(
"peer identity mismatch: expected {:?}, got {:?}",
parent.node, actual
)));
}
Err(error) => {
if let Some(hosting) = hosting {
let _ = hosting.shutdown().await;
}
return Err(WsError::Connect(error.to_string()));
}
}
}
None => None,
};
Ok(UnbTopology { hosting, parent })
}
}
pub struct UnixHosting {
cancellation: unb_runtime::CancellationToken,
task: tokio::task::JoinHandle<Result<(), WsError>>,
}
impl UnixHosting {
pub fn is_finished(&self) -> bool {
self.task.is_finished()
}
pub fn cancel(&self) {
self.cancellation.cancel();
}
pub async fn wait(&mut self) -> Result<(), WsError> {
(&mut self.task)
.await
.map_err(|error| WsError::Connect(error.to_string()))?
}
pub async fn shutdown(self) -> Result<(), WsError> {
self.cancellation.cancel();
self.task
.await
.map_err(|error| WsError::Connect(error.to_string()))?
}
}
impl Node {
pub fn host_unix_child(
self: &Arc<Self>,
path: impl AsRef<std::path::Path>,
expected_child: impl Into<String>,
) -> Result<UnixHosting, WsError> {
let expected_child = expected_child.into();
validate_node_identifier(&expected_child)
.map_err(|error| WsError::Connect(error.to_string()))?;
let listener = unb_transport::unix::UnixListener::bind(path)?;
let cancellation = self.cancellation().child_token();
let task_cancellation = cancellation.clone();
let node = self.clone();
let task = tokio::spawn(async move {
let mut connections = JoinSet::new();
loop {
tokio::select! {
biased;
() = task_cancellation.cancelled() => break,
result = connections.join_next(), if !connections.is_empty() => {
result.expect("connection task exists").map_err(|error| WsError::Connect(error.to_string()))?;
}
pipe = accept_unix(&listener) => {
let pipe = pipe?;
let connection = tokio::select! {
biased;
() = task_cancellation.cancelled() => break,
connection = node.connect_transport(&expected_child, pipe) => connection,
};
if let Ok(connection) = connection {
let cancellation = task_cancellation.clone();
connections.spawn(async move {
tokio::select! {
biased;
() = cancellation.cancelled() => {
connection.shutdown();
connection.closed().await;
}
() = connection.closed() => {}
}
});
}
}
}
}
while let Some(result) = connections.join_next().await {
result.map_err(|error| WsError::Connect(error.to_string()))?;
}
Ok(())
});
Ok(UnixHosting { cancellation, task })
}
pub fn host_unix(self: &Arc<Self>, listener: unb_transport::unix::UnixListener) -> UnixHosting {
let cancellation = self.cancellation().child_token();
let task_cancellation = cancellation.clone();
let node = self.clone();
let task = tokio::spawn(async move {
let mut connections = JoinSet::new();
loop {
tokio::select! {
biased;
() = task_cancellation.cancelled() => break,
result = connections.join_next(), if !connections.is_empty() => {
result.expect("connection task exists").map_err(|error| WsError::Connect(error.to_string()))?;
}
pipe = accept_unix(&listener) => {
let pipe = pipe?;
let connection = tokio::select! {
biased;
() = task_cancellation.cancelled() => break,
connection = node.serve_transport(pipe) => connection,
};
let cancellation = task_cancellation.clone();
connections.spawn(async move {
tokio::select! {
biased;
() = cancellation.cancelled() => {
connection.shutdown();
connection.closed().await;
}
() = connection.closed() => {}
}
});
}
}
}
while let Some(result) = connections.join_next().await {
result.map_err(|error| WsError::Connect(error.to_string()))?;
}
Ok(())
});
UnixHosting { cancellation, task }
}
}