1use std::sync::Arc;
2#[cfg(feature = "client")]
3use std::sync::OnceLock;
4#[cfg(feature = "client")]
5use std::time::{Duration, Instant};
6
7#[cfg(feature = "client")]
8use futures_util::future::BoxFuture;
9#[cfg(feature = "client")]
10use tokio::sync::watch;
11
12use crate::capabilities::{Capabilities, CapabilitySet};
13use crate::connection::{establish, AcceptOptions, Connection, Role};
14#[cfg(feature = "client")]
15use crate::error::{ConnectError, TransportError};
16use crate::error::Result;
17use crate::serve::ServeCapability;
18use crate::transport::Transport;
19
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum NodeKind {
22 Desktop,
23 Server,
24 Cli,
25 Service,
26}
27
28impl NodeKind {
29 pub fn as_str(&self) -> &'static str {
34 match self {
35 NodeKind::Desktop => "desktop",
36 NodeKind::Server => "server",
37 NodeKind::Cli => "cli",
38 NodeKind::Service => "service",
39 }
40 }
41}
42
43pub struct Node {
44 name: String,
45 kind: NodeKind,
46 capabilities: Arc<CapabilitySet>,
47 #[cfg(feature = "client")]
48 on_connect: Option<ConnectHook>,
49}
50
51pub struct NodeBuilder {
52 name: String,
53 kind: NodeKind,
54 capabilities: Vec<Arc<dyn ServeCapability>>,
55 #[cfg(feature = "client")]
56 on_connect: Option<ConnectHook>,
57}
58
59impl Node {
60 pub fn builder() -> NodeBuilder {
61 NodeBuilder {
62 name: "zyris-node".to_string(),
63 kind: NodeKind::Service,
64 capabilities: Vec::new(),
65 #[cfg(feature = "client")]
66 on_connect: None,
67 }
68 }
69
70 pub fn name(&self) -> &str {
71 &self.name
72 }
73
74 pub fn kind(&self) -> &NodeKind {
75 &self.kind
76 }
77
78 pub fn capabilities(&self) -> Capabilities {
80 Capabilities::new(self.capabilities.clone())
81 }
82
83 fn agent(&self) -> String {
84 format!(
85 "zyris/{} ({}; {})",
86 env!("CARGO_PKG_VERSION"),
87 self.name,
88 self.kind.as_str()
89 )
90 }
91
92 pub async fn connect_over(&self, transport: impl Transport) -> Result<Connection> {
93 establish(
94 Box::new(transport),
95 Role::Dial { agent: self.agent(), kind: self.kind.as_str().to_string() },
96 self.capabilities.clone(),
97 )
98 .await
99 }
100
101 #[cfg(feature = "client")]
109 pub async fn dial(&self, url: &str, token: &str) -> Result<Connection, ConnectError> {
110 if crate::tls::missing_for(url) {
113 return Err(ConnectError::NoTlsProvider);
114 }
115 let transport = crate::transport::ws::connect(url, token).await?;
116 Ok(self.connect_over(transport).await?)
117 }
118
119 #[cfg(feature = "client")]
122 fn share(&self) -> Node {
123 Node {
124 name: self.name.clone(),
125 kind: self.kind.clone(),
126 capabilities: self.capabilities.clone(),
127 on_connect: self.on_connect.clone(),
128 }
129 }
130
131 #[cfg(feature = "client")]
141 pub async fn connect(&self, url: &str, token: impl AsRef<str>) -> Result<Link, ConnectError> {
142 let node = self.share();
143 let url = url.to_string();
144 let token = token.as_ref().to_string();
145 let redial: Redial = Arc::new(move || {
146 let node = node.share();
147 let url = url.clone();
148 let token = token.clone();
149 Box::pin(async move { node.dial(&url, &token).await })
150 as BoxFuture<'static, Result<Connection, ConnectError>>
151 });
152 Link::start(redial, self.on_connect.clone()).await
153 }
154
155 pub async fn accept(
156 &self,
157 transport: impl Transport,
158 options: AcceptOptions,
159 ) -> Result<Connection> {
160 establish(
161 Box::new(transport),
162 Role::Accept { options },
163 self.capabilities.clone(),
164 )
165 .await
166 }
167}
168
169impl NodeBuilder {
170 pub fn name(mut self, name: impl Into<String>) -> Self {
171 self.name = name.into();
172 self
173 }
174
175 #[cfg(feature = "hostname")]
180 pub fn name_from_hostname(mut self) -> Self {
181 if let Some(name) = crate::machine_name() {
182 self.name = name;
183 }
184 self
185 }
186
187 pub fn kind(mut self, kind: NodeKind) -> Self {
188 self.kind = kind;
189 self
190 }
191
192 pub fn capability(mut self, capability: impl ServeCapability) -> Self {
193 self.capabilities.push(Arc::new(capability));
194 self
195 }
196
197 pub fn capability_arc(mut self, capability: Arc<dyn ServeCapability>) -> Self {
198 self.capabilities.push(capability);
199 self
200 }
201
202 #[cfg(feature = "client")]
209 pub fn on_connect<F, Fut>(mut self, hook: F) -> Self
210 where
211 F: Fn(Connection) -> Fut + Send + Sync + 'static,
212 Fut: std::future::Future<Output = ()> + Send + 'static,
213 {
214 self.on_connect = Some(Arc::new(move |conn| Box::pin(hook(conn))));
215 self
216 }
217
218 pub fn build(self) -> Result<Node> {
219 let capabilities = CapabilitySet::new();
220 for cap in self.capabilities {
221 capabilities.push(cap)?;
222 }
223 Ok(Node {
224 name: self.name,
225 kind: self.kind,
226 capabilities,
227 #[cfg(feature = "client")]
228 on_connect: self.on_connect,
229 })
230 }
231}
232
233#[cfg(feature = "client")]
237const STABLE_AFTER: Duration = Duration::from_secs(30);
238#[cfg(feature = "client")]
239const BACKOFF_MIN: Duration = Duration::from_secs(1);
240#[cfg(feature = "client")]
241const BACKOFF_MAX: Duration = Duration::from_secs(30);
242#[cfg(feature = "client")]
245const CLOSE_GRACE: Duration = Duration::from_millis(200);
246
247#[cfg(feature = "client")]
250type Redial = Arc<dyn Fn() -> BoxFuture<'static, Result<Connection, ConnectError>> + Send + Sync>;
251
252#[cfg(feature = "client")]
254type ConnectHook = Arc<dyn Fn(Connection) -> BoxFuture<'static, ()> + Send + Sync>;
255
256#[cfg(feature = "client")]
262#[derive(Debug, Clone)]
263enum Ending {
264 Asked,
265 Revoked,
266 Unauthorized,
267 Version { ours: String, theirs: Option<String> },
268 Unreachable(TransportError),
269 NoTlsProvider,
270}
271
272#[cfg(feature = "client")]
273impl Ending {
274 fn of(error: ConnectError) -> Ending {
277 match error {
278 ConnectError::Revoked => Ending::Revoked,
279 ConnectError::Unauthorized => Ending::Unauthorized,
280 ConnectError::VersionMismatch { ours, theirs } => Ending::Version { ours, theirs },
281 ConnectError::Unreachable(inner) => Ending::Unreachable(inner),
282 ConnectError::NoTlsProvider => Ending::NoTlsProvider,
283 }
284 }
285
286 fn into_result(self) -> Result<(), ConnectError> {
287 match self {
288 Ending::Asked => Ok(()),
289 Ending::Revoked => Err(ConnectError::Revoked),
290 Ending::Unauthorized => Err(ConnectError::Unauthorized),
291 Ending::Version { ours, theirs } => Err(ConnectError::VersionMismatch { ours, theirs }),
292 Ending::Unreachable(inner) => Err(ConnectError::Unreachable(inner)),
293 Ending::NoTlsProvider => Err(ConnectError::NoTlsProvider),
294 }
295 }
296}
297
298#[cfg(feature = "client")]
300fn is_fatal(error: &ConnectError) -> bool {
301 match error {
302 ConnectError::Revoked
303 | ConnectError::Unauthorized
304 | ConnectError::VersionMismatch { .. }
305 | ConnectError::NoTlsProvider => true,
306 ConnectError::Unreachable(_) => false,
307 }
308}
309
310#[cfg(feature = "client")]
316pub struct Link {
317 node_id: Arc<OnceLock<String>>,
322 stop: watch::Sender<bool>,
325 done: Arc<watch::Sender<Option<Ending>>>,
326}
327
328#[cfg(feature = "client")]
329impl Link {
330 pub fn node_id(&self) -> &str {
332 self.node_id.get().map(String::as_str).unwrap_or_default()
333 }
334
335 pub async fn wait_closed(&self) -> Result<(), ConnectError> {
338 let mut settled = self.done.subscribe();
339 loop {
340 let ending = settled.borrow().clone();
341 if let Some(ending) = ending {
342 return ending.into_result();
343 }
344 if settled.changed().await.is_err() {
345 return Ok(());
346 }
347 }
348 }
349
350 pub async fn disconnect(self) {
355 let _ = self.stop.send(true);
356 let mut settled = self.done.subscribe();
357 loop {
358 let ended = settled.borrow().is_some();
359 if ended {
360 return;
361 }
362 if settled.changed().await.is_err() {
363 return;
364 }
365 }
366 }
367
368 async fn start(redial: Redial, on_connect: Option<ConnectHook>) -> Result<Link, ConnectError> {
371 let first = match redial().await {
372 Ok(conn) => Some(conn),
373 Err(error) if is_fatal(&error) => return Err(error),
374 Err(error) => {
375 tracing::warn!(%error, "connect failed");
376 None
377 }
378 };
379 let node_id = Arc::new(OnceLock::new());
380 if let Some(conn) = &first {
383 let _ = node_id.set(conn.info().node_id.clone());
384 }
385 let (stop, stop_rx) = watch::channel(false);
386 let (done, _) = watch::channel::<Option<Ending>>(None);
387 let done = Arc::new(done);
388 tokio::spawn(maintain(redial, on_connect, first, node_id.clone(), stop_rx, done.clone()));
389 Ok(Link { node_id, stop, done })
390 }
391}
392
393#[cfg(feature = "client")]
399async fn maintain(
400 redial: Redial,
401 on_connect: Option<ConnectHook>,
402 first: Option<Connection>,
403 node_id: Arc<OnceLock<String>>,
404 mut stop: watch::Receiver<bool>,
405 done: Arc<watch::Sender<Option<Ending>>>,
406) {
407 let mut backoff = BACKOFF_MIN;
408 let mut current = first;
409 if current.is_none() {
410 let (widened, asked) = wait_then_widen(backoff, &mut stop).await;
413 backoff = widened;
414 if asked {
415 let _ = done.send(Some(Ending::Asked));
416 return;
417 }
418 }
419 loop {
420 if *stop.borrow() {
423 let _ = done.send(Some(Ending::Asked));
424 return;
425 }
426
427 let conn = match current.take() {
428 Some(conn) => conn,
429 None => match redial().await {
430 Ok(conn) => conn,
431 Err(error) if is_fatal(&error) => {
432 tracing::warn!(%error, "the server refused this node; not dialling again");
433 let _ = done.send(Some(Ending::of(error)));
434 return;
435 }
436 Err(error) => {
437 tracing::warn!(%error, "connect failed");
438 let (widened, asked) = wait_then_widen(backoff, &mut stop).await;
439 backoff = widened;
440 if asked {
441 let _ = done.send(Some(Ending::Asked));
442 return;
443 }
444 continue;
445 }
446 },
447 };
448
449 let _ = node_id.set(conn.info().node_id.clone());
450 tracing::info!(
451 node_id = %conn.info().node_id,
452 conn_id = %conn.info().conn_id,
453 "connected"
454 );
455 if let Some(hook) = &on_connect {
456 tokio::spawn(hook(conn.clone()));
457 }
458
459 let up = Instant::now();
460 let asked = tokio::select! {
461 reason = conn.closed() => {
462 tracing::warn!(%reason, "disconnected");
463 false
464 }
465 _ = stop.changed() => true,
467 };
468 if asked {
469 conn.close("node disconnecting");
470 tokio::time::sleep(CLOSE_GRACE).await;
471 let _ = done.send(Some(Ending::Asked));
472 return;
473 }
474
475 if up.elapsed() >= STABLE_AFTER {
476 backoff = BACKOFF_MIN;
477 }
478 let (widened, asked) = wait_then_widen(backoff, &mut stop).await;
479 backoff = widened;
480 if asked {
481 let _ = done.send(Some(Ending::Asked));
482 return;
483 }
484 }
485}
486
487#[cfg(feature = "client")]
488async fn wait_then_widen(backoff: Duration, stop: &mut watch::Receiver<bool>) -> (Duration, bool) {
489 let wait = jitter(backoff);
490 tracing::info!(seconds = wait.as_secs_f64(), "reconnecting");
491 let asked = tokio::select! {
492 _ = tokio::time::sleep(wait) => false,
493 _ = stop.changed() => true,
494 };
495 ((backoff * 2).min(BACKOFF_MAX), asked)
496}
497
498#[cfg(feature = "client")]
504fn jitter(base: Duration) -> Duration {
505 use std::hash::{BuildHasher, Hasher, RandomState};
506 let mut hasher = RandomState::new().build_hasher();
507 hasher.write_u8(0);
508 let factor = 0.8 + (hasher.finish() % 400) as f64 / 1000.0;
509 base.mul_f64(factor)
510}
511
512#[cfg(all(test, feature = "client"))]
513mod link_tests {
514 use std::sync::atomic::{AtomicUsize, Ordering};
515 use std::sync::{Arc, Mutex};
516 use std::time::Duration;
517
518 use futures_util::future::BoxFuture;
519
520 use super::*;
521 use crate::connection::Connection;
522
523 const A_MOMENT: Duration = Duration::from_millis(200);
526
527 fn duplex_redial(served: Arc<Mutex<Vec<Connection>>>, dials: Arc<AtomicUsize>) -> Redial {
530 Arc::new(move || {
531 let served = served.clone();
532 let dials = dials.clone();
533 Box::pin(async move {
534 dials.fetch_add(1, Ordering::SeqCst);
535 let dialer = Node::builder().name("probe").kind(NodeKind::Cli).build().unwrap();
536 let acceptor =
537 Node::builder().name("server").kind(NodeKind::Server).build().unwrap();
538 let (mine, theirs) = crate::testing::duplex(&dialer, &acceptor)
539 .await
540 .map_err(|e| ConnectError::Unreachable(TransportError::Io(e.to_string())))?;
541 served.lock().unwrap().push(theirs);
542 Ok(mine)
543 }) as BoxFuture<'static, Result<Connection, ConnectError>>
544 })
545 }
546
547 async fn wait_until(what: &str, mut settled: impl FnMut() -> bool) {
550 let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
551 while !settled() {
552 assert!(tokio::time::Instant::now() < deadline, "the link never {what}");
553 tokio::time::sleep(Duration::from_millis(25)).await;
554 }
555 }
556
557 #[tokio::test]
561 async fn a_refused_token_stops_the_link_instead_of_dialling_again() {
562 let dials = Arc::new(AtomicUsize::new(0));
563 let counted = dials.clone();
564 let redial: Redial = Arc::new(move || {
565 let counted = counted.clone();
566 Box::pin(async move {
567 counted.fetch_add(1, Ordering::SeqCst);
568 Err(ConnectError::Unauthorized)
569 }) as BoxFuture<'static, Result<Connection, ConnectError>>
570 });
571
572 let error = Link::start(redial, None).await.err().expect("a refusal is not a link");
573 assert!(matches!(error, ConnectError::Unauthorized), "got: {error}");
574
575 tokio::time::sleep(A_MOMENT).await;
576 assert_eq!(
577 dials.load(Ordering::SeqCst),
578 1,
579 "a refusal is answered once; answering it in a loop is what the type exists to stop"
580 );
581 }
582
583 #[tokio::test]
586 async fn a_link_reports_the_node_id_the_server_handed_back() {
587 let served = Arc::new(Mutex::new(Vec::new()));
588 let dials = Arc::new(AtomicUsize::new(0));
589 let link = Link::start(duplex_redial(served.clone(), dials), None)
590 .await
591 .expect("an in-process duplex comes up");
592
593 let assigned = served.lock().unwrap()[0].info().node_id.clone();
594 assert!(!assigned.is_empty(), "the acceptor assigns one");
595 assert_eq!(link.node_id(), assigned, "the link reports what the acceptor assigned");
596 }
597
598 #[tokio::test]
602 async fn the_connect_hook_runs_again_on_the_connection_that_replaces_a_dropped_one() {
603 let served = Arc::new(Mutex::new(Vec::new()));
604 let dials = Arc::new(AtomicUsize::new(0));
605 let hooked = Arc::new(AtomicUsize::new(0));
606 let counted = hooked.clone();
607 let hook: ConnectHook = Arc::new(move |_conn| {
608 let counted = counted.clone();
609 Box::pin(async move {
610 counted.fetch_add(1, Ordering::SeqCst);
611 }) as BoxFuture<'static, ()>
612 });
613
614 let _link = Link::start(duplex_redial(served.clone(), dials.clone()), Some(hook))
615 .await
616 .expect("an in-process duplex comes up");
617 wait_until("ran the hook on the first connection", || {
618 hooked.load(Ordering::SeqCst) == 1
619 })
620 .await;
621
622 served.lock().unwrap()[0].close("server restarting");
623 wait_until("ran the hook again after reconnecting", || {
624 hooked.load(Ordering::SeqCst) == 2
625 })
626 .await;
627 assert_eq!(dials.load(Ordering::SeqCst), 2, "and it got there by dialling again");
628 }
629
630 #[tokio::test]
633 async fn one_connection_dropping_is_not_the_link_finishing() {
634 let served = Arc::new(Mutex::new(Vec::new()));
635 let dials = Arc::new(AtomicUsize::new(0));
636 let link = Link::start(duplex_redial(served.clone(), dials.clone()), None)
637 .await
638 .expect("an in-process duplex comes up");
639
640 served.lock().unwrap()[0].close("server restarting");
641 wait_until("dialled again", || dials.load(Ordering::SeqCst) == 2).await;
642
643 assert!(
644 tokio::time::timeout(A_MOMENT, link.wait_closed()).await.is_err(),
645 "a link that reconnected is still up, and has to keep whoever waits on it waiting"
646 );
647 }
648
649 #[tokio::test]
653 async fn a_link_the_server_disowns_hands_the_reason_to_whoever_waits_for_it() {
654 let served = Arc::new(Mutex::new(Vec::new()));
655 let dials = Arc::new(AtomicUsize::new(0));
656 let live = duplex_redial(served.clone(), dials);
657 let attempts = Arc::new(AtomicUsize::new(0));
658 let redial: Redial = Arc::new(move || {
659 let live = live.clone();
660 let first = attempts.fetch_add(1, Ordering::SeqCst) == 0;
661 Box::pin(async move {
662 if first {
663 live().await
664 } else {
665 Err(ConnectError::Revoked)
666 }
667 }) as BoxFuture<'static, Result<Connection, ConnectError>>
668 });
669
670 let link = Link::start(redial, None).await.expect("the first dial comes up");
671 served.lock().unwrap()[0].close("credential revoked");
672
673 let ended = tokio::time::timeout(Duration::from_secs(10), link.wait_closed())
674 .await
675 .expect("a revoked credential ends the link rather than looping on it");
676 assert!(matches!(ended, Err(ConnectError::Revoked)), "got: {ended:?}");
677 }
678 #[tokio::test]
682 async fn ending_a_link_closes_the_connection_it_was_keeping_up() {
683 let served = Arc::new(Mutex::new(Vec::new()));
684 let dials = Arc::new(AtomicUsize::new(0));
685 let link = Link::start(duplex_redial(served.clone(), dials.clone()), None)
686 .await
687 .expect("an in-process duplex comes up");
688
689 assert!(
690 tokio::time::timeout(A_MOMENT, link.wait_closed()).await.is_err(),
691 "a link that is still connected has not finished"
692 );
693
694 tokio::time::timeout(Duration::from_secs(5), link.disconnect())
695 .await
696 .expect("disconnect returns once the link has actually ended");
697
698 let peer = served.lock().unwrap().remove(0);
699 tokio::time::timeout(Duration::from_secs(5), peer.closed())
700 .await
701 .expect("the other end has to see the connection go, not wait for a heartbeat to lapse");
702 assert_eq!(dials.load(Ordering::SeqCst), 1, "an ended link does not dial again");
703 }
704}