use std::sync::Arc;
#[cfg(feature = "client")]
use std::sync::OnceLock;
#[cfg(feature = "client")]
use std::time::{Duration, Instant};
#[cfg(feature = "client")]
use futures_util::future::BoxFuture;
#[cfg(feature = "client")]
use tokio::sync::watch;
use crate::capabilities::{Capabilities, CapabilitySet};
use crate::connection::{establish, AcceptOptions, Connection, Role};
#[cfg(feature = "client")]
use crate::error::{ConnectError, TransportError};
use crate::error::Result;
use crate::serve::ServeCapability;
use crate::transport::Transport;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NodeKind {
Desktop,
Server,
Cli,
Service,
}
impl NodeKind {
pub fn as_str(&self) -> &'static str {
match self {
NodeKind::Desktop => "desktop",
NodeKind::Server => "server",
NodeKind::Cli => "cli",
NodeKind::Service => "service",
}
}
}
pub struct Node {
name: String,
kind: NodeKind,
capabilities: Arc<CapabilitySet>,
#[cfg(feature = "client")]
on_connect: Option<ConnectHook>,
}
pub struct NodeBuilder {
name: String,
kind: NodeKind,
capabilities: Vec<Arc<dyn ServeCapability>>,
#[cfg(feature = "client")]
on_connect: Option<ConnectHook>,
}
impl Node {
pub fn builder() -> NodeBuilder {
NodeBuilder {
name: "zyris-node".to_string(),
kind: NodeKind::Service,
capabilities: Vec::new(),
#[cfg(feature = "client")]
on_connect: None,
}
}
pub fn name(&self) -> &str {
&self.name
}
pub fn kind(&self) -> &NodeKind {
&self.kind
}
pub fn capabilities(&self) -> Capabilities {
Capabilities::new(self.capabilities.clone())
}
fn agent(&self) -> String {
format!(
"zyris/{} ({}; {})",
env!("CARGO_PKG_VERSION"),
self.name,
self.kind.as_str()
)
}
pub async fn connect_over(&self, transport: impl Transport) -> Result<Connection> {
establish(
Box::new(transport),
Role::Dial { agent: self.agent(), kind: self.kind.as_str().to_string() },
self.capabilities.clone(),
)
.await
}
#[cfg(feature = "client")]
pub async fn dial(&self, url: &str, token: &str) -> Result<Connection, ConnectError> {
if crate::tls::missing_for(url) {
return Err(ConnectError::NoTlsProvider);
}
let transport = crate::transport::ws::connect(url, token).await?;
Ok(self.connect_over(transport).await?)
}
#[cfg(feature = "client")]
fn share(&self) -> Node {
Node {
name: self.name.clone(),
kind: self.kind.clone(),
capabilities: self.capabilities.clone(),
on_connect: self.on_connect.clone(),
}
}
#[cfg(feature = "client")]
pub async fn connect(&self, url: &str, token: impl AsRef<str>) -> Result<Link, ConnectError> {
let node = self.share();
let url = url.to_string();
let token = token.as_ref().to_string();
let redial: Redial = Arc::new(move || {
let node = node.share();
let url = url.clone();
let token = token.clone();
Box::pin(async move { node.dial(&url, &token).await })
as BoxFuture<'static, Result<Connection, ConnectError>>
});
Link::start(redial, self.on_connect.clone()).await
}
pub async fn accept(
&self,
transport: impl Transport,
options: AcceptOptions,
) -> Result<Connection> {
establish(
Box::new(transport),
Role::Accept { options },
self.capabilities.clone(),
)
.await
}
}
impl NodeBuilder {
pub fn name(mut self, name: impl Into<String>) -> Self {
self.name = name.into();
self
}
#[cfg(feature = "hostname")]
pub fn name_from_hostname(mut self) -> Self {
if let Some(name) = crate::machine_name() {
self.name = name;
}
self
}
pub fn kind(mut self, kind: NodeKind) -> Self {
self.kind = kind;
self
}
pub fn capability(mut self, capability: impl ServeCapability) -> Self {
self.capabilities.push(Arc::new(capability));
self
}
pub fn capability_arc(mut self, capability: Arc<dyn ServeCapability>) -> Self {
self.capabilities.push(capability);
self
}
#[cfg(feature = "client")]
pub fn on_connect<F, Fut>(mut self, hook: F) -> Self
where
F: Fn(Connection) -> Fut + Send + Sync + 'static,
Fut: std::future::Future<Output = ()> + Send + 'static,
{
self.on_connect = Some(Arc::new(move |conn| Box::pin(hook(conn))));
self
}
pub fn build(self) -> Result<Node> {
let capabilities = CapabilitySet::new();
for cap in self.capabilities {
capabilities.push(cap)?;
}
Ok(Node {
name: self.name,
kind: self.kind,
capabilities,
#[cfg(feature = "client")]
on_connect: self.on_connect,
})
}
}
#[cfg(feature = "client")]
const STABLE_AFTER: Duration = Duration::from_secs(30);
#[cfg(feature = "client")]
const BACKOFF_MIN: Duration = Duration::from_secs(1);
#[cfg(feature = "client")]
const BACKOFF_MAX: Duration = Duration::from_secs(30);
#[cfg(feature = "client")]
const CLOSE_GRACE: Duration = Duration::from_millis(200);
#[cfg(feature = "client")]
type Redial = Arc<dyn Fn() -> BoxFuture<'static, Result<Connection, ConnectError>> + Send + Sync>;
#[cfg(feature = "client")]
type ConnectHook = Arc<dyn Fn(Connection) -> BoxFuture<'static, ()> + Send + Sync>;
#[cfg(feature = "client")]
#[derive(Debug, Clone)]
enum Ending {
Asked,
Revoked,
Unauthorized,
Version { ours: String, theirs: Option<String> },
Unreachable(TransportError),
NoTlsProvider,
}
#[cfg(feature = "client")]
impl Ending {
fn of(error: ConnectError) -> Ending {
match error {
ConnectError::Revoked => Ending::Revoked,
ConnectError::Unauthorized => Ending::Unauthorized,
ConnectError::VersionMismatch { ours, theirs } => Ending::Version { ours, theirs },
ConnectError::Unreachable(inner) => Ending::Unreachable(inner),
ConnectError::NoTlsProvider => Ending::NoTlsProvider,
}
}
fn into_result(self) -> Result<(), ConnectError> {
match self {
Ending::Asked => Ok(()),
Ending::Revoked => Err(ConnectError::Revoked),
Ending::Unauthorized => Err(ConnectError::Unauthorized),
Ending::Version { ours, theirs } => Err(ConnectError::VersionMismatch { ours, theirs }),
Ending::Unreachable(inner) => Err(ConnectError::Unreachable(inner)),
Ending::NoTlsProvider => Err(ConnectError::NoTlsProvider),
}
}
}
#[cfg(feature = "client")]
fn is_fatal(error: &ConnectError) -> bool {
match error {
ConnectError::Revoked
| ConnectError::Unauthorized
| ConnectError::VersionMismatch { .. }
| ConnectError::NoTlsProvider => true,
ConnectError::Unreachable(_) => false,
}
}
#[cfg(feature = "client")]
pub struct Link {
node_id: Arc<OnceLock<String>>,
stop: watch::Sender<bool>,
done: Arc<watch::Sender<Option<Ending>>>,
}
#[cfg(feature = "client")]
impl Link {
pub fn node_id(&self) -> &str {
self.node_id.get().map(String::as_str).unwrap_or_default()
}
pub async fn wait_closed(&self) -> Result<(), ConnectError> {
let mut settled = self.done.subscribe();
loop {
let ending = settled.borrow().clone();
if let Some(ending) = ending {
return ending.into_result();
}
if settled.changed().await.is_err() {
return Ok(());
}
}
}
pub async fn disconnect(self) {
let _ = self.stop.send(true);
let mut settled = self.done.subscribe();
loop {
let ended = settled.borrow().is_some();
if ended {
return;
}
if settled.changed().await.is_err() {
return;
}
}
}
async fn start(redial: Redial, on_connect: Option<ConnectHook>) -> Result<Link, ConnectError> {
let first = match redial().await {
Ok(conn) => Some(conn),
Err(error) if is_fatal(&error) => return Err(error),
Err(error) => {
tracing::warn!(%error, "connect failed");
None
}
};
let node_id = Arc::new(OnceLock::new());
if let Some(conn) = &first {
let _ = node_id.set(conn.info().node_id.clone());
}
let (stop, stop_rx) = watch::channel(false);
let (done, _) = watch::channel::<Option<Ending>>(None);
let done = Arc::new(done);
tokio::spawn(maintain(redial, on_connect, first, node_id.clone(), stop_rx, done.clone()));
Ok(Link { node_id, stop, done })
}
}
#[cfg(feature = "client")]
async fn maintain(
redial: Redial,
on_connect: Option<ConnectHook>,
first: Option<Connection>,
node_id: Arc<OnceLock<String>>,
mut stop: watch::Receiver<bool>,
done: Arc<watch::Sender<Option<Ending>>>,
) {
let mut backoff = BACKOFF_MIN;
let mut current = first;
if current.is_none() {
let (widened, asked) = wait_then_widen(backoff, &mut stop).await;
backoff = widened;
if asked {
let _ = done.send(Some(Ending::Asked));
return;
}
}
loop {
if *stop.borrow() {
let _ = done.send(Some(Ending::Asked));
return;
}
let conn = match current.take() {
Some(conn) => conn,
None => match redial().await {
Ok(conn) => conn,
Err(error) if is_fatal(&error) => {
tracing::warn!(%error, "the server refused this node; not dialling again");
let _ = done.send(Some(Ending::of(error)));
return;
}
Err(error) => {
tracing::warn!(%error, "connect failed");
let (widened, asked) = wait_then_widen(backoff, &mut stop).await;
backoff = widened;
if asked {
let _ = done.send(Some(Ending::Asked));
return;
}
continue;
}
},
};
let _ = node_id.set(conn.info().node_id.clone());
tracing::info!(
node_id = %conn.info().node_id,
conn_id = %conn.info().conn_id,
"connected"
);
if let Some(hook) = &on_connect {
tokio::spawn(hook(conn.clone()));
}
let up = Instant::now();
let asked = tokio::select! {
reason = conn.closed() => {
tracing::warn!(%reason, "disconnected");
false
}
_ = stop.changed() => true,
};
if asked {
conn.close("node disconnecting");
tokio::time::sleep(CLOSE_GRACE).await;
let _ = done.send(Some(Ending::Asked));
return;
}
if up.elapsed() >= STABLE_AFTER {
backoff = BACKOFF_MIN;
}
let (widened, asked) = wait_then_widen(backoff, &mut stop).await;
backoff = widened;
if asked {
let _ = done.send(Some(Ending::Asked));
return;
}
}
}
#[cfg(feature = "client")]
async fn wait_then_widen(backoff: Duration, stop: &mut watch::Receiver<bool>) -> (Duration, bool) {
let wait = jitter(backoff);
tracing::info!(seconds = wait.as_secs_f64(), "reconnecting");
let asked = tokio::select! {
_ = tokio::time::sleep(wait) => false,
_ = stop.changed() => true,
};
((backoff * 2).min(BACKOFF_MAX), asked)
}
#[cfg(feature = "client")]
fn jitter(base: Duration) -> Duration {
use std::hash::{BuildHasher, Hasher, RandomState};
let mut hasher = RandomState::new().build_hasher();
hasher.write_u8(0);
let factor = 0.8 + (hasher.finish() % 400) as f64 / 1000.0;
base.mul_f64(factor)
}
#[cfg(all(test, feature = "client"))]
mod link_tests {
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use futures_util::future::BoxFuture;
use super::*;
use crate::connection::Connection;
const A_MOMENT: Duration = Duration::from_millis(200);
fn duplex_redial(served: Arc<Mutex<Vec<Connection>>>, dials: Arc<AtomicUsize>) -> Redial {
Arc::new(move || {
let served = served.clone();
let dials = dials.clone();
Box::pin(async move {
dials.fetch_add(1, Ordering::SeqCst);
let dialer = Node::builder().name("probe").kind(NodeKind::Cli).build().unwrap();
let acceptor =
Node::builder().name("server").kind(NodeKind::Server).build().unwrap();
let (mine, theirs) = crate::testing::duplex(&dialer, &acceptor)
.await
.map_err(|e| ConnectError::Unreachable(TransportError::Io(e.to_string())))?;
served.lock().unwrap().push(theirs);
Ok(mine)
}) as BoxFuture<'static, Result<Connection, ConnectError>>
})
}
async fn wait_until(what: &str, mut settled: impl FnMut() -> bool) {
let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
while !settled() {
assert!(tokio::time::Instant::now() < deadline, "the link never {what}");
tokio::time::sleep(Duration::from_millis(25)).await;
}
}
#[tokio::test]
async fn a_refused_token_stops_the_link_instead_of_dialling_again() {
let dials = Arc::new(AtomicUsize::new(0));
let counted = dials.clone();
let redial: Redial = Arc::new(move || {
let counted = counted.clone();
Box::pin(async move {
counted.fetch_add(1, Ordering::SeqCst);
Err(ConnectError::Unauthorized)
}) as BoxFuture<'static, Result<Connection, ConnectError>>
});
let error = Link::start(redial, None).await.err().expect("a refusal is not a link");
assert!(matches!(error, ConnectError::Unauthorized), "got: {error}");
tokio::time::sleep(A_MOMENT).await;
assert_eq!(
dials.load(Ordering::SeqCst),
1,
"a refusal is answered once; answering it in a loop is what the type exists to stop"
);
}
#[tokio::test]
async fn a_link_reports_the_node_id_the_server_handed_back() {
let served = Arc::new(Mutex::new(Vec::new()));
let dials = Arc::new(AtomicUsize::new(0));
let link = Link::start(duplex_redial(served.clone(), dials), None)
.await
.expect("an in-process duplex comes up");
let assigned = served.lock().unwrap()[0].info().node_id.clone();
assert!(!assigned.is_empty(), "the acceptor assigns one");
assert_eq!(link.node_id(), assigned, "the link reports what the acceptor assigned");
}
#[tokio::test]
async fn the_connect_hook_runs_again_on_the_connection_that_replaces_a_dropped_one() {
let served = Arc::new(Mutex::new(Vec::new()));
let dials = Arc::new(AtomicUsize::new(0));
let hooked = Arc::new(AtomicUsize::new(0));
let counted = hooked.clone();
let hook: ConnectHook = Arc::new(move |_conn| {
let counted = counted.clone();
Box::pin(async move {
counted.fetch_add(1, Ordering::SeqCst);
}) as BoxFuture<'static, ()>
});
let _link = Link::start(duplex_redial(served.clone(), dials.clone()), Some(hook))
.await
.expect("an in-process duplex comes up");
wait_until("ran the hook on the first connection", || {
hooked.load(Ordering::SeqCst) == 1
})
.await;
served.lock().unwrap()[0].close("server restarting");
wait_until("ran the hook again after reconnecting", || {
hooked.load(Ordering::SeqCst) == 2
})
.await;
assert_eq!(dials.load(Ordering::SeqCst), 2, "and it got there by dialling again");
}
#[tokio::test]
async fn one_connection_dropping_is_not_the_link_finishing() {
let served = Arc::new(Mutex::new(Vec::new()));
let dials = Arc::new(AtomicUsize::new(0));
let link = Link::start(duplex_redial(served.clone(), dials.clone()), None)
.await
.expect("an in-process duplex comes up");
served.lock().unwrap()[0].close("server restarting");
wait_until("dialled again", || dials.load(Ordering::SeqCst) == 2).await;
assert!(
tokio::time::timeout(A_MOMENT, link.wait_closed()).await.is_err(),
"a link that reconnected is still up, and has to keep whoever waits on it waiting"
);
}
#[tokio::test]
async fn a_link_the_server_disowns_hands_the_reason_to_whoever_waits_for_it() {
let served = Arc::new(Mutex::new(Vec::new()));
let dials = Arc::new(AtomicUsize::new(0));
let live = duplex_redial(served.clone(), dials);
let attempts = Arc::new(AtomicUsize::new(0));
let redial: Redial = Arc::new(move || {
let live = live.clone();
let first = attempts.fetch_add(1, Ordering::SeqCst) == 0;
Box::pin(async move {
if first {
live().await
} else {
Err(ConnectError::Revoked)
}
}) as BoxFuture<'static, Result<Connection, ConnectError>>
});
let link = Link::start(redial, None).await.expect("the first dial comes up");
served.lock().unwrap()[0].close("credential revoked");
let ended = tokio::time::timeout(Duration::from_secs(10), link.wait_closed())
.await
.expect("a revoked credential ends the link rather than looping on it");
assert!(matches!(ended, Err(ConnectError::Revoked)), "got: {ended:?}");
}
#[tokio::test]
async fn ending_a_link_closes_the_connection_it_was_keeping_up() {
let served = Arc::new(Mutex::new(Vec::new()));
let dials = Arc::new(AtomicUsize::new(0));
let link = Link::start(duplex_redial(served.clone(), dials.clone()), None)
.await
.expect("an in-process duplex comes up");
assert!(
tokio::time::timeout(A_MOMENT, link.wait_closed()).await.is_err(),
"a link that is still connected has not finished"
);
tokio::time::timeout(Duration::from_secs(5), link.disconnect())
.await
.expect("disconnect returns once the link has actually ended");
let peer = served.lock().unwrap().remove(0);
tokio::time::timeout(Duration::from_secs(5), peer.closed())
.await
.expect("the other end has to see the connection go, not wait for a heartbeat to lapse");
assert_eq!(dials.load(Ordering::SeqCst), 1, "an ended link does not dial again");
}
}