pb_mapper_client/sdk/
handle.rs

1//! Handles for a running tunnel: [`Registration`] and [`Connection`].
2//!
3//! Both wrap the same `LiveTunnel` — a cancellation token, the worker's join
4//! handle, and the status channel it publishes to. A handle observes and stops
5//! its tunnel; it never drives the traffic itself.
6
7use std::sync::Mutex;
8use std::time::Duration;
9
10use tokio::sync::watch;
11use tokio::task::JoinHandle;
12use tokio_util::sync::CancellationToken;
13
14use super::{Error, Result, TunnelStatus};
15
16pub(crate) struct LiveTunnel {
17    shutdown: CancellationToken,
18    join: Mutex<Option<JoinHandle<()>>>,
19    status: watch::Receiver<TunnelStatus>,
20}
21
22impl LiveTunnel {
23    pub(crate) fn new(
24        shutdown: CancellationToken,
25        join: JoinHandle<()>,
26        status: watch::Receiver<TunnelStatus>,
27    ) -> Self {
28        Self {
29            shutdown,
30            join: Mutex::new(Some(join)),
31            status,
32        }
33    }
34
35    pub(crate) fn status(&self) -> TunnelStatus {
36        self.status.borrow().clone()
37    }
38
39    pub(crate) fn subscribe(&self) -> watch::Receiver<TunnelStatus> {
40        self.status.clone()
41    }
42
43    pub(crate) async fn wait_ready(&self) -> Result<()> {
44        wait_for_connected(&mut self.status.clone()).await
45    }
46
47    pub(crate) async fn wait_ready_timeout(&self, timeout: Duration) -> Result<()> {
48        match tokio::time::timeout(timeout, self.wait_ready()).await {
49            Ok(result) => result,
50            Err(_) => Err(Error::ReadyTimeout { timeout }),
51        }
52    }
53
54    pub(crate) async fn stop(&self) -> Result<()> {
55        self.shutdown.cancel();
56        let handle = self
57            .join
58            .lock()
59            .unwrap_or_else(|poisoned| poisoned.into_inner())
60            .take();
61        if let Some(mut handle) = handle {
62            match tokio::time::timeout(Duration::from_secs(5), &mut handle).await {
63                Ok(Ok(())) => {}
64                Ok(Err(join_error)) if join_error.is_cancelled() => {}
65                Ok(Err(join_error)) => {
66                    return Err(Error::protocol(format!("tunnel task failed: {join_error}")));
67                }
68                Err(_) => {
69                    handle.abort();
70                    let _ = handle.await;
71                }
72            }
73        }
74        Ok(())
75    }
76}
77
78impl Drop for LiveTunnel {
79    fn drop(&mut self) {
80        self.shutdown.cancel();
81        if let Some(handle) = self
82            .join
83            .lock()
84            .unwrap_or_else(|poisoned| poisoned.into_inner())
85            .take()
86        {
87            handle.abort();
88        }
89    }
90}
91
92async fn wait_for_connected(status: &mut watch::Receiver<TunnelStatus>) -> Result<()> {
93    loop {
94        match status.borrow().clone() {
95            TunnelStatus::Connected => return Ok(()),
96            TunnelStatus::Failed(reason) => return Err(Error::TunnelFailed { reason }),
97            TunnelStatus::Stopped => return Err(Error::Stopped),
98            TunnelStatus::Starting | TunnelStatus::Retrying => {}
99        }
100        status.changed().await.map_err(|_| Error::Stopped)?;
101    }
102}
103
104/// Declares one end of a live tunnel as a public handle.
105///
106/// [`Registration`] and [`Connection`] share a lifecycle down to the last method:
107/// the same status, the same readiness wait, the same stop. They stay distinct
108/// types so neither can be passed where the other is meant, and this macro is
109/// what keeps the two from drifting apart.
110macro_rules! tunnel_handle {
111    ($(#[$doc:meta])* $name:ident) => {
112        $(#[$doc])*
113        pub struct $name {
114            inner: LiveTunnel,
115            key: String,
116        }
117
118        impl $name {
119            pub(crate) fn new(inner: LiveTunnel, key: String) -> Self {
120                Self { inner, key }
121            }
122
123            /// The service key this tunnel is bound to.
124            pub fn key(&self) -> &str {
125                &self.key
126            }
127
128            /// The latest status its worker reported.
129            pub fn status(&self) -> TunnelStatus {
130                self.inner.status()
131            }
132
133            /// Subscribe to status changes. Useful for N-API event bridges.
134            pub fn subscribe(&self) -> watch::Receiver<TunnelStatus> {
135                self.inner.subscribe()
136            }
137
138            /// Resolve once the tunnel is connected, or fails, or stops.
139            pub async fn wait_ready(&self) -> Result<()> {
140                self.inner.wait_ready().await
141            }
142
143            /// [`Self::wait_ready`], bounded by `timeout`.
144            pub async fn wait_ready_timeout(&self, timeout: Duration) -> Result<()> {
145                self.inner.wait_ready_timeout(timeout).await
146            }
147
148            /// Cancel the worker and wait for it to unwind.
149            pub async fn stop(&self) -> Result<()> {
150                self.inner.stop().await
151            }
152        }
153    };
154}
155
156tunnel_handle!(
157    /// A live `register` tunnel: a local service published on the relay.
158    Registration
159);
160
161tunnel_handle!(
162    /// A live `connect` tunnel: a local listener forwarding to a registered service.
163    Connection
164);
pb_mapper_client/sdk/handle.rs

pb_mapper_client/sdk/
handle.rs
