1use std::net::SocketAddr;
2use std::path::PathBuf;
3use std::time::Duration;
4
5use std::sync::Arc;
6use tokio::task::JoinSet;
7use unb_core::validate_node_identifier;
8use unb_runtime::Pipe;
9
10use crate::host::HostConfig;
11use crate::{Hosting, Node};
12use unb_runtime::WsError;
13
14struct TopologySession {
15 wire: Arc<unb_runtime::Wire>,
16}
17
18impl TopologySession {
19 async fn connect(node: &Arc<Node>, peer: &str, pipe: Pipe) -> Result<Self, WsError> {
20 Ok(TopologySession {
21 wire: node.connect_transport(peer, pipe).await?,
22 })
23 }
24
25 async fn closed(&self) {
26 self.wire.closed().await;
27 }
28
29 fn shutdown(&self) {
30 self.wire.shutdown();
31 }
32}
33
34pub async fn connect_unix(path: impl AsRef<std::path::Path>) -> Result<Pipe, WsError> {
35 Ok(Pipe::Piped {
36 pipe: unb_transport::unix::connect(path).await?,
37 initiator: true,
38 })
39}
40
41pub async fn accept_unix(listener: &unb_transport::unix::UnixListener) -> Result<Pipe, WsError> {
42 Ok(Pipe::Piped {
43 pipe: listener.accept().await?,
44 initiator: false,
45 })
46}
47
48#[derive(Clone, Debug)]
49pub struct ParentLink {
50 pub node: String,
51 pub path: PathBuf,
52 pub reconnect: ReconnectPolicy,
53}
54
55impl ParentLink {
56 pub fn unix(node: impl Into<String>, path: impl Into<PathBuf>) -> ParentLink {
57 ParentLink {
58 node: node.into(),
59 path: path.into(),
60 reconnect: ReconnectPolicy::default(),
61 }
62 }
63
64 pub fn reconnect_policy(mut self, reconnect: ReconnectPolicy) -> ParentLink {
65 self.reconnect = reconnect;
66 self
67 }
68
69 fn validate(&self) -> Result<(), WsError> {
70 validate_node_identifier(&self.node)
71 .map_err(|error| WsError::Connect(error.to_string()))?;
72 if self.path.as_os_str().is_empty() {
73 return Err(WsError::Connect("parent Unix path is empty".into()));
74 }
75 self.reconnect.validate()
76 }
77}
78
79#[derive(Clone, Debug)]
80pub struct ReconnectPolicy {
81 pub initial_delay: Duration,
82 pub max_delay: Duration,
83}
84
85impl ReconnectPolicy {
86 pub fn new(initial_delay: Duration, max_delay: Duration) -> ReconnectPolicy {
87 ReconnectPolicy {
88 initial_delay,
89 max_delay,
90 }
91 }
92
93 pub fn validate(&self) -> Result<(), WsError> {
94 if self.initial_delay.is_zero() {
95 return Err(WsError::Connect(
96 "parent reconnect initial delay must be nonzero".into(),
97 ));
98 }
99 if self.max_delay < self.initial_delay {
100 return Err(WsError::Connect(
101 "parent reconnect maximum delay must not be less than its initial delay".into(),
102 ));
103 }
104 Ok(())
105 }
106}
107
108impl Default for ReconnectPolicy {
109 fn default() -> Self {
110 ReconnectPolicy {
111 initial_delay: Duration::from_millis(100),
112 max_delay: Duration::from_secs(5),
113 }
114 }
115}
116
117#[derive(Default)]
118pub struct TopologyConfig {
119 pub host: Option<HostConfig>,
120 pub parent: Option<ParentLink>,
121}
122
123impl TopologyConfig {
124 pub fn new() -> TopologyConfig {
125 TopologyConfig::default()
126 }
127
128 pub fn host(host: HostConfig) -> TopologyConfig {
129 TopologyConfig {
130 host: Some(host),
131 parent: None,
132 }
133 }
134
135 pub fn parent(parent: ParentLink) -> TopologyConfig {
136 TopologyConfig {
137 host: None,
138 parent: Some(parent),
139 }
140 }
141
142 pub fn with_parent(mut self, parent: ParentLink) -> TopologyConfig {
143 self.parent = Some(parent);
144 self
145 }
146
147 pub fn validate(&self) -> Result<(), WsError> {
148 if self.host.is_none() && self.parent.is_none() {
149 return Err(WsError::Connect(
150 "topology requires a host or parent".into(),
151 ));
152 }
153 if let Some(host) = &self.host {
154 host.validate()
155 .map_err(|error| WsError::Connect(error.to_string()))?;
156 }
157 if let Some(parent) = &self.parent {
158 parent.validate()?;
159 }
160 Ok(())
161 }
162}
163
164pub struct UnbTopology {
165 hosting: Option<Hosting>,
166 parent: Option<ParentSupervisor>,
167}
168
169struct ParentSupervisor {
170 cancellation: unb_runtime::CancellationToken,
171 task: Option<tokio::task::JoinHandle<Result<(), WsError>>>,
172}
173
174impl ParentSupervisor {
175 fn cancel(&self) {
176 self.cancellation.cancel();
177 }
178
179 async fn wait(&mut self) -> Result<(), WsError> {
180 if self.task.is_none() {
181 return Ok(());
182 }
183 let joined = {
184 let task = self.task.as_mut().expect("parent task present");
185 task.await
186 };
187 self.task = None;
188 joined.map_err(|error| WsError::Connect(error.to_string()))?
189 }
190
191 async fn shutdown(self) -> Result<(), WsError> {
192 self.cancellation.cancel();
193 match self.task {
194 Some(task) => task
195 .await
196 .map_err(|error| WsError::Connect(error.to_string()))?,
197 None => Ok(()),
198 }
199 }
200}
201
202impl UnbTopology {
203 pub fn is_finished(&self) -> bool {
204 self.hosting.as_ref().is_some_and(Hosting::is_finished)
205 || self
206 .parent
207 .as_ref()
208 .and_then(|parent| parent.task.as_ref())
209 .is_some_and(tokio::task::JoinHandle::is_finished)
210 }
211
212 pub fn health(&self) -> crate::host::HealthStatus {
213 let hosting = self.hosting.as_ref().map(Hosting::health);
214 let parent_link_ready = match &self.parent {
215 None => true,
216 Some(parent) => parent.task.as_ref().is_some_and(|task| !task.is_finished()),
217 };
218 crate::host::HealthStatus {
219 process_alive: true,
220 websocket_bound: hosting.is_some_and(|health| health.websocket_bound),
221 websocket_addr: hosting.and_then(|health| health.websocket_addr),
222 webtransport_bound: hosting.is_some_and(|health| health.webtransport_bound),
223 webtransport_addr: hosting.and_then(|health| health.webtransport_addr),
224 listeners_running: hosting.is_some_and(|health| health.listeners_running),
225 parent_link_ready,
226 child_link_ready: true,
227 }
228 }
229
230 pub fn websocket_addr(&self) -> Option<SocketAddr> {
231 self.hosting.as_ref().and_then(Hosting::websocket_addr)
232 }
233
234 pub fn webtransport_addr(&self) -> Option<SocketAddr> {
235 self.hosting.as_ref().and_then(Hosting::webtransport_addr)
236 }
237
238 pub async fn shutdown(self) -> Result<(), WsError> {
239 if let Some(parent) = &self.parent {
240 parent.cancel();
241 }
242 if let Some(hosting) = &self.hosting {
243 hosting.cancel();
244 }
245 let parent = async {
246 match self.parent {
247 Some(parent) => parent.shutdown().await,
248 None => Ok(()),
249 }
250 };
251 let hosting = async {
252 match self.hosting {
253 Some(hosting) => hosting
254 .shutdown()
255 .await
256 .map_err(|error| WsError::Connect(error.to_string())),
257 None => Ok(()),
258 }
259 };
260 let (parent, hosting) = tokio::join!(parent, hosting);
261 parent?;
262 hosting
263 }
264
265 pub async fn wait(&mut self) -> Result<(), WsError> {
266 let host_failure = |error: crate::host::HostError| WsError::Connect(error.to_string());
267 match (&mut self.hosting, &mut self.parent) {
268 (Some(hosting), Some(parent)) => tokio::select! {
269 biased;
270 result = hosting.wait() => result.map_err(host_failure),
271 result = parent.wait() => result,
272 },
273 (Some(hosting), None) => hosting.wait().await.map_err(host_failure),
274 (None, Some(parent)) => parent.wait().await,
275 (None, None) => Ok(()),
276 }
277 }
278}
279
280impl Node {
281 pub async fn start_topology(
282 self: &Arc<Self>,
283 config: TopologyConfig,
284 ) -> Result<UnbTopology, WsError> {
285 config.validate()?;
286 let hosting = match config.host {
287 Some(host) => Some(
288 host.start(self)
289 .await
290 .map_err(|error| WsError::Connect(error.to_string()))?,
291 ),
292 None => None,
293 };
294 let parent = match config.parent {
295 Some(parent) => {
296 let connected = async {
297 let pipe = connect_unix(&parent.path).await?;
298 TopologySession::connect(self, &parent.node, pipe).await
299 }
300 .await;
301 match connected {
302 Ok(connection) => {
303 let cancellation = self.cancellation().child_token();
304 let task_cancellation = cancellation.clone();
305 let node = self.clone();
306 let task = tokio::spawn(async move {
307 supervise_parent(node, parent, connection, task_cancellation).await
308 });
309 Some(ParentSupervisor {
310 cancellation,
311 task: Some(task),
312 })
313 }
314 Err(error) => {
315 if let Some(hosting) = hosting {
316 let _ = hosting.shutdown().await;
317 }
318 return Err(error);
319 }
320 }
321 }
322 None => None,
323 };
324 Ok(UnbTopology { hosting, parent })
325 }
326}
327
328async fn supervise_parent(
329 node: Arc<Node>,
330 parent: ParentLink,
331 mut connection: TopologySession,
332 cancellation: unb_runtime::CancellationToken,
333) -> Result<(), WsError> {
334 let mut delay = parent.reconnect.initial_delay;
335 loop {
336 tokio::select! {
337 biased;
338 () = cancellation.cancelled() => {
339 connection.shutdown();
340 connection.closed().await;
341 return Ok(());
342 }
343 () = connection.closed() => {}
344 }
345 loop {
346 tokio::select! {
347 biased;
348 () = cancellation.cancelled() => return Ok(()),
349 () = tokio::time::sleep(delay) => {}
350 }
351 let connected = tokio::select! {
352 biased;
353 () = cancellation.cancelled() => return Ok(()),
354 result = async {
355 let pipe = connect_unix(&parent.path).await?;
356 TopologySession::connect(&node, &parent.node, pipe).await
357 } => result,
358 };
359 match connected {
360 Ok(next) => {
361 connection = next;
362 delay = parent.reconnect.initial_delay;
363 break;
364 }
365 Err(_) => {
366 delay = delay.saturating_mul(2).min(parent.reconnect.max_delay);
367 }
368 }
369 }
370 }
371}
372
373pub struct UnixHosting {
374 cancellation: unb_runtime::CancellationToken,
375 task: tokio::task::JoinHandle<Result<(), WsError>>,
376}
377
378impl UnixHosting {
379 pub fn is_finished(&self) -> bool {
380 self.task.is_finished()
381 }
382
383 pub fn cancel(&self) {
384 self.cancellation.cancel();
385 }
386
387 pub async fn wait(&mut self) -> Result<(), WsError> {
388 (&mut self.task)
389 .await
390 .map_err(|error| WsError::Connect(error.to_string()))?
391 }
392
393 pub async fn shutdown(self) -> Result<(), WsError> {
394 self.cancellation.cancel();
395 self.task
396 .await
397 .map_err(|error| WsError::Connect(error.to_string()))?
398 }
399}
400
401impl Node {
402 pub fn host_unix_child(
403 self: &Arc<Self>,
404 path: impl AsRef<std::path::Path>,
405 expected_child: impl Into<String>,
406 ) -> Result<UnixHosting, WsError> {
407 let expected_child = expected_child.into();
408 validate_node_identifier(&expected_child)
409 .map_err(|error| WsError::Connect(error.to_string()))?;
410 let listener = unb_transport::unix::UnixListener::bind(path)?;
411 let cancellation = self.cancellation().child_token();
412 let task_cancellation = cancellation.clone();
413 let node = self.clone();
414 let task = tokio::spawn(async move {
415 let mut connections = JoinSet::new();
416 loop {
417 tokio::select! {
418 biased;
419 () = task_cancellation.cancelled() => break,
420 result = connections.join_next(), if !connections.is_empty() => {
421 result.expect("connection task exists").map_err(|error| WsError::Connect(error.to_string()))?;
422 }
423 pipe = accept_unix(&listener) => {
424 let pipe = pipe?;
425 let connection = tokio::select! {
426 biased;
427 () = task_cancellation.cancelled() => break,
428 connection = node.connect_transport(&expected_child, pipe) => connection,
429 };
430 if let Ok(connection) = connection {
431 let cancellation = task_cancellation.clone();
432 connections.spawn(async move {
433 tokio::select! {
434 biased;
435 () = cancellation.cancelled() => {
436 connection.shutdown();
437 connection.closed().await;
438 }
439 () = connection.closed() => {}
440 }
441 });
442 }
443 }
444 }
445 }
446 while let Some(result) = connections.join_next().await {
447 result.map_err(|error| WsError::Connect(error.to_string()))?;
448 }
449 Ok(())
450 });
451 Ok(UnixHosting { cancellation, task })
452 }
453
454 pub fn host_unix(self: &Arc<Self>, listener: unb_transport::unix::UnixListener) -> UnixHosting {
455 let cancellation = self.cancellation().child_token();
456 let task_cancellation = cancellation.clone();
457 let node = self.clone();
458 let task = tokio::spawn(async move {
459 let mut connections = JoinSet::new();
460 loop {
461 tokio::select! {
462 biased;
463 () = task_cancellation.cancelled() => break,
464 result = connections.join_next(), if !connections.is_empty() => {
465 result.expect("connection task exists").map_err(|error| WsError::Connect(error.to_string()))?;
466 }
467 pipe = accept_unix(&listener) => {
468 let pipe = pipe?;
469 let connection = tokio::select! {
470 biased;
471 () = task_cancellation.cancelled() => break,
472 connection = node.serve_transport(pipe) => connection,
473 };
474 let cancellation = task_cancellation.clone();
475 connections.spawn(async move {
476 tokio::select! {
477 biased;
478 () = cancellation.cancelled() => {
479 connection.shutdown();
480 connection.closed().await;
481 }
482 () = connection.closed() => {}
483 }
484 });
485 }
486 }
487 }
488 while let Some(result) = connections.join_next().await {
489 result.map_err(|error| WsError::Connect(error.to_string()))?;
490 }
491 Ok(())
492 });
493 UnixHosting { cancellation, task }
494 }
495}