From f7fa6407306a26e934782db4325ef7ef1cd7dc2f Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Fri, 25 Sep 2026 12:21:53 -0700 Subject: [PATCH 1/2] quest(drain): claim drain-exit Co-Authored-By: Claude Opus 5.5 From c20dd4bc04638246341806a7f8f4923668e78fe8 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Fri, 25 Sep 2026 12:59:24 -0700 Subject: [PATCH 2/2] feat(relay): end a drain once every session has left Relay::run no longer sleeps out the whole drain window. It returns as soon as every established session has left (the deadline force-closes the rest) and logs which bound ended the drain, with the number of sessions the deadline force-closed. /metrics gains moq_relay_draining_sessions. The relay's own outbound cluster sessions are not counted: nothing sends them a GOAWAY, and the peer sees them close when the process exits. Co-Authored-By: Claude Opus 5.5 --- doc/bin/relay/config.md | 10 +- doc/bin/relay/http.md | 4 +- doc/bin/relay/index.md | 5 +- quest/m1/drain/README.md | 11 +- quest/m1/drain/drain-exit.md | 19 ---- rs/moq-relay/src/config.rs | 5 +- rs/moq-relay/src/connection.rs | 3 +- rs/moq-relay/src/internal.rs | 35 ++++++- rs/moq-relay/src/relay.rs | 58 ++++++++--- rs/moq-relay/src/shutdown.rs | 145 ++++++++++++++++++++++++-- rs/moq-relay/src/websocket.rs | 1 + rs/moq-relay/tests/shutdown_signal.rs | 101 +++++++++++++++--- 12 files changed, 322 insertions(+), 75 deletions(-) delete mode 100644 quest/m1/drain/drain-exit.md diff --git a/doc/bin/relay/config.md b/doc/bin/relay/config.md index c2588abc45..4b9680d03f 100644 --- a/doc/bin/relay/config.md +++ b/doc/bin/relay/config.md @@ -231,8 +231,14 @@ The first SIGTERM or SIGINT starts a drain: every session is sent a GOAWAY asking it to reconnect, and is force-closed if it is still connected when the window ends. A session that connects during the drain, such as a client with a cached DNS answer, is sent a GOAWAY immediately, with only the time left in -the window. The relay exits one second after the window ends, or immediately -on a second signal. `0` skips the GOAWAY and closes every session at once. +the window. The relay exits as soon as every session has left, when the window +ends, or immediately on a second signal. `0` skips the GOAWAY and closes every +session at once. + +The exit is logged with how long the drain took, as either +`drain complete: every session left` or `drain deadline force-closed sessions` +with the number `forced`. A session still in its handshake when the last one +leaves is not waited for. Only moq-lite-04+ and moq-transport clients act on a GOAWAY; older ones are closed when the window ends. An embedder can take over the signals and start the drain itself; see [Embed](/bin/relay/#embed). diff --git a/doc/bin/relay/http.md b/doc/bin/relay/http.md index 28cad8de5c..1ace5b461c 100644 --- a/doc/bin/relay/http.md +++ b/doc/bin/relay/http.md @@ -52,7 +52,9 @@ split by `tier` and `role`, plus accept-loop counters per TCP listener. Alert on `moq_relay_accept_failures_total{class="exhausted"}`, which means the process ran out of a resource `accept` needs. Content dropped for drifting past a subscriber's budget is counted separately as `moq_relay_stale_bytes_total` -and friends. Host CPU and memory belong to a node exporter. +and friends. During a [shutdown drain](/bin/relay/config#shutdown), +`moq_relay_draining_sessions` counts the sessions sent a GOAWAY that have not +left yet. Host CPU and memory belong to a node exporter. With `--runtime-io-uring`, each QUIC worker thread also reports its own `moq_relay_uring_*` counters under a `worker` label: datagrams and syscalls diff --git a/doc/bin/relay/index.md b/doc/bin/relay/index.md index 9e6393d557..6c18cf5186 100644 --- a/doc/bin/relay/index.md +++ b/doc/bin/relay/index.md @@ -78,8 +78,9 @@ TLS, and a certificate fingerprint for client pinning. The accessors borrow and `run` consumes the relay, so clone `cluster`, `auth`, `client`, `stats`, `shutdown`, and `shutdown_trigger` for application tasks before calling it. `trigger.start()` drains every session with a GOAWAY, -including any that connect afterwards, and `run` returns once the drain window -elapses, with the listeners released and the workers joined. `run` also starts +including any that connect afterwards, and `run` returns once every session +has left or the drain window elapses, with the listeners released and the +workers joined. `run` also starts the drain on SIGTERM or SIGINT. An application that owns those signals, for example to withdraw the node from DNS and wait out the TTL before draining, calls `with_signals(false)` and fires the trigger itself. Build routes from `web().routes()` (or diff --git a/quest/m1/drain/README.md b/quest/m1/drain/README.md index c8e0a45e04..c5396c5d24 100644 --- a/quest/m1/drain/README.md +++ b/quest/m1/drain/README.md @@ -1,4 +1,4 @@ -# Graceful relay drains (GOAWAY) +# [M] Graceful relay drains (GOAWAY) ## Goal @@ -30,7 +30,9 @@ is moq.pro's (downstream) fleet drain work, which consumes these quests. The relay's drain hook has landed: `Relay::with_signals(false)` hands SIGTERM to the embedder, and its `shutdown_trigger` GOAWAYs every session, arrivals -included, against one deadline. +included, against one deadline. `Relay::run` returns as soon as every session +has left, logging whether the deadline force-closed any, and +`moq_relay_draining_sessions` shows the drain's progress. **Clients (landed).** The JS reconnector migrates like the Rust one, preserving the app-visible session while resolving DNS again before dialing. @@ -43,11 +45,6 @@ by the stop deadline and encoder reconnect. track through an in-tree relay drained with the drain hook migrates to a second relay behind the same name without a dropped group. -## Quests - -- [Drain exit](/quest/m1/drain/drain-exit.md) - a drain ends as soon as every - session has left, and reports whether that or the deadline ended it - ## Related - [pop-skipping](/quest/m1/pop-skipping/README.md) - its same-PoP link price and full eligible pairing become important when a deployment adds a second relay per PoP diff --git a/quest/m1/drain/drain-exit.md b/quest/m1/drain/drain-exit.md deleted file mode 100644 index 5917a476dd..0000000000 --- a/quest/m1/drain/drain-exit.md +++ /dev/null @@ -1,19 +0,0 @@ -# [S] Drain exit - -## Goal - -`Relay::run` returns as soon as every session has left a drain, instead of -always waiting out the window, and the relay reports which bound ended it: -every session left, or the deadline force-closed the stragglers (and how -many). An orchestrator bounding its stop time can then prove from the log and -`/metrics` which one it hit. - -## Plan - -Today `drain` in `rs/moq-relay/src/relay.rs` sleeps the whole window plus a -second whatever the sessions do. Counting only needs the sessions that go -through `shutdown::Observer::drain_session` (QUIC on either runtime, and -WebSocket), not the `session::Registry`, which skips LAN peers. - -Open question: whether the relay's own outbound cluster sessions count, since -they are not drained by GOAWAY at all. diff --git a/rs/moq-relay/src/config.rs b/rs/moq-relay/src/config.rs index 3fca841c59..69ef3ccc4a 100644 --- a/rs/moq-relay/src/config.rs +++ b/rs/moq-relay/src/config.rs @@ -91,8 +91,9 @@ pub struct Config { /// How long accepted sessions may keep running after a shutdown signal, e.g. /// "10s" or "500ms". The first signal sends every session a GOAWAY and waits - /// this long for clients to reconnect elsewhere before force-closing them; a - /// second signal exits immediately. Zero closes them at once, with no GOAWAY + /// up to this long for clients to reconnect elsewhere before force-closing + /// them, exiting as soon as they have all left; a second signal exits + /// immediately. Zero closes them at once, with no GOAWAY /// they would have no time to act on. Defaults to 10 seconds. #[usage(skip)] #[serde(with = "crate::duration::serde_duration")] diff --git a/rs/moq-relay/src/connection.rs b/rs/moq-relay/src/connection.rs index 8d5e460c84..5faa19dbea 100644 --- a/rs/moq-relay/src/connection.rs +++ b/rs/moq-relay/src/connection.rs @@ -279,7 +279,7 @@ pub(crate) fn authorize( /// the session ([`auth::Lease::ended`]) the session closes with the reason, and /// the session's own close is reported back through the lease as the `end` event. /// Either way, a relay shutdown drains the session with a GOAWAY instead of -/// cutting it off. +/// cutting it off, and does not exit before this returns or the drain deadline. /// /// The session handle is `Send + Sync` whatever transport carries it, so this /// runs on the shared runtime even for sessions a pinned QUIC worker drives. @@ -289,6 +289,7 @@ pub async fn supervise( mut shutdown: crate::shutdown::Observer, registration: Option, ) -> anyhow::Result<()> { + let _serving = shutdown.serve(); loop { let nudged = async { match ®istration { diff --git a/rs/moq-relay/src/internal.rs b/rs/moq-relay/src/internal.rs index c3e576a8a8..493fa77472 100644 --- a/rs/moq-relay/src/internal.rs +++ b/rs/moq-relay/src/internal.rs @@ -9,7 +9,8 @@ //! - `/metrics` - this node's own traffic counters as Prometheus text //! exposition, plus the accept-loop health of its TCP listeners //! ([`with_listeners`](Internal::with_listeners)) and the per-worker health of -//! its io_uring runtime ([`with_uring`](Internal::with_uring)). A distinct plane +//! its io_uring runtime ([`with_uring`](Internal::with_uring)), and the +//! progress of a shutdown drain ([`with_shutdown`](Internal::with_shutdown)). A distinct plane //! from both the customer `web` surface and the MoQ `.stats` broadcast: the same //! atomics, but a different transport and audience (an ops scraper, not a //! customer or the dashboard/billing aggregators). The runtime counters are @@ -89,6 +90,7 @@ pub struct Internal { health: moq_tokio::accept::Health, listeners: Vec, uring: Vec, + shutdown: Option, } #[derive(Clone)] @@ -98,6 +100,7 @@ struct InternalState { sessions: crate::session::Registry, listeners: Vec, uring: Vec, + shutdown: Option, } impl Internal { @@ -122,6 +125,7 @@ impl Internal { health, listeners, uring: Vec::new(), + shutdown: None, } } @@ -172,6 +176,12 @@ impl Internal { self } + /// Report the sessions a shutdown drain is still waiting on at `/metrics`. + pub fn with_shutdown(mut self, shutdown: crate::shutdown::Observer) -> Self { + self.shutdown = Some(shutdown); + self + } + /// Attach the relay cluster used to serve the `/nodes` topology snapshot. pub fn with_cluster(mut self, cluster: &crate::cluster::Cluster) -> Self { self.nodes = Some(cluster.nodes.clone()); @@ -205,6 +215,7 @@ impl Internal { sessions: self.sessions.clone(), listeners: self.listeners.clone(), uring: self.uring.clone(), + shutdown: self.shutdown.clone(), }) } @@ -275,7 +286,10 @@ async fn serve_health() -> Response { /// current cumulative snapshot; a downstream scraper derives rates and live /// counts (`open - closed`). async fn serve_metrics(State(state): State) -> Response { - let body = render_metrics(&state.stats.snapshot(), &state.listeners, &state.uring); + let mut body = render_metrics(&state.stats.snapshot(), &state.listeners, &state.uring); + if let Some(shutdown) = &state.shutdown { + render_drain(&mut body, shutdown.tally()); + } ([(http::header::CONTENT_TYPE, "text/plain; version=0.0.4")], body).into_response() } @@ -454,6 +468,21 @@ fn render_metrics( out } +/// The sessions a shutdown drain is still waiting on: 0 until the drain starts, +/// and back to 0 when every session has left, which is when the relay exits. +/// A scrape that last saw it above 0 shortly before the deadline means the +/// deadline force-closed the rest; the exit log records how many. +fn render_drain(out: &mut String, tally: crate::shutdown::Tally) { + use std::fmt::Write as _; + + let _ = writeln!( + out, + "# HELP moq_relay_draining_sessions Sessions sent a shutdown GOAWAY that have not left yet." + ); + let _ = writeln!(out, "# TYPE moq_relay_draining_sessions gauge"); + let _ = writeln!(out, "moq_relay_draining_sessions {}", tally.draining); +} + /// The accept-loop health of every listener on the node. /// /// The counters are the load-bearing half: a process out of descriptors cannot @@ -863,6 +892,7 @@ mod tests { sessions: crate::session::Registry::new(), listeners: Vec::new(), uring: Vec::new(), + shutdown: None, }; let Json(snapshot) = serve_nodes(State(state)).await; @@ -880,6 +910,7 @@ mod tests { sessions: crate::session::Registry::new(), listeners: Vec::new(), uring: Vec::new(), + shutdown: None, }; let Json(snapshot) = serve_nodes(State(state)).await; diff --git a/rs/moq-relay/src/relay.rs b/rs/moq-relay/src/relay.rs index 953111b772..3b62122cbe 100644 --- a/rs/moq-relay/src/relay.rs +++ b/rs/moq-relay/src/relay.rs @@ -268,7 +268,8 @@ impl Relay { let cluster = cluster.with_stats(stats.clone()); // Graceful shutdown: the first signal drains every accepted session with a - // GOAWAY; a second signal (or the drain window elapsing) exits. + // GOAWAY; the relay exits once they have all left, at the drain deadline, + // or on a second signal. let (shutdown_trigger, shutdown) = shutdown::Observer::new(drain_timeout); let sessions = crate::session::Registry::new(); let (ready, _) = tokio::sync::watch::channel(false); @@ -286,6 +287,7 @@ impl Relay { let internal = internal::Internal::new(config.internal, cluster.stats.clone()) .with_cluster(&cluster) .with_sessions(sessions.clone()) + .with_shutdown(shutdown.clone()) .with_listeners(web.accept_health()) .with_listeners(server.accept_health()); // Bound but not yet serving: registering here (rather than after the @@ -394,8 +396,9 @@ impl Relay { } /// Starts graceful shutdown: every session, including any accepted - /// afterwards, drains with a GOAWAY and [`Self::run`] returns once the drain - /// window elapses. Clone it before `run` consumes the relay. + /// afterwards, drains with a GOAWAY and [`Self::run`] returns once every + /// session has left or the drain window elapses. Clone it before `run` + /// consumes the relay. pub fn shutdown_trigger(&self) -> &shutdown::Trigger { &self.shutdown_trigger } @@ -462,9 +465,10 @@ impl Relay { /// Serve until something fails or shutdown completes: accept sessions, run /// the cluster, and serve both HTTP surfaces. Notifies systemd once - /// everything is up. Returns once the drain window elapses after a signal - /// (see [`Self::with_signals`]) or [`shutdown::Trigger::start`], with every - /// listener released and every worker joined. + /// everything is up. Returns once a drain started by a signal (see + /// [`Self::with_signals`]) or [`shutdown::Trigger::start`] ends, as soon as + /// every session has left or at the drain deadline, with every listener + /// released and every worker joined. /// /// This is also the embedding loop. Extra routes go on via [`Self::with_web`] /// / [`Self::with_internal`] before calling this; cloned handles outlive it. @@ -656,9 +660,10 @@ impl Relay { /// Two-stage shutdown: the first signal, or an embedder firing /// [`shutdown::Trigger::start`], starts the drain broadcast (every session sends -/// GOAWAY and waits for its peer to leave); a second signal, or that recorded -/// deadline plus one second, returns from [`Relay::run`]. Without `signals` -/// only the trigger and that deadline count. +/// GOAWAY and waits for its peer to leave). Returns from [`Relay::run`] once +/// every session has left, which the drain deadline forces, or on a second +/// signal, logging which ended it. +/// Without `signals` only the trigger and the sessions count. async fn drain(trigger: shutdown::Trigger, mut shutdown: shutdown::Observer, signals: bool) -> anyhow::Result<()> { let window = shutdown.drain_timeout; let signal = || async move { @@ -679,19 +684,38 @@ async fn drain(trigger: shutdown::Trigger, mut shutdown: shutdown::Observer, sig _ = shutdown.started() => tracing::info!(?window, "shutdown requested; draining sessions"), } - // One extra second past the deadline fixed when the trigger fired, so - // per-session force-closes fire first. That instant may be earlier than - // this future was polled (the embedder can start the drain during startup), - // and a fresh window here would keep the process up past the time sessions - // were told. + // The deadline fixed when the trigger fired, which may be earlier than this + // future was polled (the embedder can start the drain during startup); a + // fresh window here would keep the process up past the time sessions were + // told. Each session is force-closed at it, so `drained` resolves by then; + // the extra second only bounds a session whose close never completes. let deadline = shutdown.deadline().context("drain started without a deadline")?; - let grace = (deadline + std::time::Duration::from_secs(1)).saturating_duration_since(std::time::Instant::now()); tokio::select! { res = signal() => { res?; - tracing::warn!("second shutdown signal; exiting immediately"); + tracing::warn!(open = shutdown.tally().live, "second shutdown signal; exiting immediately"); + return Ok(()); + } + _ = shutdown.drained() => {} + _ = tokio::time::sleep_until(deadline + std::time::Duration::from_secs(1)) => {} + } + + let elapsed = tokio::time::Instant::now().saturating_duration_since(deadline - window); + match shutdown.tally() { + shutdown::Tally { live: 0, forced: 0, .. } => { + tracing::info!(?elapsed, "drain complete: every session left; exiting") + } + shutdown::Tally { live: 0, forced, .. } => { + tracing::warn!(?elapsed, forced, "drain deadline force-closed sessions; exiting") + } + shutdown::Tally { live, forced, .. } => { + tracing::warn!( + ?elapsed, + forced, + open = live, + "drain deadline passed with sessions still open; exiting" + ) } - _ = tokio::time::sleep(grace) => tracing::info!("drain window elapsed; exiting"), } Ok(()) } diff --git a/rs/moq-relay/src/shutdown.rs b/rs/moq-relay/src/shutdown.rs index a1987c104e..cf7db81675 100644 --- a/rs/moq-relay/src/shutdown.rs +++ b/rs/moq-relay/src/shutdown.rs @@ -5,10 +5,15 @@ //! The drain has one deadline, fixed when it starts. A session accepted after //! that (a cached DNS resolve, a pool alias) is sent a GOAWAY at once, carrying //! only the time left, so no session outlives the window the relay promised. +//! +//! The drain ends as soon as every session has left, or at that deadline. +//! Sessions are counted once established, so one still in +//! its handshake when the last established session leaves is cut off with the +//! process. -use std::time::{Duration, Instant}; +use std::{sync::Arc, time::Duration}; -use tokio::sync::watch; +use tokio::{sync::watch, time::Instant}; /// Fires the relay-wide shutdown broadcast. Held by `Relay::run` for the OS /// signal path; an embedder clones one to stop the relay from its own task. @@ -34,6 +39,17 @@ impl Trigger { } } +/// The sessions a drain is waiting on. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct Tally { + /// Established sessions still open, drained or not. + pub live: usize, + /// Sessions sent a GOAWAY that have not left yet. + pub draining: usize, + /// Sessions still open when their drain deadline passed, and so closed by it. + pub forced: usize, +} + /// A per-connection handle observing the relay-wide shutdown broadcast. /// /// Cheap to clone; each accepted session waits on [`started`](Self::started) @@ -41,6 +57,7 @@ impl Trigger { #[derive(Clone)] pub struct Observer { rx: watch::Receiver>, + tally: Arc>, /// How long a drained session may keep running before it is force-closed. pub drain_timeout: Duration, } @@ -49,7 +66,7 @@ impl Observer { /// Create the trigger and its observer half. pub fn new(drain_timeout: Duration) -> (Trigger, Self) { let (tx, rx) = watch::channel(None); - (Trigger { tx, drain_timeout }, Self { rx, drain_timeout }) + (Trigger { tx, drain_timeout }, Self::from(rx, drain_timeout)) } /// A handle that never fires, for callers without shutdown coordination @@ -59,9 +76,14 @@ impl Observer { // Leak-free: dropping the sender doesn't resolve `started` (it waits for // a deadline, not for channel closure). drop(tx); + Self::from(rx, crate::DEFAULT_DRAIN_TIMEOUT) + } + + fn from(rx: watch::Receiver>, drain_timeout: Duration) -> Self { Self { rx, - drain_timeout: crate::DEFAULT_DRAIN_TIMEOUT, + tally: Arc::new(watch::channel(Tally::default()).0), + drain_timeout, } } @@ -80,6 +102,36 @@ impl Observer { *self.rx.borrow() } + /// Count an established session until the returned guard drops, so the drain + /// waits for it to leave. + pub(crate) fn serve(&self) -> Serving { + self.tally.send_modify(|tally| tally.live += 1); + Serving { + tally: self.tally.clone(), + } + } + + /// The sessions counted so far. + pub(crate) fn tally(&self) -> Tally { + *self.tally.borrow() + } + + /// Resolve once no [`serve`](Self::serve) guard is left. + pub(crate) async fn drained(&self) { + // The sender lives in `self`, so the channel cannot close under the wait. + let _ = self.tally.subscribe().wait_for(|tally| tally.live == 0).await; + } + + /// Count a session as draining until the returned guard drops, and as forced + /// if it is still open at `deadline`. + fn draining(&self, deadline: Instant) -> Draining { + self.tally.send_modify(|tally| tally.draining += 1); + Draining { + tally: self.tally.clone(), + deadline, + } + } + /// Drain `session` with an empty-URI GOAWAY ("reconnect to me"), waiting for /// the peer to leave. /// @@ -93,10 +145,12 @@ impl Observer { /// With no time left (a zero [`drain_timeout`](Self::drain_timeout), or a /// session accepted after the deadline) the session is closed at once. pub async fn drain_session(&self, session: &moq_net::Session) { - let remaining = match *self.rx.borrow() { - Some(deadline) => deadline.saturating_duration_since(Instant::now()), - None => self.drain_timeout, - }; + let now = Instant::now(); + let deadline = self.deadline().unwrap_or(now + self.drain_timeout); + let remaining = deadline.saturating_duration_since(now); + // Dropped when this returns or is cancelled (a WebSocket driver ending + // first), so the count holds whichever way the session goes. + let _draining = self.draining(deadline); // No grace left, so there is nothing to drain. Close now rather than send a // GOAWAY: the peer would have no time to act on it, and a zero timeout means @@ -119,3 +173,78 @@ impl Observer { session.closed().await; } } + +/// An established session the drain waits on; see [`Observer::serve`]. +pub(crate) struct Serving { + tally: Arc>, +} + +impl Drop for Serving { + fn drop(&mut self) { + self.tally.send_modify(|tally| tally.live -= 1); + } +} + +/// A session sent a GOAWAY; see [`Observer::draining`]. +struct Draining { + tally: Arc>, + deadline: Instant, +} + +impl Drop for Draining { + fn drop(&mut self) { + // The driver arms its force-close after sending the GOAWAY, so it never + // fires before `deadline` and every session it closes is counted. A peer + // leaving in that sliver was still there at the deadline too. + let forced = Instant::now() >= self.deadline; + self.tally.send_modify(|tally| { + tally.draining -= 1; + tally.forced += usize::from(forced); + }); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test(start_paused = true)] + async fn drained_waits_for_every_session() { + let (trigger, shutdown) = Observer::new(Duration::from_secs(10)); + let first = shutdown.serve(); + let second = shutdown.serve(); + trigger.start(); + + let drained = tokio::spawn({ + let shutdown = shutdown.clone(); + async move { shutdown.drained().await } + }); + drop(first); + tokio::task::yield_now().await; + assert!(!drained.is_finished(), "drained with a session still open"); + + drop(second); + drained.await.expect("drained task"); + assert_eq!(shutdown.tally(), Tally::default()); + } + + #[tokio::test(start_paused = true)] + async fn a_session_open_at_the_deadline_is_forced() { + let window = Duration::from_secs(10); + let (trigger, shutdown) = Observer::new(window); + trigger.start(); + let deadline = shutdown.deadline().expect("started"); + + let left = shutdown.draining(deadline); + let straggler = shutdown.draining(deadline); + assert_eq!(shutdown.tally().draining, 2); + + drop(left); + tokio::time::sleep(window).await; + drop(straggler); + + let tally = shutdown.tally(); + assert_eq!(tally.draining, 0); + assert_eq!(tally.forced, 1, "only the session open at the deadline is forced"); + } +} diff --git a/rs/moq-relay/src/websocket.rs b/rs/moq-relay/src/websocket.rs index 6093d0ba95..b2ebba9428 100644 --- a/rs/moq-relay/src/websocket.rs +++ b/rs/moq-relay/src/websocket.rs @@ -197,6 +197,7 @@ where // The handshake is done, so this is a MoQ session now: only now can a push // be serviced, and only now does the session appear in the live table. let registration = pending.map(|(sessions, request)| sessions.register(request)); + let _serving = shutdown.serve(); loop { let nudged = async { diff --git a/rs/moq-relay/tests/shutdown_signal.rs b/rs/moq-relay/tests/shutdown_signal.rs index b23ae043ea..6e505cd3fa 100644 --- a/rs/moq-relay/tests/shutdown_signal.rs +++ b/rs/moq-relay/tests/shutdown_signal.rs @@ -17,6 +17,9 @@ //! itself; a session that still arrives mid-drain is sent a GOAWAY at once, //! carrying only what is left of the window. //! +//! The drain ends as soon as every session has left rather than waiting out the +//! window, which a stop-time budget depends on. +//! //! Each signal test raises or handles process signals, so they rely on //! nextest's process-per-test isolation. `a_trigger_before_run_keeps_the_deadline` //! fires the trigger before `run` instead of a signal. @@ -28,8 +31,8 @@ use std::{net::TcpListener, time::Duration}; use moq_relay::{Config, Relay, auth}; /// Long enough that "exited immediately" and "waited out the window" cannot be -/// confused, short enough to keep the test quick: `Relay::run` sleeps this plus -/// one second before exiting. +/// confused, short enough to keep the test quick: a session that never leaves +/// keeps `Relay::run` up this long. const DRAIN_TIMEOUT: Duration = Duration::from_secs(3); /// Run `test` on a current-thread runtime with a large stack. @@ -66,6 +69,11 @@ fn a_session_arriving_mid_drain_gets_what_is_left() { run_test(a_session_arriving_mid_drain_gets_what_is_left_inner); } +#[test] +fn a_drain_ends_once_every_session_leaves() { + run_test(a_drain_ends_once_every_session_leaves_inner); +} + #[test] fn a_trigger_before_run_keeps_the_deadline() { run_test(a_trigger_before_run_keeps_the_deadline_inner); @@ -88,6 +96,9 @@ async fn sigint_drains_sessions_before_exiting_inner() { let connection = connect(&client, port).await; let draining = connection.draining().expect("connected"); + // The one-shot client leaves on the GOAWAY; this peer is what keeps the drain + // open for its whole window. + let _straggler = straggler(port).await; let signalled = std::time::Instant::now(); // SAFETY: `raise` is async-signal-safe, and SIGINT's disposition is tokio's @@ -142,6 +153,7 @@ async fn an_embedder_owns_the_signals_inner() { let connection = connect(&client, port).await; let draining = connection.draining().expect("connected"); + let _straggler = straggler(port).await; // SAFETY: `raise` is async-signal-safe, and SIGINT's disposition is tokio's // handler, registered above. @@ -193,15 +205,15 @@ async fn a_session_arriving_mid_drain_gets_what_is_left_inner() { let client = client(vec!["moq-transport-17".parse().expect("parse version")]); let established = connect(&client, port).await; + let draining = established.draining().expect("connected"); + // Keeps the relay up for the arrival below once `established` leaves. + let _straggler = straggler(port).await; trigger.start(); let deadline = std::time::Instant::now() + DRAIN_TIMEOUT; - let goaway = tokio::time::timeout( - Duration::from_secs(5), - established.draining().expect("connected").recv(), - ) - .await - .expect("no GOAWAY within 5s of the trigger") - .expect("session closed without a GOAWAY"); + let goaway = tokio::time::timeout(Duration::from_secs(5), draining.recv()) + .await + .expect("no GOAWAY within 5s of the trigger") + .expect("session closed without a GOAWAY"); assert_eq!(goaway.uri(), "", "expected a reconnect-to-me GOAWAY"); let timeout = goaway.timeout().expect("the GOAWAY carries its deadline"); assert!( @@ -236,6 +248,39 @@ async fn a_session_arriving_mid_drain_gets_what_is_left_inner() { .expect("relay exited with an error"); } +async fn a_drain_ends_once_every_session_leaves_inner() { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + // A window the test would time out long before, so only an early exit passes. + let (port, mut config) = relay_config(); + config.drain_timeout = Duration::from_secs(600); + let internal = free_port(); + config.internal.listen = Some(format!("127.0.0.1:{internal}").parse().expect("parse addr")); + let relay = Relay::load(config).await.expect("load relay").with_signals(false); + let trigger = relay.shutdown_trigger().clone(); + let run = tokio::spawn(relay.run()); + wait_listening(port).await; + + let left = connect(&client(Vec::new()), port).await; + let straggler = straggler(port).await; + trigger.start(); + + // The one-shot client leaves on its GOAWAY; the straggler is still draining. + tokio::time::timeout(Duration::from_secs(5), left.closed()) + .await + .expect("the one-shot client did not leave on the GOAWAY") + .expect("the one-shot client failed"); + wait_metric(internal, "moq_relay_draining_sessions 1").await; + assert!(!run.is_finished(), "the relay exited with a session still draining"); + + drop(straggler); + tokio::time::timeout(Duration::from_secs(5), run) + .await + .expect("relay kept running after every session left") + .expect("relay task panicked") + .expect("relay exited with an error"); +} + async fn a_trigger_before_run_keeps_the_deadline_inner() { let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); @@ -275,20 +320,30 @@ fn client(version: Vec) -> moq_tokio::Client { .with_reconnect(false) } +/// A session that stays until the drain deadline closes it: moq-lite-03 has no +/// GOAWAY message, so the relay drains it without the peer ever knowing. +async fn straggler(port: u16) -> moq_tokio::Connection { + connect(&client(vec!["moq-lite-03".parse().expect("parse version")]), port).await +} + /// A session to the relay on `port`. async fn connect(client: &moq_tokio::Client, port: u16) -> moq_tokio::Connection { let url: url::Url = format!("tcp://127.0.0.1:{port}/").parse().expect("parse url"); client.connect(url).established().await.expect("connect") } +/// A free loopback TCP port. The listener is bound by `Relay::run`, not here, +/// so this leaves the usual probe/bind gap; on loopback it is not worth +/// retrying around. +fn free_port() -> u16 { + let probe = TcpListener::bind("127.0.0.1:0").expect("bind probe"); + probe.local_addr().expect("local addr").port() +} + /// A stream-only relay on a free loopback TCP port, fully public, with a short /// drain window. Returns the port and the config to hand [`Relay::load`]. fn relay_config() -> (u16, Config) { - // The listener is bound by `Relay::run`, not here, so this leaves the usual - // probe/bind gap; on loopback it is not worth retrying around. - let probe = TcpListener::bind("127.0.0.1:0").expect("bind probe"); - let port = probe.local_addr().expect("local addr").port(); - drop(probe); + let port = free_port(); // Fully public auth: any no-JWT stream client gets the whole root. let mut auth = auth::Config::default(); @@ -302,6 +357,24 @@ fn relay_config() -> (u16, Config) { (port, config) } +/// Wait until the internal listener on `port` reports `line` at `/metrics`. +async fn wait_metric(port: u16, line: &str) { + let deadline = std::time::Instant::now() + Duration::from_secs(5); + loop { + let metrics = reqwest::get(format!("http://127.0.0.1:{port}/metrics")) + .await + .expect("scrape metrics") + .text() + .await + .expect("read metrics"); + if metrics.lines().any(|l| l == line) { + break; + } + assert!(std::time::Instant::now() < deadline, "never saw {line} in:\n{metrics}"); + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + async fn wait_listening(port: u16) { let deadline = std::time::Instant::now() + Duration::from_secs(5); loop {