use std::net::SocketAddr;
use std::path::PathBuf;
use std::time::Duration;
use std::sync::Arc;
use tokio::task::JoinSet;
use unb_core::validate_node_identifier;
use unb_runtime::Pipe;
use crate::host::HostConfig;
use crate::{Hosting, Node};
use unb_runtime::WsError;
struct TopologySession {
wire: Arc<unb_runtime::Wire>,
}
impl TopologySession {
async fn connect(node: &Arc<Node>, peer: &str, pipe: Pipe) -> Result<Self, WsError> {
Ok(TopologySession {
wire: node.connect_transport(peer, pipe).await?,
})
}
async fn closed(&self) {
self.wire.closed().await;
}
fn shutdown(&self) {
self.wire.shutdown();
}
}
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,
pub reconnect: ReconnectPolicy,
}
impl ParentLink {
pub fn unix(node: impl Into<String>, path: impl Into<PathBuf>) -> ParentLink {
ParentLink {
node: node.into(),
path: path.into(),
reconnect: ReconnectPolicy::default(),
}
}
pub fn reconnect_policy(mut self, reconnect: ReconnectPolicy) -> ParentLink {
self.reconnect = reconnect;
self
}
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()));
}
self.reconnect.validate()
}
}
#[derive(Clone, Debug)]
pub struct ReconnectPolicy {
pub initial_delay: Duration,
pub max_delay: Duration,
}
impl ReconnectPolicy {
pub fn new(initial_delay: Duration, max_delay: Duration) -> ReconnectPolicy {
ReconnectPolicy {
initial_delay,
max_delay,
}
}
pub fn validate(&self) -> Result<(), WsError> {
if self.initial_delay.is_zero() {
return Err(WsError::Connect(
"parent reconnect initial delay must be nonzero".into(),
));
}
if self.max_delay < self.initial_delay {
return Err(WsError::Connect(
"parent reconnect maximum delay must not be less than its initial delay".into(),
));
}
Ok(())
}
}
impl Default for ReconnectPolicy {
fn default() -> Self {
ReconnectPolicy {
initial_delay: Duration::from_millis(100),
max_delay: Duration::from_secs(5),
}
}
}
#[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<ParentSupervisor>,
}
struct ParentSupervisor {
cancellation: unb_runtime::CancellationToken,
task: Option<tokio::task::JoinHandle<Result<(), WsError>>>,
}
impl ParentSupervisor {
fn cancel(&self) {
self.cancellation.cancel();
}
async fn wait(&mut self) -> Result<(), WsError> {
if self.task.is_none() {
return Ok(());
}
let joined = {
let task = self.task.as_mut().expect("parent task present");
task.await
};
self.task = None;
joined.map_err(|error| WsError::Connect(error.to_string()))?
}
async fn shutdown(self) -> Result<(), WsError> {
self.cancellation.cancel();
match self.task {
Some(task) => task
.await
.map_err(|error| WsError::Connect(error.to_string()))?,
None => Ok(()),
}
}
}
impl UnbTopology {
pub fn is_finished(&self) -> bool {
self.hosting.as_ref().is_some_and(Hosting::is_finished)
|| self
.parent
.as_ref()
.and_then(|parent| parent.task.as_ref())
.is_some_and(tokio::task::JoinHandle::is_finished)
}
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.task.as_ref().is_some_and(|task| !task.is_finished()),
};
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.cancel();
}
if let Some(hosting) = &self.hosting {
hosting.cancel();
}
let parent = async {
match self.parent {
Some(parent) => parent.shutdown().await,
None => Ok(()),
}
};
let hosting = async {
match self.hosting {
Some(hosting) => hosting
.shutdown()
.await
.map_err(|error| WsError::Connect(error.to_string())),
None => Ok(()),
}
};
let (parent, hosting) = tokio::join!(parent, hosting);
parent?;
hosting
}
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),
result = parent.wait() => result,
},
(Some(hosting), None) => hosting.wait().await.map_err(host_failure),
(None, Some(parent)) => parent.wait().await,
(None, None) => Ok(()),
}
}
}
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 connected = async {
let pipe = connect_unix(&parent.path).await?;
TopologySession::connect(self, &parent.node, pipe).await
}
.await;
match connected {
Ok(connection) => {
let cancellation = self.cancellation().child_token();
let task_cancellation = cancellation.clone();
let node = self.clone();
let task = tokio::spawn(async move {
supervise_parent(node, parent, connection, task_cancellation).await
});
Some(ParentSupervisor {
cancellation,
task: Some(task),
})
}
Err(error) => {
if let Some(hosting) = hosting {
let _ = hosting.shutdown().await;
}
return Err(error);
}
}
}
None => None,
};
Ok(UnbTopology { hosting, parent })
}
}
async fn supervise_parent(
node: Arc<Node>,
parent: ParentLink,
mut connection: TopologySession,
cancellation: unb_runtime::CancellationToken,
) -> Result<(), WsError> {
let mut delay = parent.reconnect.initial_delay;
loop {
tokio::select! {
biased;
() = cancellation.cancelled() => {
connection.shutdown();
connection.closed().await;
return Ok(());
}
() = connection.closed() => {}
}
loop {
tokio::select! {
biased;
() = cancellation.cancelled() => return Ok(()),
() = tokio::time::sleep(delay) => {}
}
let connected = tokio::select! {
biased;
() = cancellation.cancelled() => return Ok(()),
result = async {
let pipe = connect_unix(&parent.path).await?;
TopologySession::connect(&node, &parent.node, pipe).await
} => result,
};
match connected {
Ok(next) => {
connection = next;
delay = parent.reconnect.initial_delay;
break;
}
Err(_) => {
delay = delay.saturating_mul(2).min(parent.reconnect.max_delay);
}
}
}
}
}
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 }
}
}