feat(relay): end a drain once every session has left

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).

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

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

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

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")]

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<crate::session::Registration>,
 ) -> anyhow::Result<()> {
+	let _serving = shutdown.serve();
 	loop {
 		let nudged = async {
 			match &registration {

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<moq_tokio::accept::Health>,
 	uring: Vec<UringWorker>,
+	shutdown: Option<crate::shutdown::Observer>,
 }
 
 #[derive(Clone)]
@@ -98,6 +100,7 @@ struct InternalState {
 	sessions: crate::session::Registry,
 	listeners: Vec<moq_tokio::accept::Health>,
 	uring: Vec<UringWorker>,
+	shutdown: Option<crate::shutdown::Observer>,
 }
 
 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<InternalState>) -> 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;

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(())
 }