Skip to main content

seyd_transport/
relay.rs

1// Copyright 2026 Anton Gravestam
2// SPDX-License-Identifier: Apache-2.0
3//! The cloud relay backend of [`Session`] (ADR 0010).
4//!
5//! When no candidate connected, the pilot may ask the signal server to carry
6//! the session: each party opens one WebSocket to the relay, the cloud pairs
7//! them by session id and forwards binary frames verbatim. The frames are the
8//! very bytes that would have been QUIC datagrams and control-stream bytes:
9//!
10//! | byte 0 | rest |
11//! |---|---|
12//! | `1` | one datagram (a chunk, exactly as `send_datagram` would have sent it) |
13//! | `2` | a segment of the control stream (NDJSON bytes, any split) |
14//!
15//! Nothing above the transport knows the difference: the engine's admission
16//! control still measures a queued backlog against the profile's threshold,
17//! and the pilot's reassembler still assembles chunks. The relay is TCP, so
18//! loss is retransmitted rather than repaired by FEC, and a congested leg
19//! shows up as backlog here (`send_buffer_queued`) instead of as QUIC loss.
20
21use crate::session::{ControlReader, ControlWriter, PathStats, SendError, Session};
22use bytes::{Bytes, BytesMut};
23use futures_util::{SinkExt, StreamExt};
24use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
25use std::sync::{Arc, Mutex};
26use std::time::{Duration, Instant};
27use tokio::io::{AsyncReadExt, AsyncWriteExt};
28use tokio::sync::{mpsc, oneshot, watch};
29use tokio_tungstenite::tungstenite::Message;
30
31pub const KIND_DATAGRAM: u8 = 1;
32pub const KIND_CONTROL: u8 = 2;
33
34/// How long the cloud gets to pair us with the pilot.
35const ATTACH_TIMEOUT: Duration = Duration::from_secs(10);
36/// A relayed datagram may be as large as a WebSocket frame; the engine's
37/// chunk length is what actually bounds it.
38const MAX_DATAGRAM: usize = 60_000;
39
40enum Out {
41    /// A complete frame (kind byte included) and, for datagrams, the payload
42    /// bytes to release from the queued count once it has been written.
43    Frame(Bytes, usize),
44    Close(String),
45}
46
47#[derive(Default)]
48struct RttSample {
49    rtt_ms: f64,
50    min_rtt_ms: f64,
51}
52
53struct Rate {
54    at: Instant,
55    bytes: u64,
56}
57
58pub struct Relay {
59    out: mpsc::UnboundedSender<Out>,
60    /// Datagram payload bytes handed to the writer task and not yet written.
61    queued: Arc<AtomicUsize>,
62    capacity: usize,
63    bytes_tx: Arc<AtomicU64>,
64    rtt: Arc<Mutex<RttSample>>,
65    rate: Mutex<Rate>,
66    closed: watch::Receiver<bool>,
67}
68
69impl Relay {
70    pub fn max_datagram_size(&self) -> usize {
71        MAX_DATAGRAM
72    }
73
74    pub fn send_buffer_space(&self) -> usize {
75        self.capacity
76            .saturating_sub(self.queued.load(Ordering::Relaxed))
77    }
78
79    pub fn send_buffer_queued(&self) -> usize {
80        self.queued.load(Ordering::Relaxed)
81    }
82
83    pub fn send_datagram(&self, payload: Bytes) -> Result<(), SendError> {
84        if *self.closed.borrow() {
85            return Err(SendError::Closed);
86        }
87        if payload.len() > MAX_DATAGRAM {
88            return Err(SendError::TooLarge);
89        }
90        // Reserve before queueing so two frames in flight cannot both pass the
91        // check against the same free space.
92        let len = payload.len();
93        let queued = self.queued.fetch_add(len, Ordering::Relaxed);
94        if queued + len > self.capacity {
95            self.queued.fetch_sub(len, Ordering::Relaxed);
96            return Err(SendError::Blocked);
97        }
98        let mut frame = BytesMut::with_capacity(len + 1);
99        frame.extend_from_slice(&[KIND_DATAGRAM]);
100        frame.extend_from_slice(&payload);
101        if self.out.send(Out::Frame(frame.freeze(), len)).is_err() {
102            self.queued.fetch_sub(len, Ordering::Relaxed);
103            return Err(SendError::Closed);
104        }
105        Ok(())
106    }
107
108    pub fn stats(&self) -> PathStats {
109        let rtt = self.rtt.lock().unwrap();
110        let now = Instant::now();
111        let mut rate = self.rate.lock().unwrap();
112        let bytes = self.bytes_tx.load(Ordering::Relaxed);
113        let dt = now.duration_since(rate.at).as_secs_f64();
114        let delivery_kbps = if dt > 0.05 {
115            ((bytes.saturating_sub(rate.bytes)) as f64 * 8.0 / dt / 1000.0) as u64
116        } else {
117            0
118        };
119        if dt > 0.05 {
120            rate.at = now;
121            rate.bytes = bytes;
122        }
123        PathStats {
124            rtt_ms: rtt.rtt_ms,
125            min_rtt_ms: rtt.min_rtt_ms,
126            cwnd: 0,
127            delivery_kbps,
128            lost_packets: 0,
129            current_mtu: 0,
130        }
131    }
132
133    pub fn close(&self, reason: &str) {
134        let _ = self.out.send(Out::Close(reason.to_string()));
135    }
136
137    pub async fn closed(&self) {
138        let mut rx = self.closed.clone();
139        while !*rx.borrow() {
140            if rx.changed().await.is_err() {
141                return;
142            }
143        }
144    }
145}
146
147/// Dial the relay, attach to `session_id`, and wait until the cloud has
148/// paired us with the pilot.
149pub(crate) async fn connect(
150    id: u64,
151    url: &str,
152    session_id: &str,
153    token: &str,
154    capacity: usize,
155) -> anyhow::Result<Session> {
156    let (ws, _) = tokio_tungstenite::connect_async(url).await?;
157    let (mut sink, mut stream) = ws.split();
158    let attach = serde_json::json!({
159        "type": "relay-attach", "session_id": session_id, "token": token, "party": "robot",
160    });
161    sink.send(Message::Text(attach.to_string().into())).await?;
162    let reply = tokio::time::timeout(ATTACH_TIMEOUT, async {
163        loop {
164            match stream.next().await {
165                Some(Ok(Message::Text(t))) => {
166                    let v: serde_json::Value = serde_json::from_str(&t)?;
167                    return Ok::<_, anyhow::Error>(v);
168                }
169                Some(Ok(Message::Close(_))) | None => {
170                    anyhow::bail!("relay closed before attaching")
171                }
172                Some(Ok(_)) => continue,
173                Some(Err(e)) => return Err(e.into()),
174            }
175        }
176    })
177    .await
178    .map_err(|_| anyhow::anyhow!("relay attach timed out"))??;
179    match reply.get("type").and_then(|t| t.as_str()) {
180        Some("relay-attached") => {}
181        Some("denied") => anyhow::bail!(
182            "relay denied: {}",
183            reply
184                .get("reason")
185                .and_then(|r| r.as_str())
186                .unwrap_or("unknown")
187        ),
188        other => anyhow::bail!("unexpected relay reply {other:?}"),
189    }
190
191    let (out_tx, mut out_rx) = mpsc::unbounded_channel::<Out>();
192    let (dg_tx, dg_rx) = mpsc::channel::<Bytes>(1024);
193    let (ctl_tx, ctl_rx) = oneshot::channel::<(ControlReader, ControlWriter)>();
194    let (closed_tx, closed_rx) = watch::channel(false);
195    let queued = Arc::new(AtomicUsize::new(0));
196    let bytes_tx = Arc::new(AtomicU64::new(0));
197    let rtt = Arc::new(Mutex::new(RttSample::default()));
198
199    // The control stream is two byte pipes: what the engine writes is framed
200    // out as `KIND_CONTROL`; what arrives as `KIND_CONTROL` is fed to what
201    // the engine reads. The pilot speaks first, exactly as over QUIC.
202    let (ctl_in_engine, mut ctl_in_task) = tokio::io::duplex(64 * 1024);
203    let (ctl_out_engine, mut ctl_out_task) = tokio::io::duplex(64 * 1024);
204    let reader: ControlReader = tokio::io::BufReader::new(Box::new(ctl_in_engine));
205    let writer: ControlWriter = Box::new(ctl_out_engine);
206    let _ = ctl_tx.send((reader, writer));
207
208    let relay = Relay {
209        out: out_tx,
210        queued: queued.clone(),
211        capacity,
212        bytes_tx: bytes_tx.clone(),
213        rtt: rtt.clone(),
214        rate: Mutex::new(Rate {
215            at: Instant::now(),
216            bytes: 0,
217        }),
218        closed: closed_rx,
219    };
220
221    tokio::spawn(async move {
222        let mut ping = tokio::time::interval(Duration::from_secs(1));
223        ping.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
224        let epoch = Instant::now();
225        let mut ctl_buf = vec![0u8; 16 * 1024];
226        let mut dropped = 0u64;
227        let reason: String = loop {
228            tokio::select! {
229                out = out_rx.recv() => {
230                    match out {
231                        Some(Out::Frame(frame, release)) => {
232                            let len = frame.len() as u64;
233                            let res = sink.send(Message::Binary(frame)).await;
234                            if release > 0 {
235                                queued.fetch_sub(release, Ordering::Relaxed);
236                            }
237                            if res.is_err() {
238                                break "write failed".to_string();
239                            }
240                            bytes_tx.fetch_add(len, Ordering::Relaxed);
241                        }
242                        Some(Out::Close(reason)) => {
243                            let _ = sink.send(Message::Close(None)).await;
244                            break reason;
245                        }
246                        None => break "session dropped".to_string(),
247                    }
248                }
249                n = ctl_out_task.read(&mut ctl_buf) => {
250                    match n {
251                        Ok(0) | Err(_) => break "control writer closed".to_string(),
252                        Ok(n) => {
253                            let mut frame = BytesMut::with_capacity(n + 1);
254                            frame.extend_from_slice(&[KIND_CONTROL]);
255                            frame.extend_from_slice(&ctl_buf[..n]);
256                            if sink.send(Message::Binary(frame.freeze())).await.is_err() {
257                                break "write failed".to_string();
258                            }
259                        }
260                    }
261                }
262                msg = stream.next() => {
263                    match msg {
264                        Some(Ok(Message::Binary(b))) => {
265                            match b.first().copied() {
266                                Some(KIND_DATAGRAM) => {
267                                    if dg_tx.try_send(b.slice(1..)).is_err() {
268                                        dropped += 1;
269                                    }
270                                }
271                                Some(KIND_CONTROL) => {
272                                    let ok = ctl_in_task.write_all(&b[1..]).await.is_ok();
273                                    if !ok {
274                                        break "control reader closed".to_string();
275                                    }
276                                }
277                                _ => {}
278                            }
279                        }
280                        Some(Ok(Message::Pong(p))) => {
281                            if p.len() == 8 {
282                                let sent = u64::from_be_bytes(p[..8].try_into().unwrap());
283                                let now = epoch.elapsed().as_micros() as u64;
284                                let ms = now.saturating_sub(sent) as f64 / 1000.0;
285                                let mut r = rtt.lock().unwrap();
286                                r.rtt_ms = ms;
287                                if r.min_rtt_ms == 0.0 || ms < r.min_rtt_ms {
288                                    r.min_rtt_ms = ms;
289                                }
290                            }
291                        }
292                        Some(Ok(Message::Text(t))) => {
293                            // `relay-closed` and friends: the cloud ending the session.
294                            tracing::debug!(session = id, msg = %t, "relay text");
295                            if t.contains("relay-closed") || t.contains("denied") {
296                                break "relay closed".to_string();
297                            }
298                        }
299                        Some(Ok(Message::Close(_))) | None => break "relay socket closed".to_string(),
300                        Some(Ok(_)) => {}
301                        Some(Err(e)) => break format!("relay error: {e}"),
302                    }
303                }
304                _ = ping.tick() => {
305                    let ts = (epoch.elapsed().as_micros() as u64).to_be_bytes();
306                    if sink.send(Message::Ping(Bytes::copy_from_slice(&ts))).await.is_err() {
307                        break "ping failed".to_string();
308                    }
309                }
310            }
311        };
312        if dropped > 0 {
313            tracing::debug!(
314                session = id,
315                dropped,
316                "relayed datagrams dropped (receiver slow)"
317            );
318        }
319        tracing::info!(session = id, %reason, "relay session ended");
320        let _ = closed_tx.send(true);
321    });
322
323    Ok(Session::from_relay(id, relay, dg_rx, ctl_rx))
324}