From 2a85de2680f3028d307a17156e7514fc500424ae Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Tue, 18 Aug 2026 14:51:30 -0700 Subject: [PATCH 1/2] feat(relay): serve QUIC from pinned per-core workers Milestone 1 of the thread-per-core runtime plan (#2875), first half. The relay runs one work-stealing runtime over one UDP socket, so every packet can cross threads and every wakeup is a candidate context switch. `runtime.workers` switches QUIC to the opposite shape: N threads, each pinned to a core, each running a `current_thread` runtime and owning one socket in a `SO_REUSEPORT` group on the listen address. A connection lands on whichever worker the kernel steers its first packet to and stays there, so its QUIC driver, session, and tasks never leave that thread. Everything that is not QUIC (web, internal, tcp/unix, clustering, signals, cert reload) stays on the shared runtime, and the workers reach the rest of the relay through the shared Cluster. Off by default. The kernel picks a group member by hashing the packet's 4-tuple, which holds a connection to one worker only for as long as the client keeps its address: a NAT rebinding or network change hashes somewhere else and reaches a worker that has never seen the connection ID. Steering by connection ID is the second half of M1 and needs a CID codec plus a reuseport filter; until then this is a measurement knob, which is also why the shared runtime stays the default. moq-tokio gains the two primitives the relay drives: - `bind::Udp`, an options struct for `bind::udp`, carrying `with_reuse_port`. It fails with `Unsupported` off Linux rather than binding a group the platform will not balance: macOS and the BSDs accept `SO_REUSEPORT` and then feed a unicast flow to one member, so the workers would come up healthy with one of them serving everything. - `listen::Shard` and `listen::Listeners`, so one process can build a QUIC-only `Server` per group member while a stream-only `Server` keeps tcp/unix on its main runtime. `Listeners` also makes "no QUIC" expressible, which was previously only implied by configuring a stream listener. `--listen-tls-generate` is rejected with workers, since each would generate and serve a certificate of its own; the fingerprint endpoint reads worker 0's, and every worker loads the same files. Verified on Linux (container, kernel 6.19 aarch64), since neither reuseport load balancing nor the integration test compiles on macOS. Co-Authored-By: Claude Opus 5 --- Cargo.lock | 12 + doc/bin/relay/config.md | 31 +++ rs/moq-relay/Cargo.toml | 1 + rs/moq-relay/src/config.rs | 31 +++ rs/moq-relay/src/lib.rs | 2 + rs/moq-relay/src/relay.rs | 46 +++- rs/moq-relay/src/runtime.rs | 316 ++++++++++++++++++++++++++ rs/moq-relay/tests/runtime_workers.rs | 152 +++++++++++++ rs/moq-tokio/Cargo.toml | 2 +- rs/moq-tokio/src/bind.rs | 144 +++++++++++- rs/moq-tokio/src/listen.rs | 131 +++++++++++ rs/moq-tokio/src/noq.rs | 5 +- rs/moq-tokio/src/quiche.rs | 6 +- rs/moq-tokio/src/quinn.rs | 5 +- rs/moq-tokio/src/server.rs | 6 +- 15 files changed, 865 insertions(+), 25 deletions(-) create mode 100644 rs/moq-relay/src/runtime.rs create mode 100644 rs/moq-relay/tests/runtime_workers.rs diff --git a/Cargo.lock b/Cargo.lock index ac9e80f0a..d3fcac86e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1469,6 +1469,17 @@ dependencies = [ "libc", ] +[[package]] +name = "core_affinity" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a034b3a7b624016c6e13f5df875747cc25f884156aad2abd12b6c46797971342" +dependencies = [ + "libc", + "num_cpus", + "winapi", +] + [[package]] name = "coreaudio-rs" version = "0.14.2" @@ -4923,6 +4934,7 @@ dependencies = [ "bytes", "bytesize", "clap", + "core_affinity", "futures", "http-body", "http-cache-reqwest", diff --git a/doc/bin/relay/config.md b/doc/bin/relay/config.md index 24055bf89..408adb9f7 100644 --- a/doc/bin/relay/config.md +++ b/doc/bin/relay/config.md @@ -45,6 +45,37 @@ Logging configuration. level = "info" ``` +### \[runtime] + +How the relay lays its QUIC work out over threads. By default one work-stealing +runtime serves every connection off one UDP socket. + +```toml +[runtime] +# Serve QUIC from this many single-threaded workers instead of the shared +# runtime. Each worker is a thread with its own socket on the listen address +# (SO_REUSEPORT), and a connection is handled start to finish by one worker, so +# its packets never cross threads. Everything else (HTTP, WebSocket, tcp/unix, +# clustering) stays on the shared runtime. Omit to keep QUIC there too. +# +# Linux only, and incompatible with listen.tls.generate: each worker would +# generate a certificate of its own. Point at real certificate files instead. +# +# The kernel steers a packet to a worker by hashing its source address and port, +# so a client that changes address mid-connection (a NAT rebinding, a network +# change) lands on a worker that has never seen it and loses the connection. +# Treat this as a benchmarking knob until steering follows the connection ID. +workers = 8 + +# Pin each worker to a CPU core. Default: true. +pin = true +``` + +The shared runtime is still there for everything that is not QUIC, and it still +sizes its thread pool to the machine. Set `TOKIO_WORKER_THREADS` in the +environment to bound it, so the workers are not competing with a full second +pool for the same cores. + ### \[listen] QUIC/WebTransport server settings. Optionally add plaintext qmux stream diff --git a/rs/moq-relay/Cargo.toml b/rs/moq-relay/Cargo.toml index e0faa9c35..8b077bdaa 100644 --- a/rs/moq-relay/Cargo.toml +++ b/rs/moq-relay/Cargo.toml @@ -47,6 +47,7 @@ axum-server = { version = "0.8", features = ["tls-rustls"] } bytes = "1" bytesize = "2.4.2" clap = { version = "4", features = ["derive"] } +core_affinity = "0.8" futures = "0.3" http-body = "1" http-cache-reqwest = { version = "1.0.0-alpha.6", features = ["manager-moka", "url-standard"], default-features = false } diff --git a/rs/moq-relay/src/config.rs b/rs/moq-relay/src/config.rs index 6ea281cde..919f10fd3 100644 --- a/rs/moq-relay/src/config.rs +++ b/rs/moq-relay/src/config.rs @@ -35,6 +35,12 @@ pub struct Config { #[serde(default)] pub log: moq_tokio::Log, + /// How QUIC work is laid out over threads. One shared runtime unless + /// `runtime.workers` is set. + #[command(flatten)] + #[serde(default)] + pub runtime: crate::RuntimeConfig, + /// Cluster configuration. #[command(flatten)] #[serde(default)] @@ -369,6 +375,31 @@ depth = 2 /// Regression test for the clap+TOML clobber bug applied to `[cache]`. Both /// fields are `Option` so a TOML-configured cache size survives the /// CLI re-parse when no `--cache-*` flag is passed. + #[test] + fn cli_does_not_clobber_toml_runtime() { + let _env = EnvGuard::clear(&["MOQ_RUNTIME_WORKERS", "MOQ_RUNTIME_PIN"]); + + let toml = r#" +[runtime] +workers = 8 +pin = false +"#; + let dir = std::env::temp_dir().join("moq-relay-config-test"); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("runtime-toml-wins.toml"); + std::fs::write(&path, toml).unwrap(); + + let args = vec![std::ffi::OsString::from("moq-relay"), std::ffi::OsString::from(&path)]; + let config = Config::parse_and_merge(args).expect("config load"); + + assert_eq!(config.runtime.workers, Some(8)); + assert_eq!( + config.runtime.pin, + Some(false), + "TOML's runtime.pin=false must survive the CLI re-parse" + ); + } + #[test] fn cli_does_not_clobber_toml_cache() { let _env = EnvGuard::clear(&["MOQ_CACHE_CAPACITY", "MOQ_CACHE_HEADROOM", "MOQ_CACHE_DURATION"]); diff --git a/rs/moq-relay/src/lib.rs b/rs/moq-relay/src/lib.rs index a22dae154..e5f91ab6f 100644 --- a/rs/moq-relay/src/lib.rs +++ b/rs/moq-relay/src/lib.rs @@ -20,6 +20,7 @@ mod internal; mod listener; mod nodes; mod relay; +mod runtime; mod shutdown; mod stats; #[cfg(test)] @@ -50,6 +51,7 @@ pub use config::*; pub use connection::*; pub use internal::*; pub use relay::*; +pub use runtime::*; pub use shutdown::*; pub use stats::*; pub use web::*; diff --git a/rs/moq-relay/src/relay.rs b/rs/moq-relay/src/relay.rs index 5761fd44f..198e8238e 100644 --- a/rs/moq-relay/src/relay.rs +++ b/rs/moq-relay/src/relay.rs @@ -24,7 +24,9 @@ use anyhow::Context; -use crate::{Auth, Cluster, Config, Connection, DEFAULT_DRAIN_TIMEOUT, Internal, Shutdown, ShutdownTrigger, Web}; +use crate::{ + Auth, Cluster, Config, Connection, DEFAULT_DRAIN_TIMEOUT, Internal, Shutdown, ShutdownTrigger, Web, Workers, +}; /// A fully assembled relay: the listeners and the shared cluster behind them. /// @@ -35,8 +37,10 @@ use crate::{Auth, Cluster, Config, Connection, DEFAULT_DRAIN_TIMEOUT, Internal, /// /// `#[non_exhaustive]`, so destructure with a trailing `..` (a pattern naming /// every field does not compile outside this crate) or move out the fields you -/// need one at a time. Dropping any field you do not name is safe: the pieces -/// that must outlive setup are owned by [`Self::cluster`]. +/// need one at a time. Dropping a field you do not name is safe for all but +/// [`Self::workers`]: the rest of what must outlive setup is owned by +/// [`Self::cluster`], while the workers own their threads and sockets, so +/// dropping them un-binds QUIC. #[non_exhaustive] pub struct Relay { /// The QUIC/WebTransport server, already bound. Feed it to [`serve`]. @@ -73,6 +77,11 @@ pub struct Relay { /// Starts graceful shutdown for [`Self::shutdown`]. pub shutdown_trigger: ShutdownTrigger, + + /// The thread-per-core QUIC workers, already bound. Drive them with + /// [`Workers::serve`]; inert unless `runtime.workers` is configured, in which + /// case [`Self::server`] carries no QUIC listener of its own. + pub workers: Workers, } impl Relay { @@ -89,15 +98,27 @@ impl Relay { let mtls_enabled = !config.listen.tls.root.is_empty(); let server_versions = config.listen.versions(); + // Bind the QUIC workers first: they own the listen address when configured, + // so the server below must not also try to bind it. + let workers = Workers::bind(&config).context("failed to start the QUIC workers")?; + + let mut listen = config.listen.clone(); + if workers.enabled() { + listen.listeners = moq_tokio::listen::Listeners::Stream; + } + #[allow(unused_mut)] - let mut server = config.listen.init(config.quic.clone())?; + let mut server = listen.init(config.quic.clone())?; let client = config.connect.clone().init(config.quic.clone())?; // `None` for a stream-only server (no QUIC); any other error is real. - let addr = match server.local_addr() { - Ok(addr) => Some(addr), - Err(moq_tokio::Error::NoBackend(_)) => None, - Err(err) => return Err(err).context("failed to resolve the QUIC bind address"), + let addr = match workers.local_addr() { + Some(addr) => Some(addr), + None => match server.local_addr() { + Ok(addr) => Some(addr), + Err(moq_tokio::Error::NoBackend(_)) => None, + Err(err) => return Err(err).context("failed to resolve the QUIC bind address"), + }, }; #[cfg(feature = "iroh")] @@ -140,7 +161,9 @@ impl Relay { let drain_timeout = config.drain_timeout.unwrap_or(DEFAULT_DRAIN_TIMEOUT); let (shutdown_trigger, shutdown) = Shutdown::new(drain_timeout); // Create a web server too. mTLS for HTTPS is opt-in via `--web-https-root`. - let web = Web::new(auth.clone(), cluster.clone(), server.certificates(), config.web) + // The workers hold the certificates when they own QUIC; the server has none. + let certificates = workers.certificates().unwrap_or_else(|| server.certificates()); + let web = Web::new(auth.clone(), cluster.clone(), certificates, config.web) .with_shutdown(shutdown.clone()) .with_versions(server_versions); @@ -170,6 +193,7 @@ impl Relay { addr, shutdown, shutdown_trigger, + workers, }) } @@ -187,6 +211,7 @@ impl Relay { web, shutdown, shutdown_trigger, + workers, .. } = self; @@ -209,7 +234,8 @@ impl Relay { Err(err) = started.run() => Err(err).context("cluster failed"), Err(err) = web.run() => Err(err).context("web server failed"), Err(err) = internal.run() => Err(err).context("internal server failed"), - Err(err) = serve(server, cluster, auth, shutdown.clone()) => Err(err).context("server failed"), + Err(err) = serve(server, cluster.clone(), auth.clone(), shutdown.clone()) => Err(err).context("server failed"), + Err(err) = workers.serve(cluster, auth, shutdown.clone()) => Err(err).context("QUIC workers failed"), Err(err) = jemalloc => Err(err).context("jemalloc profiler failed"), res = drain_on_signal(shutdown_trigger, shutdown.drain_timeout) => res, else => Ok(()), diff --git a/rs/moq-relay/src/runtime.rs b/rs/moq-relay/src/runtime.rs new file mode 100644 index 000000000..d6f845c30 --- /dev/null +++ b/rs/moq-relay/src/runtime.rs @@ -0,0 +1,316 @@ +//! Thread-per-core QUIC workers. +//! +//! By default the relay runs one work-stealing runtime and one UDP socket: +//! every packet can cross threads, and every wakeup is a candidate context +//! switch. [`RuntimeConfig::workers`] switches QUIC to the opposite shape. +//! Each worker is a thread of its own, pinned to a core, running a +//! `current_thread` runtime and owning one socket in a `SO_REUSEPORT` group. A +//! connection lands on whichever worker the kernel steers its first packet to +//! and stays there: its QUIC driver, its session, and its tasks never leave +//! that thread, so the locks they take are uncontended and nothing is stolen. +//! +//! Everything that is not QUIC (the web and internal listeners, `tcp`/`unix`, +//! clustering, signals, cert reload) stays on the main runtime. The workers +//! reach the rest of the relay through the shared [`Cluster`], the same way +//! sessions on one runtime already do. +//! +//! # Steering +//! +//! The kernel picks a member of the reuseport group by hashing the packet's +//! 4-tuple, which pins a connection to one worker for as long as the client +//! keeps its address. A client that migrates (a NAT rebinding, a network +//! change) hashes somewhere else and its packets reach a worker that has never +//! seen the connection ID, so the connection is lost rather than migrated. +//! Steering by connection ID instead is the next piece of the work; until then +//! this mode is for measurement and for deployments where that trade is +//! acceptable, which is why it is off by default. + +use std::net::SocketAddr; + +use anyhow::Context; +use clap::Args; +use serde::{Deserialize, Serialize}; + +use crate::{Auth, Cluster, Config, Shutdown}; + +/// How the relay lays its QUIC work out over threads. +#[derive(Args, Clone, Debug, Default, Deserialize, Serialize)] +#[serde(default, deny_unknown_fields)] +#[non_exhaustive] +#[group(id = "runtime-config")] +pub struct RuntimeConfig { + /// Serve QUIC from this many single-threaded workers instead of the shared + /// runtime, each pinned to a core with its own socket on the listen address. + /// + /// A connection is handled start to finish by one worker, which trades the + /// shared runtime's load balancing for no cross-thread traffic per packet. + /// Linux-only, and mutually exclusive with `--listen-tls-generate`, since + /// each worker would otherwise generate and serve a certificate of its own. + /// Unset (the default) keeps QUIC on the shared runtime. + #[arg(long = "runtime-workers", env = "MOQ_RUNTIME_WORKERS")] + pub workers: Option, + + /// Pin each worker to a CPU core, defaulting to on. + /// + /// Pinning is the point of the mode: it keeps a connection's caches warm on + /// one core and stops the scheduler migrating a busy worker. Turn it off to + /// measure what pinning alone is worth, or when sharing the machine with + /// something that manages CPU placement itself. + #[arg(long = "runtime-pin", env = "MOQ_RUNTIME_PIN")] + pub pin: Option, +} + +/// The QUIC workers, bound and waiting to serve. +/// +/// [`Workers::bind`] opens every worker's socket, so a returned value is +/// listening; nothing is accepted until [`Workers::serve`]. Inert (and cheap) +/// when [`RuntimeConfig::workers`] is unset, which is how [`crate::Relay`] can +/// hold one unconditionally. +pub struct Workers { + bound: Vec, + certificates: Option, + addr: Option, +} + +/// One worker thread, bound and parked until it is handed the cluster to serve. +struct Bound { + shard: u16, + start: tokio::sync::oneshot::Sender, + done: tokio::sync::oneshot::Receiver>, +} + +/// What a worker needs to start accepting, known only after the cluster exists. +struct Start { + cluster: Cluster, + auth: Auth, + shutdown: Shutdown, +} + +impl Workers { + /// Bind one socket per worker, or nothing at all when workers are disabled. + /// + /// Returns once every worker is listening, so a port conflict or a bad + /// certificate is an error here rather than a worker that quietly died. + pub fn bind(config: &Config) -> anyhow::Result { + let Some(count) = config.runtime.workers.filter(|count| *count > 0) else { + return Ok(Self::disabled()); + }; + + anyhow::ensure!( + config.listen.tls.generate.is_empty(), + "--runtime-workers cannot be combined with --listen-tls-generate: \ + each worker would generate a different certificate. Pass --listen-tls-cert/--listen-tls-key instead." + ); + + let pin = config.runtime.pin.unwrap_or(true).then(cores).unwrap_or_default(); + + let mut bound = Vec::with_capacity(count as usize); + let mut certificates = None; + let mut addr = None; + + for shard in 0..count { + let listen = shard_config(config, shard, count)?; + let quic = config.quic.clone(); + let core = pin.get(shard as usize % pin.len().max(1)).copied(); + + let (ready_tx, ready_rx) = std::sync::mpsc::channel(); + let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let (done_tx, done_rx) = tokio::sync::oneshot::channel(); + + std::thread::Builder::new() + .name(format!("moq-quic-{shard}")) + .spawn(move || run(shard, core, listen, quic, ready_tx, start_rx, done_tx)) + .with_context(|| format!("failed to spawn QUIC worker {shard}"))?; + + // A worker that fails to bind drops its sender, so the recv error and + // the bind error are the same event; report the bind error. + let ready = ready_rx + .recv() + .map_err(|_| anyhow::anyhow!("QUIC worker {shard} exited before binding"))? + .with_context(|| format!("QUIC worker {shard} failed to listen"))?; + + certificates.get_or_insert(ready.certificates); + addr.get_or_insert(ready.addr); + + bound.push(Bound { + shard, + start: start_tx, + done: done_rx, + }); + } + + tracing::info!(workers = count, pinned = !pin.is_empty(), "started QUIC workers"); + + Ok(Self { + bound, + // Every worker loads the same certificate files, so any one of them + // answers for the fingerprint endpoint. + certificates, + addr, + }) + } + + /// A pool that owns nothing, for a relay serving QUIC on the shared runtime. + fn disabled() -> Self { + Self { + bound: Vec::new(), + certificates: None, + addr: None, + } + } + + /// Whether QUIC is being served by workers rather than the shared runtime. + pub fn enabled(&self) -> bool { + !self.bound.is_empty() + } + + /// The address every worker is bound to, or `None` when they are disabled. + pub fn local_addr(&self) -> Option { + self.addr + } + + /// The certificates the workers are serving, tracking hot reloads, or `None` + /// when they are disabled and the shared server holds them instead. + pub fn certificates(&self) -> Option { + self.certificates.clone() + } + + /// Release every worker to accept into `cluster`, then resolve when one stops. + /// + /// Pends forever when the workers are disabled, so it composes into a + /// `select!` either way. A worker only stops on error or on shutdown, so the + /// first one to finish ends the relay the same way the shared accept loop does. + pub async fn serve(self, cluster: Cluster, auth: Auth, shutdown: Shutdown) -> anyhow::Result<()> { + if self.bound.is_empty() { + return std::future::pending().await; + } + + let mut running = futures::stream::FuturesUnordered::new(); + for worker in self.bound { + let start = Start { + cluster: cluster.clone(), + auth: auth.clone(), + shutdown: shutdown.clone(), + }; + let shard = worker.shard; + // The worker is parked on this receiver; it only drops if the thread + // died, which its `done` half reports with the real error. + let _ = worker.start.send(start); + running.push(async move { + match worker.done.await { + Ok(res) => res.with_context(|| format!("QUIC worker {shard} failed")), + Err(_) => Err(anyhow::anyhow!("QUIC worker {shard} exited without reporting")), + } + }); + } + + use futures::StreamExt; + running + .next() + .await + .unwrap_or_else(|| Err(anyhow::anyhow!("no QUIC workers"))) + } +} + +/// What a worker reports back once it is listening. +struct Ready { + certificates: moq_tokio::tls::Certificates, + addr: SocketAddr, +} + +/// The listen config for one worker: QUIC only, on its own socket in the group. +fn shard_config(config: &Config, index: u16, count: u16) -> anyhow::Result { + let mut listen = config.listen.clone(); + listen.listeners = moq_tokio::listen::Listeners::Quic; + listen.shard = Some( + moq_tokio::listen::Shard::new(index, count) + .with_context(|| format!("invalid QUIC worker {index} of {count}"))?, + ); + Ok(listen) +} + +/// One worker thread: bind, report, wait for the cluster, then serve until it stops. +fn run( + shard: u16, + core: Option, + listen: moq_tokio::listen::Config, + quic: moq_tokio::quic::Config, + ready: std::sync::mpsc::Sender>, + start: tokio::sync::oneshot::Receiver, + done: tokio::sync::oneshot::Sender>, +) { + if let Some(core) = core { + pin(shard, core); + } + + let runtime = match tokio::runtime::Builder::new_current_thread().enable_all().build() { + Ok(runtime) => runtime, + Err(err) => { + let _ = ready.send(Err(anyhow::Error::new(err).context("failed to build worker runtime"))); + return; + } + }; + + runtime.block_on(async move { + // Built inside the runtime on purpose: the QUIC backend spawns its socket + // driver where it is constructed, which is what keeps this worker's packets + // on this thread. + let server = match listen.init(quic) { + Ok(server) => server, + Err(err) => { + let _ = ready.send(Err(err.into())); + return; + } + }; + + let addr = match server.local_addr() { + Ok(addr) => addr, + Err(err) => { + let _ = ready.send(Err(err.into())); + return; + } + }; + + tracing::debug!(shard, %addr, "QUIC worker listening"); + if ready + .send(Ok(Ready { + certificates: server.certificates(), + addr, + })) + .is_err() + { + return; + } + + let Ok(start) = start.await else { + return; + }; + let _ = done.send(crate::serve(server, start.cluster, start.auth, start.shutdown).await); + }); +} + +/// A core to pin a worker to. Kept behind a type so the non-Linux build has +/// something to name without depending on the pinning crate's types. +type CoreId = core_affinity::CoreId; + +/// The cores workers may be pinned to, in the order they are handed out. +/// +/// Empty when the platform will not report them, which disables pinning rather +/// than failing: the mode's other half (one runtime and one socket per worker) +/// is still worth having. +fn cores() -> Vec { + let cores = core_affinity::get_core_ids().unwrap_or_default(); + if cores.is_empty() { + tracing::warn!("could not enumerate CPU cores; QUIC workers will not be pinned"); + } + cores +} + +/// Pin the calling thread to `core`, warning if the platform refuses. +fn pin(shard: u16, core: CoreId) { + if core_affinity::set_for_current(core) { + tracing::debug!(shard, core = core.id, "pinned QUIC worker"); + } else { + tracing::warn!(shard, core = core.id, "failed to pin QUIC worker"); + } +} diff --git a/rs/moq-relay/tests/runtime_workers.rs b/rs/moq-relay/tests/runtime_workers.rs new file mode 100644 index 000000000..a1c60b419 --- /dev/null +++ b/rs/moq-relay/tests/runtime_workers.rs @@ -0,0 +1,152 @@ +//! The thread-per-core QUIC mode, end to end through a real relay. +//! +//! Linux-only, because the mode is: every worker binds the listen address with +//! `SO_REUSEPORT`, and no other platform load-balances a unicast UDP port +//! across the group. +#![cfg(target_os = "linux")] + +use std::net::{SocketAddr, UdpSocket}; +use std::time::Duration; + +use moq_relay::{Config, PublicConfig, Relay}; +use moq_tokio::moq_net::{self, Origin}; + +const TIMEOUT: Duration = Duration::from_secs(10); +const WORKERS: u16 = 4; + +/// A UDP port nothing is bound to. +/// +/// Every worker binds the same port, so this cannot be `:0`: each would pick an +/// ephemeral port of its own and they would not form a group. +fn free_udp_port() -> u16 { + let probe = UdpSocket::bind("127.0.0.1:0").expect("bind probe"); + let port = probe.local_addr().expect("local addr").port(); + drop(probe); + port +} + +/// A self-signed certificate on disk. Workers refuse `listen.tls.generate`, +/// since each would generate one of its own and serve a different identity. +fn certificate(dir: &std::path::Path) -> (std::path::PathBuf, std::path::PathBuf) { + let key = rcgen::KeyPair::generate().expect("keypair"); + let params = rcgen::CertificateParams::new(vec!["localhost".to_string()]).expect("cert params"); + let cert = params.self_signed(&key).expect("self-signed cert"); + let cert_path = dir.join("cert.pem"); + let key_path = dir.join("key.pem"); + std::fs::write(&cert_path, cert.pem()).expect("write cert"); + std::fs::write(&key_path, key.serialize_pem()).expect("write key"); + (cert_path, key_path) +} + +fn client() -> moq_tokio::Client { + let mut config = moq_tokio::connect::Config::default(); + config.tls.insecure = Some(true); + config.once = Some(true); + config.bind = Some("127.0.0.1:0".parse().expect("parse bind")); + config.init(Default::default()).expect("client init") +} + +async fn connect(client: moq_tokio::Client, url: url::Url) -> moq_tokio::Connection { + tokio::time::timeout(TIMEOUT, client.with_reconnect(false).connect(url).established()) + .await + .expect("connect timeout") + .expect("connect failed") +} + +/// A publisher and several subscribers over QUIC, served by a pool of workers. +/// +/// The relay hands each connection to whichever worker the kernel steered its +/// first packet to, so the subscribers are spread over threads the publisher is +/// probably not on: the frame has to travel through the shared origin to reach +/// them. That crossing is the part of the mode that could deadlock or drop, and +/// it is what this asserts, whichever way the placement lands. +#[tokio::test] +async fn workers_serve_quic_and_share_one_origin() { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let dir = tempfile::tempdir().expect("tempdir"); + let (cert, key) = certificate(dir.path()); + let port = free_udp_port(); + + let mut config = Config::default(); + config.listen.bind = Some(format!("127.0.0.1:{port}")); + config.listen.tls.cert = vec![cert]; + config.listen.tls.key = vec![key]; + config.runtime.workers = Some(WORKERS); + // A CI container may restrict which cores it may run on, and pinning is not + // what this test is about. + config.runtime.pin = Some(false); + #[allow(deprecated)] + let public = PublicConfig::Simple(vec![String::new()]); + config.auth.public = Some(public); + + let relay = Relay::load(config).await.expect("load relay"); + let expected: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap(); + assert_eq!(relay.addr, Some(expected), "workers bound a different address"); + + let workers = tokio::spawn( + relay + .workers + .serve(relay.cluster.clone(), relay.auth.clone(), relay.shutdown.clone()), + ); + + let url: url::Url = format!("https://127.0.0.1:{port}/workers").parse().expect("parse url"); + + // ── publisher ─────────────────────────────────────────────────── + let origin = moq_tokio::origin::spawn(Origin::random()); + let mut broadcast = origin + .create_broadcast("test", moq_net::broadcast::Route::new().with_announce(true)) + .expect("create broadcast"); + let mut track = broadcast.create_track("video", None).expect("create track"); + let mut group = track.append_group().expect("append group"); + group + .write_frame(moq_net::Timestamp::ZERO, b"hello".as_ref()) + .expect("write frame"); + group.finish().expect("finish group"); + + let publisher = connect(client().with_publisher(&origin), url.clone()).await; + + // ── subscribers ───────────────────────────────────────────────── + let mut subscribers = Vec::new(); + for _ in 0..WORKERS { + let origin = moq_tokio::origin::spawn(Origin::random()); + let announced = origin.consume().announced(); + let connection = connect(client().with_subscriber(origin), url.clone()).await; + subscribers.push((connection, announced)); + } + + for (index, (_connection, announced)) in subscribers.iter_mut().enumerate() { + let moq_net::announce::Update { path, broadcast } = tokio::time::timeout(TIMEOUT, announced.next()) + .await + .unwrap_or_else(|_| panic!("subscriber {index} announcement timeout")) + .expect("origin closed"); + assert_eq!(path.as_str(), "test"); + let broadcast = broadcast.expect("expected announce, got unannounce"); + + let mut subscription = broadcast + .track("video") + .unwrap() + .subscribe(None) + .await + .expect("subscribe"); + let mut group = tokio::time::timeout(TIMEOUT, subscription.recv_group()) + .await + .unwrap_or_else(|_| panic!("subscriber {index} recv_group timeout")) + .expect("recv_group failed") + .expect("track closed prematurely"); + let frame = tokio::time::timeout(TIMEOUT, group.read_frame()) + .await + .unwrap_or_else(|_| panic!("subscriber {index} read_frame timeout")) + .expect("read_frame failed") + .expect("group closed prematurely"); + assert_eq!(&frame.payload[..], b"hello"); + } + + assert!(!workers.is_finished(), "a QUIC worker stopped while serving"); + + drop(track); + drop(broadcast); + drop(publisher); + drop(subscribers); + workers.abort(); +} diff --git a/rs/moq-tokio/Cargo.toml b/rs/moq-tokio/Cargo.toml index 33fcefb94..86f146461 100644 --- a/rs/moq-tokio/Cargo.toml +++ b/rs/moq-tokio/Cargo.toml @@ -74,7 +74,7 @@ rustls-webpki = { version = "0.103", features = ["aws-lc-rs"], optional = true } serde = { version = "1", features = ["derive"] } serde_with = { version = "3", features = ["hex"] } sha2 = { workspace = true, optional = true } -socket2 = "0.6" +socket2 = { version = "0.6", features = ["all"] } subtle = { workspace = true, optional = true } thiserror = "2" tikv-jemalloc-ctl = { version = "0.6", optional = true } diff --git a/rs/moq-tokio/src/bind.rs b/rs/moq-tokio/src/bind.rs index 6f992826f..22a573115 100644 --- a/rs/moq-tokio/src/bind.rs +++ b/rs/moq-tokio/src/bind.rs @@ -10,6 +10,7 @@ //! See . use socket2::{Domain, Protocol, Socket, TcpKeepalive, Type}; +use std::io; use std::net::{SocketAddr, TcpListener, UdpSocket}; use std::time::Duration; @@ -23,15 +24,75 @@ use std::time::Duration; const KEEPALIVE_IDLE: Duration = Duration::from_secs(30); const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(10); +/// What to bind a UDP socket to, and how. +/// +/// [`Udp::new`] takes the address and leaves everything else at the platform +/// defaults; reach for the `with_*` methods for the rest. +#[derive(Clone, Copy, Debug)] +pub struct Udp { + addr: SocketAddr, + reuse_port: bool, +} + +impl Udp { + /// Bind this address, with the platform defaults for everything else. + pub fn new(addr: SocketAddr) -> Self { + Self { + addr, + reuse_port: false, + } + } + + /// Share the port with the other sockets bound this way (`SO_REUSEPORT`). + /// + /// Linux spreads inbound datagrams across every socket in the group, which is + /// what lets one worker per core own a socket on the same port instead of + /// funnelling every packet through one. Enable it on *every* member: the first + /// socket to bind without it owns the port outright and the rest fail. + /// + /// Linux-only, and [`udp`] fails with [`io::ErrorKind::Unsupported`] elsewhere + /// rather than binding a group the platform does not balance. macOS and the + /// BSDs accept the option and then deliver a unicast flow to a single member, + /// so the workers would come up looking healthy with one of them serving + /// everything. + pub fn with_reuse_port(mut self, enabled: bool) -> Self { + self.reuse_port = enabled; + self + } +} + /// Bind a UDP socket, making an IPv6 socket dual-stack so it also serves IPv4. -pub fn udp(addr: SocketAddr) -> std::io::Result { +pub fn udp(options: Udp) -> io::Result { + let addr = options.addr; + let domain = if addr.is_ipv4() { Domain::IPV4 } else { Domain::IPV6 }; let socket = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?; make_dual_stack(&socket, addr); + if options.reuse_port { + set_reuse_port(&socket)?; + } socket.bind(&addr.into())?; Ok(socket.into()) } +/// Set `SO_REUSEPORT`, or report that this platform cannot load-balance a port. +/// +/// Not best-effort like the other socket options here: a caller asking for this +/// is building a group of sockets that only works if the kernel spreads traffic +/// over it, so a silent no-op would leave one member serving everything. +#[cfg(target_os = "linux")] +fn set_reuse_port(socket: &Socket) -> io::Result<()> { + socket.set_reuse_port(true) +} + +#[cfg(not(target_os = "linux"))] +fn set_reuse_port(_socket: &Socket) -> io::Result<()> { + Err(io::Error::new( + io::ErrorKind::Unsupported, + "SO_REUSEPORT load balancing is Linux-only", + )) +} + /// Whether `socket` also reaches IPv4, through IPv4-mapped addresses. /// /// [`udp`] clears `IPV6_V6ONLY` best-effort, so this reads back what the @@ -51,7 +112,7 @@ pub(crate) fn udp_is_dual_stack(socket: &UdpSocket) -> bool { /// /// The returned listener is non-blocking, ready to be adopted by an async runtime /// (`tokio::net::TcpListener::from_std`, `axum_server::from_tcp`). -pub fn tcp(addr: SocketAddr) -> std::io::Result { +pub fn tcp(addr: SocketAddr) -> io::Result { let domain = if addr.is_ipv4() { Domain::IPV4 } else { Domain::IPV6 }; let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))?; make_dual_stack(&socket, addr); @@ -115,7 +176,7 @@ mod tests { fn udp_ipv6_is_dual_stack() { // An IPv6 wildcard bind should come back dual-stack so IPv4 traffic // reaches it. socket2 lets us read the option back to confirm. - let socket = match udp("[::]:0".parse().unwrap()) { + let socket = match udp(Udp::new("[::]:0".parse().unwrap())) { Ok(socket) => socket, Err(err) if skip_if_no_ipv6(&err) => return, Err(err) => panic!("failed to bind IPv6 UDP socket: {err}"), @@ -126,10 +187,85 @@ mod tests { #[test] fn udp_ipv4_still_binds() { - let socket = udp("127.0.0.1:0".parse().unwrap()).unwrap(); + let socket = udp(Udp::new("127.0.0.1:0".parse().unwrap())).unwrap(); assert!(socket.local_addr().unwrap().is_ipv4()); } + /// A reuseport group is only useful if every member can hold the same port, + /// so bind a second socket to the first one's address and check it takes. + #[test] + #[cfg(target_os = "linux")] + fn udp_reuse_port_shares_a_port() { + let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + let first = udp(Udp::new(addr).with_reuse_port(true)).unwrap(); + let bound = first.local_addr().unwrap(); + let second = udp(Udp::new(bound).with_reuse_port(true)).unwrap(); + assert_eq!(second.local_addr().unwrap(), bound); + } + + /// Without the option the second bind must lose the port, which is what makes + /// a missed `with_reuse_port` on one member a startup failure rather than a + /// silently lopsided group. + #[test] + #[cfg(target_os = "linux")] + fn udp_without_reuse_port_keeps_the_port() { + let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + let first = udp(Udp::new(addr)).unwrap(); + let bound = first.local_addr().unwrap(); + assert!(udp(Udp::new(bound).with_reuse_port(true)).is_err()); + } + + /// The point of the option is that the kernel spreads traffic over the group, + /// which is what a worker-per-core listener is built on. Send from a spread of + /// source ports and check that more than one member is fed. + #[test] + #[cfg(target_os = "linux")] + fn udp_reuse_port_spreads_datagrams() { + const MEMBERS: usize = 4; + const SENDERS: usize = 64; + + let mut group = Vec::with_capacity(MEMBERS); + let mut addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + for _ in 0..MEMBERS { + let socket = udp(Udp::new(addr).with_reuse_port(true)).unwrap(); + addr = socket.local_addr().unwrap(); + socket.set_nonblocking(true).unwrap(); + group.push(socket); + } + + // Each sender gets its own ephemeral source port, which is what the + // kernel hashes on. + for _ in 0..SENDERS { + let sender = udp(Udp::new("127.0.0.1:0".parse().unwrap())).unwrap(); + sender.send_to(b"quic", addr).unwrap(); + } + + let mut fed = 0; + let mut total = 0; + for socket in &group { + let mut received = 0; + let mut buf = [0u8; 8]; + while socket.recv_from(&mut buf).is_ok() { + received += 1; + } + total += received; + fed += usize::from(received > 0); + } + + assert_eq!(total, SENDERS, "every datagram reached exactly one member"); + assert!(fed > 1, "only {fed} of {MEMBERS} members were fed"); + } + + /// Elsewhere the request fails loudly instead of binding a group the kernel + /// won't balance. + #[test] + #[cfg(not(target_os = "linux"))] + fn udp_reuse_port_is_linux_only() { + let addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + let err = udp(Udp::new(addr).with_reuse_port(true)).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::Unsupported); + } + #[test] fn tcp_ipv6_is_dual_stack() { let listener = match tcp("[::]:0".parse().unwrap()) { diff --git a/rs/moq-tokio/src/listen.rs b/rs/moq-tokio/src/listen.rs index 4b57e49ed..a66227ff7 100644 --- a/rs/moq-tokio/src/listen.rs +++ b/rs/moq-tokio/src/listen.rs @@ -121,6 +121,94 @@ pub struct Config { #[arg(skip)] #[serde(default, skip_serializing_if = "Option::is_none")] pub quic: Option, + + /// Which of the configured listeners this server opens. + /// + /// Programmatic only: there is no flag or TOML key, because splitting the + /// listeners across servers is a decision about which threads run what, not + /// about what the process listens on. + #[arg(skip)] + #[serde(skip)] + pub listeners: Listeners, + + /// This server's slot in a `SO_REUSEPORT` group, when several share a port. + /// + /// Programmatic only, for the same reason as [`Self::listeners`]: the group is + /// a set of sockets one process binds itself, not something an operator names. + #[arg(skip)] + #[serde(skip)] + pub shard: Option, +} + +/// Which listeners a [`Server`](crate::Server) opens out of the ones configured. +/// +/// [`All`](Self::All) is the default and what a single-threaded process wants. +/// The split exists so a process can run QUIC on dedicated threads (one +/// [`Quic`](Self::Quic) server per [`Shard`]) while a [`Stream`](Self::Stream) +/// server keeps `tcp`/`unix` on its main runtime. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum Listeners { + /// Every listener the config asks for. + #[default] + All, + + /// The QUIC (UDP) listener only, even when a `tcp`/`unix` listener is configured. + Quic, + + /// The stream (`tcp`/`unix`) listeners only, and no QUIC even when nothing else + /// is configured. A server with neither accepts nothing, which is the honest + /// outcome rather than a surprise UDP bind. + Stream, +} + +impl Listeners { + /// Whether a QUIC backend should be built. + pub(crate) fn quic(self) -> bool { + matches!(self, Self::All | Self::Quic) + } + + /// Whether the `tcp`/`unix` listeners should be opened. + #[cfg_attr( + not(any(feature = "tcp", all(feature = "uds", unix))), + expect(dead_code, reason = "no stream listener is compiled in") + )] + pub(crate) fn stream(self) -> bool { + matches!(self, Self::All | Self::Stream) + } +} + +/// One server's slot in a group of sockets sharing a port via `SO_REUSEPORT`. +/// +/// Every member binds the same address and the kernel spreads inbound +/// connections across the group, so N servers on N threads can serve one port +/// without a shared socket between them. [`index`](Self::index) names this +/// member and is stable for the life of the group, so an owner can attribute a +/// per-thread measurement back to a socket. +/// +/// Membership is fixed once the group is bound: the kernel picks a member by +/// position, so a socket joining or leaving reshuffles every mapping. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Shard { + index: u16, + count: u16, +} + +impl Shard { + /// Slot `index` of a group of `count` sockets, or `None` if that slot does not + /// exist (`count` of zero, or an `index` at or past the end). + pub fn new(index: u16, count: u16) -> Option { + (index < count).then_some(Self { index, count }) + } + + /// This member's position in the group, from zero. + pub fn index(self) -> u16 { + self.index + } + + /// How many sockets share the port. + pub fn count(self) -> u16 { + self.count + } } /// The `--server-*` flags from before the accept side was named `listen`. @@ -366,6 +454,49 @@ mod tests { assert!(config.validate().is_ok()); } + /// A slot has to exist in the group it names. + #[test] + fn shard_slots_are_bounded() { + assert_eq!(Shard::new(0, 1).map(|shard| shard.count()), Some(1)); + assert_eq!(Shard::new(3, 4).map(|shard| shard.index()), Some(3)); + assert!(Shard::new(4, 4).is_none()); + assert!(Shard::new(0, 0).is_none()); + } + + /// A stream-only server opens no QUIC listener even with a bind configured, + /// which is what lets another thread hold that address. + #[tokio::test] + async fn stream_listeners_leave_quic_alone() { + let config = Config { + bind: Some("127.0.0.1:0".to_string()), + listeners: Listeners::Stream, + ..Default::default() + }; + + let server = config.init(Default::default()).unwrap(); + assert!(matches!(server.local_addr(), Err(crate::Error::NoBackend(_)))); + } + + /// The mirror image: a QUIC-only server ignores a configured stream listener, + /// so N of them can share one config without N tcp binds fighting over a port. + #[cfg(feature = "tcp")] + #[tokio::test] + async fn quic_listeners_leave_streams_alone() { + let mut config = Config { + bind: Some("127.0.0.1:0".to_string()), + listeners: Listeners::Quic, + ..Default::default() + }; + config.tcp.bind = Some("127.0.0.1:0".parse().unwrap()); + config.tls.generate = vec!["localhost".to_string()]; + + let server = config.init(Default::default()).unwrap(); + assert!(server.local_addr().is_ok()); + // Only the stream listeners report accept(2) health, so an empty set is + // how a server with none of them reads. + assert!(server.accept_health().is_empty()); + } + /// The canonical spellings, which is what `--help` teaches. #[test] fn canonical_spellings_parse() { diff --git a/rs/moq-tokio/src/noq.rs b/rs/moq-tokio/src/noq.rs index 78e04744b..ac48b4966 100644 --- a/rs/moq-tokio/src/noq.rs +++ b/rs/moq-tokio/src/noq.rs @@ -256,7 +256,7 @@ pub(crate) struct NoqClient { impl NoqClient { pub fn new(config: &connect::Config, quic: &crate::quic::Config) -> Result { - let socket = crate::bind::udp(config.resolved_bind()).map_err(Error::BindSocket)?; + let socket = crate::bind::udp(crate::bind::Udp::new(config.resolved_bind())).map_err(Error::BindSocket)?; let dual_stack = crate::bind::udp_is_dual_stack(&socket); let mut transport = noq::TransportConfig::default(); @@ -552,7 +552,8 @@ impl NoqServer { })); } - let socket = crate::bind::udp(listen).map_err(Error::BindSocket)?; + let socket = crate::bind::udp(crate::bind::Udp::new(listen).with_reuse_port(config.shard.is_some())) + .map_err(Error::BindSocket)?; // Create the generic QUIC endpoint. let quic = noq::Endpoint::new(endpoint_config, Some(tls), socket, runtime).map_err(Error::CreateEndpoint)?; diff --git a/rs/moq-tokio/src/quiche.rs b/rs/moq-tokio/src/quiche.rs index 4d2b4dd45..fb394a792 100644 --- a/rs/moq-tokio/src/quiche.rs +++ b/rs/moq-tokio/src/quiche.rs @@ -322,7 +322,7 @@ impl QuicheClient { // Bind the first attempt now so candidate selection uses the socket's actual // family and dual-stack state. Later attempts bind their own ephemeral // sockets because each quiche connection must own its socket. - let socket = crate::bind::udp(self.bind)?; + let socket = crate::bind::udp(crate::bind::Udp::new(self.bind))?; let local = socket.local_addr()?; let dual_stack = crate::bind::udp_is_dual_stack(&socket); @@ -359,7 +359,7 @@ impl QuicheClient { async move { let socket = match socket { Some(socket) => socket, - None => crate::bind::udp(this.bind)?, + None => crate::bind::udp(crate::bind::Udp::new(this.bind))?, }; let builder = this.builder(socket, alpns, &verification, roots, &host)?; let conn = builder @@ -566,7 +566,7 @@ impl QuicheServer { let listen = crate::util::resolve(config.bind.as_deref(), crate::server::DEFAULT_BIND).map_err(Error::ResolveBind)?; - let socket = crate::bind::udp(listen)?; + let socket = crate::bind::udp(crate::bind::Udp::new(listen).with_reuse_port(config.shard.is_some()))?; let (chain, key) = if !config.tls.generate.is_empty() { generate_quiche_cert(&config.tls.generate)? diff --git a/rs/moq-tokio/src/quinn.rs b/rs/moq-tokio/src/quinn.rs index 0a0dad25b..bceee64b5 100644 --- a/rs/moq-tokio/src/quinn.rs +++ b/rs/moq-tokio/src/quinn.rs @@ -269,7 +269,7 @@ pub(crate) struct QuinnClient { impl QuinnClient { pub fn new(config: &connect::Config, quic: &crate::quic::Config) -> Result { - let socket = crate::bind::udp(config.resolved_bind()).map_err(Error::BindSocket)?; + let socket = crate::bind::udp(crate::bind::Udp::new(config.resolved_bind())).map_err(Error::BindSocket)?; let dual_stack = crate::bind::udp_is_dual_stack(&socket); let quic = quic.resolve(); @@ -563,7 +563,8 @@ impl QuinnServer { endpoint_config.cid_generator(move || Box::new(ServerIdGenerator::new(server_id.clone(), nonce_len))); } - let socket = crate::bind::udp(listen).map_err(Error::BindSocket)?; + let socket = crate::bind::udp(crate::bind::Udp::new(listen).with_reuse_port(config.shard.is_some())) + .map_err(Error::BindSocket)?; // Create the generic QUIC endpoint. let quic = quinn::Endpoint::new(endpoint_config, Some(tls), socket, runtime).map_err(Error::CreateEndpoint)?; diff --git a/rs/moq-tokio/src/server.rs b/rs/moq-tokio/src/server.rs index 532d0e5d9..eac65733f 100644 --- a/rs/moq-tokio/src/server.rs +++ b/rs/moq-tokio/src/server.rs @@ -113,7 +113,7 @@ impl Server { // `--listen`) doesn't also open UDP/443. quic.validate()?; - let build_quic = config.bind.is_some() || !config.has_stream_listener(); + let build_quic = config.listeners.quic() && (config.bind.is_some() || !config.has_stream_listener()); #[cfg(not(any(feature = "noq", feature = "quinn", feature = "quiche")))] if config.bind.is_some() { return Err(Error::NoBackend( @@ -167,11 +167,11 @@ impl Server { #[cfg(any(feature = "tcp", all(feature = "uds", unix)))] let mut stream_binds = Vec::new(); #[cfg(feature = "tcp")] - if let Some(addr) = config.tcp.bind { + if let Some(addr) = config.tcp.bind.filter(|_| config.listeners.stream()) { stream_binds.push(StreamBind::Tcp(addr)); } #[cfg(all(feature = "uds", unix))] - if let Some(path) = config.unix.bind.clone() { + if let Some(path) = config.unix.bind.clone().filter(|_| config.listeners.stream()) { stream_binds.push(StreamBind::Unix(path)); } // `None` (or an all-empty allowlist) means the listener enforces nothing. From f7d666bd13aca6451890934d80c89b5217122f0e Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Tue, 18 Aug 2026 21:36:01 -0700 Subject: [PATCH 2/2] feat(relay): steer the worker group by connection ID Address hashing is the wrong key for QUIC. The kernel picks a reuseport member by hashing the packet's 4-tuple, so a client that changes address (a NAT rebinding, a network change, plain connection migration) hashes to a different worker and its packets arrive at a socket that has never heard of the connection. It dies instead of migrating. Without steering, `runtime.workers` is not a mode with a caveat, it is a mode that drops connections. So the group is steered by connection ID, in the new `moq_tokio::steer`: - Each worker issues connection IDs whose first byte is congruent to its index modulo the group size, keeping 256/count of that byte's randomness and all of the other 19 bytes. - A classic-BPF filter on the group reads that byte back and returns it modulo the group size, which is the index the kernel selects with. `SO_ATTACH_REUSEPORT_CBPF` runs with the UDP payload at offset 0, so the whole rule is seven instructions: branch on the header form bit, load the connection ID's first byte from offset 1 or 6, reduce, return. A client's first packets carry an ID the client invented, which encodes nothing. Those hash arbitrarily, and that is correct: whoever takes the Initial owns the connection and issues IDs naming itself, so every later packet steers to it, and a retransmitted Initial repeats the same client-chosen ID and reaches the same worker rather than starting a second handshake elsewhere. Classic BPF rather than the SK_REUSEPORT + SOCKARRAY eBPF program: the rule is small enough, and cBPF needs no CAP_BPF, no BPF toolchain in the build, and no map to keep alive across restarts. It selects by position, so the group must be built once, in order, and never resized, which is what `Workers::bind` already does. This needs a backend whose connection IDs we choose. quinn and noq allow it; quiche exposes no such hook through `web-transport-quiche`, so it refuses to start with workers rather than silently hashing addresses. QUIC-LB wants the same bytes, so that combination is refused too. Also from review of the first commit: - Workers were detached: `serve` consumed the pool and kept only the `done` receivers, so cancelling it left the threads accepting into a cluster nobody was driving, and a replacement pool would join the same reuseport group as the orphans. The pool now keeps its `JoinHandle`s and a per-worker stop channel, `serve` borrows instead of consuming, and `Drop` signals every worker and joins every thread. - An ephemeral listen port gave each worker a port of its own instead of a shared one, leaving all but the first unreachable behind an address that read as bound. Every worker's bound address is now compared against the first. - The module claimed cert reload stays on the shared runtime; each worker in fact watches the files itself. Documented, with the non-atomic rotation window that implies (#2924). Verified on Linux (container, kernel 6.19 aarch64). The steering test sends short- and long-header packets carrying each member's connection ID from a fresh source port every time and asserts each lands on its own socket; with the filter detached it fails on the first packet. Co-Authored-By: Claude Opus 5 --- doc/bin/relay/config.md | 22 +- rs/moq-relay/src/relay.rs | 6 +- rs/moq-relay/src/runtime.rs | 192 +++++++++++++---- rs/moq-relay/tests/runtime_workers.rs | 109 ++++++++-- rs/moq-tokio/src/lib.rs | 1 + rs/moq-tokio/src/noq.rs | 38 +++- rs/moq-tokio/src/quiche.rs | 15 +- rs/moq-tokio/src/quinn.rs | 55 ++++- rs/moq-tokio/src/steer.rs | 284 ++++++++++++++++++++++++++ 9 files changed, 653 insertions(+), 69 deletions(-) create mode 100644 rs/moq-tokio/src/steer.rs diff --git a/doc/bin/relay/config.md b/doc/bin/relay/config.md index 408adb9f7..6b5f16ff2 100644 --- a/doc/bin/relay/config.md +++ b/doc/bin/relay/config.md @@ -58,13 +58,23 @@ runtime serves every connection off one UDP socket. # its packets never cross threads. Everything else (HTTP, WebSocket, tcp/unix, # clustering) stays on the shared runtime. Omit to keep QUIC there too. # -# Linux only, and incompatible with listen.tls.generate: each worker would -# generate a certificate of its own. Point at real certificate files instead. +# Packets are steered to their worker by connection ID, not by address, so a +# client that migrates (a NAT rebinding, a network change) stays with the worker +# that owns its connection. # -# The kernel steers a packet to a worker by hashing its source address and port, -# so a client that changes address mid-connection (a NAT rebinding, a network -# change) lands on a worker that has never seen it and loses the connection. -# Treat this as a benchmarking knob until steering follows the connection ID. +# Linux only. Needs a backend whose connection IDs can name the owning worker, +# so listen.backend must be quinn (the default) or noq; quiche refuses to start. +# Cannot be combined with listen.quic_lb_id, which wants the same bytes of the +# connection ID. +# +# The listen address needs an explicit non-zero port: an ephemeral bind gives +# each worker a port of its own instead of a shared one. +# +# Incompatible with listen.tls.generate, since each worker would generate a +# certificate of its own. Point at real certificate files instead. Each worker +# loads and watches those files itself, so a rotation is not atomic across the +# group: for as long as the reloads take, two workers can be serving different +# certificates. See https://github.com/moq-dev/moq/issues/2924. workers = 8 # Pin each worker to a CPU core. Default: true. diff --git a/rs/moq-relay/src/relay.rs b/rs/moq-relay/src/relay.rs index 198e8238e..ef0dd698f 100644 --- a/rs/moq-relay/src/relay.rs +++ b/rs/moq-relay/src/relay.rs @@ -40,7 +40,9 @@ use crate::{ /// need one at a time. Dropping a field you do not name is safe for all but /// [`Self::workers`]: the rest of what must outlive setup is owned by /// [`Self::cluster`], while the workers own their threads and sockets, so -/// dropping them un-binds QUIC. +/// dropping them stops serving QUIC and releases the port. Keep the pool alive +/// for as long as you want it served, including while [`Workers::serve`] runs, +/// which borrows rather than consuming it for exactly that reason. #[non_exhaustive] pub struct Relay { /// The QUIC/WebTransport server, already bound. Feed it to [`serve`]. @@ -211,7 +213,7 @@ impl Relay { web, shutdown, shutdown_trigger, - workers, + mut workers, .. } = self; diff --git a/rs/moq-relay/src/runtime.rs b/rs/moq-relay/src/runtime.rs index d6f845c30..e78134307 100644 --- a/rs/moq-relay/src/runtime.rs +++ b/rs/moq-relay/src/runtime.rs @@ -10,20 +10,34 @@ //! that thread, so the locks they take are uncontended and nothing is stolen. //! //! Everything that is not QUIC (the web and internal listeners, `tcp`/`unix`, -//! clustering, signals, cert reload) stays on the main runtime. The workers -//! reach the rest of the relay through the shared [`Cluster`], the same way -//! sessions on one runtime already do. +//! clustering, signals) stays on the main runtime. The workers reach the rest +//! of the relay through the shared [`Cluster`], the same way sessions on one +//! runtime already do. +//! +//! # TLS +//! +//! Each worker builds its own listener, so each loads the certificate files +//! itself and watches them for reloads independently. They serve the same +//! identity, but a rotation is not atomic across the group: for as long as the +//! reloads take, two workers can be serving different certificates, and a +//! worker whose watcher failed to start keeps serving the old one. The +//! fingerprint published at `/certificate.sha256` is the first worker's, which +//! is why a generated certificate is refused here (see [`Workers::bind`]). //! //! # Steering //! -//! The kernel picks a member of the reuseport group by hashing the packet's -//! 4-tuple, which pins a connection to one worker for as long as the client -//! keeps its address. A client that migrates (a NAT rebinding, a network -//! change) hashes somewhere else and its packets reach a worker that has never -//! seen the connection ID, so the connection is lost rather than migrated. -//! Steering by connection ID instead is the next piece of the work; until then -//! this mode is for measurement and for deployments where that trade is -//! acceptable, which is why it is off by default. +//! Left alone, the kernel picks a member of a reuseport group by hashing the +//! packet's 4-tuple, which is the wrong key for QUIC: a client that changes +//! address would hash somewhere else and its packets would arrive at a worker +//! that has never heard of the connection. So the group is steered by +//! connection ID instead (`moq_tokio`'s `steer` module): each worker issues +//! connection IDs that name it, and a filter on the group reads them back. A +//! connection follows its ID, not its address, which is what QUIC means by +//! migration. +//! +//! That is also why this needs a backend whose connection IDs we can choose. +//! Quinn and noq allow it; the quiche backend does not, and refuses to start +//! with workers rather than silently falling back to address hashing. use std::net::SocketAddr; @@ -44,9 +58,14 @@ pub struct RuntimeConfig { /// /// A connection is handled start to finish by one worker, which trades the /// shared runtime's load balancing for no cross-thread traffic per packet. - /// Linux-only, and mutually exclusive with `--listen-tls-generate`, since - /// each worker would otherwise generate and serve a certificate of its own. - /// Unset (the default) keeps QUIC on the shared runtime. + /// The group is steered by connection ID, so a client that changes address + /// stays with its worker. + /// + /// Linux-only, and needs a backend whose connection IDs can carry the + /// worker's identity (quinn or noq, not quiche). The listen address needs an + /// explicit non-zero port, and `--listen-tls-generate` is refused because + /// each worker would generate a certificate of its own. Unset (the default) + /// keeps QUIC on the shared runtime. #[arg(long = "runtime-workers", env = "MOQ_RUNTIME_WORKERS")] pub workers: Option, @@ -67,16 +86,34 @@ pub struct RuntimeConfig { /// when [`RuntimeConfig::workers`] is unset, which is how [`crate::Relay`] can /// hold one unconditionally. pub struct Workers { - bound: Vec, + workers: Vec, certificates: Option, addr: Option, } -/// One worker thread, bound and parked until it is handed the cluster to serve. -struct Bound { +/// One worker thread and the three channels that own its lifetime. +/// +/// Each is an `Option` because [`Workers`] hands them out or drops them at +/// different points, and the worker reads a dropped sender as an instruction: +/// no [`Self::start`] means never serve, no [`Self::stop`] means stop serving. +struct Worker { shard: u16, - start: tokio::sync::oneshot::Sender, - done: tokio::sync::oneshot::Receiver>, + + /// Joined by [`Workers::drop`]. Held rather than detached so the socket is + /// provably released before the pool is gone, instead of by a thread that + /// outlives it. + thread: std::thread::JoinHandle<()>, + + /// Sends the worker the cluster to accept into. Dropping it instead tells a + /// worker that never got started to exit. + start: Option>, + + /// Never sent on: dropping it is the stop signal, which is what makes the + /// accept loop cancellable from outside the worker's own runtime. + stop: Option>, + + /// The result of the worker's accept loop. + done: Option>>, } /// What a worker needs to start accepting, known only after the cluster exists. @@ -104,9 +141,9 @@ impl Workers { let pin = config.runtime.pin.unwrap_or(true).then(cores).unwrap_or_default(); - let mut bound = Vec::with_capacity(count as usize); + let mut workers = Vec::with_capacity(count as usize); let mut certificates = None; - let mut addr = None; + let mut addr: Option = None; for shard in 0..count { let listen = shard_config(config, shard, count)?; @@ -115,11 +152,14 @@ impl Workers { let (ready_tx, ready_rx) = std::sync::mpsc::channel(); let (start_tx, start_rx) = tokio::sync::oneshot::channel(); + let (stop_tx, stop_rx) = tokio::sync::oneshot::channel(); let (done_tx, done_rx) = tokio::sync::oneshot::channel(); - std::thread::Builder::new() + // Held, not detached: `Workers` is what releases these sockets, and it + // can only promise that if it can join the threads holding them. + let thread = std::thread::Builder::new() .name(format!("moq-quic-{shard}")) - .spawn(move || run(shard, core, listen, quic, ready_tx, start_rx, done_tx)) + .spawn(move || run(shard, core, listen, quic, ready_tx, start_rx, stop_rx, done_tx)) .with_context(|| format!("failed to spawn QUIC worker {shard}"))?; // A worker that fails to bind drops its sender, so the recv error and @@ -129,20 +169,34 @@ impl Workers { .map_err(|_| anyhow::anyhow!("QUIC worker {shard} exited before binding"))? .with_context(|| format!("QUIC worker {shard} failed to listen"))?; + // The group only balances a port every member actually holds. An + // ephemeral bind gives each worker a port of its own instead, leaving + // all but the first unreachable behind an address that looks bound. + match addr { + Some(first) if first != ready.addr => anyhow::bail!( + "QUIC worker {shard} bound {} instead of {first}; \ + every worker must share one port, so --listen needs an explicit non-zero port", + ready.addr + ), + Some(_) => {} + None => addr = Some(ready.addr), + } + certificates.get_or_insert(ready.certificates); - addr.get_or_insert(ready.addr); - bound.push(Bound { + workers.push(Worker { shard, - start: start_tx, - done: done_rx, + thread, + start: Some(start_tx), + stop: Some(stop_tx), + done: Some(done_rx), }); } tracing::info!(workers = count, pinned = !pin.is_empty(), "started QUIC workers"); Ok(Self { - bound, + workers, // Every worker loads the same certificate files, so any one of them // answers for the fingerprint endpoint. certificates, @@ -153,7 +207,7 @@ impl Workers { /// A pool that owns nothing, for a relay serving QUIC on the shared runtime. fn disabled() -> Self { Self { - bound: Vec::new(), + workers: Vec::new(), certificates: None, addr: None, } @@ -161,7 +215,7 @@ impl Workers { /// Whether QUIC is being served by workers rather than the shared runtime. pub fn enabled(&self) -> bool { - !self.bound.is_empty() + !self.workers.is_empty() } /// The address every worker is bound to, or `None` when they are disabled. @@ -180,13 +234,18 @@ impl Workers { /// Pends forever when the workers are disabled, so it composes into a /// `select!` either way. A worker only stops on error or on shutdown, so the /// first one to finish ends the relay the same way the shared accept loop does. - pub async fn serve(self, cluster: Cluster, auth: Auth, shutdown: Shutdown) -> anyhow::Result<()> { - if self.bound.is_empty() { + /// + /// Borrows rather than consuming, so the caller still holds the pool when this + /// resolves or is cancelled. That is what makes the sockets go away: dropping + /// the returned future stops nothing by itself, since each worker's accept loop + /// lives on a thread of its own. Dropping the [`Workers`] is what stops them. + pub async fn serve(&mut self, cluster: Cluster, auth: Auth, shutdown: Shutdown) -> anyhow::Result<()> { + if self.workers.is_empty() { return std::future::pending().await; } let mut running = futures::stream::FuturesUnordered::new(); - for worker in self.bound { + for worker in &mut self.workers { let start = Start { cluster: cluster.clone(), auth: auth.clone(), @@ -195,11 +254,19 @@ impl Workers { let shard = worker.shard; // The worker is parked on this receiver; it only drops if the thread // died, which its `done` half reports with the real error. - let _ = worker.start.send(start); + if let Some(sender) = worker.start.take() { + let _ = sender.send(start); + } + let done = worker.done.take(); running.push(async move { - match worker.done.await { - Ok(res) => res.with_context(|| format!("QUIC worker {shard} failed")), - Err(_) => Err(anyhow::anyhow!("QUIC worker {shard} exited without reporting")), + match done { + Some(done) => match done.await { + Ok(res) => res.with_context(|| format!("QUIC worker {shard} failed")), + Err(_) => Err(anyhow::anyhow!("QUIC worker {shard} exited without reporting")), + }, + // Already awaited by an earlier `serve`, so this worker has + // nothing left to report; leave the outcome to its siblings. + None => std::future::pending().await, } }); } @@ -212,6 +279,44 @@ impl Workers { } } +impl std::fmt::Debug for Workers { + /// Hand-written because the certificate handle behind [`Workers::certificates`] + /// is opaque, and the shard count plus the shared address is what identifies a + /// pool anyway. + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Workers") + .field("workers", &self.workers.len()) + .field("addr", &self.addr) + .finish() + } +} + +impl Drop for Workers { + /// Stop every worker and wait for its thread, so the sockets are released + /// before the pool is gone. + /// + /// Detaching instead would leave threads accepting into a cluster their owner + /// has finished with, and a replacement pool would join the same reuseport + /// group as the orphans and lose a share of its traffic to them. Blocking here + /// is deliberate: the stop signal goes out first, so each thread is already on + /// its way out by the time it is joined. + fn drop(&mut self) { + // Every signal first, then the joins: a worker cannot exit while it still + // holds a live sender, so joining as we go would serialize the teardown. + for worker in &mut self.workers { + worker.start.take(); + worker.stop.take(); + } + + for worker in self.workers.drain(..) { + let shard = worker.shard; + if worker.thread.join().is_err() { + tracing::error!(shard, "QUIC worker panicked"); + } + } + } +} + /// What a worker reports back once it is listening. struct Ready { certificates: moq_tokio::tls::Certificates, @@ -230,6 +335,7 @@ fn shard_config(config: &Config, index: u16, count: u16) -> anyhow::Result, @@ -237,6 +343,7 @@ fn run( quic: moq_tokio::quic::Config, ready: std::sync::mpsc::Sender>, start: tokio::sync::oneshot::Receiver, + stop: tokio::sync::oneshot::Receiver<()>, done: tokio::sync::oneshot::Sender>, ) { if let Some(core) = core { @@ -285,7 +392,16 @@ fn run( let Ok(start) = start.await else { return; }; - let _ = done.send(crate::serve(server, start.cluster, start.auth, start.shutdown).await); + + // `stop` resolves when its sender drops, which is the only way to get out + // of the accept loop: it runs on this thread, so the owner cancelling its + // own future would leave this one running. + tokio::select! { + res = crate::serve(server, start.cluster, start.auth, start.shutdown) => { + let _ = done.send(res); + } + _ = stop => tracing::debug!(shard, "QUIC worker stopping"), + } }); } diff --git a/rs/moq-relay/tests/runtime_workers.rs b/rs/moq-relay/tests/runtime_workers.rs index a1c60b419..845cd0826 100644 --- a/rs/moq-relay/tests/runtime_workers.rs +++ b/rs/moq-relay/tests/runtime_workers.rs @@ -8,7 +8,7 @@ use std::net::{SocketAddr, UdpSocket}; use std::time::Duration; -use moq_relay::{Config, PublicConfig, Relay}; +use moq_relay::{Config, PublicConfig, Relay, Workers}; use moq_tokio::moq_net::{self, Origin}; const TIMEOUT: Duration = Duration::from_secs(10); @@ -38,6 +38,23 @@ fn certificate(dir: &std::path::Path) -> (std::path::PathBuf, std::path::PathBuf (cert_path, key_path) } +/// A relay config serving QUIC from `workers` pinned-less worker threads. +/// +/// Pinning is off because a CI container may restrict which cores it may run on, +/// and none of these tests are about placement. +fn worker_config(cert: &std::path::Path, key: &std::path::Path, port: u16, workers: u16) -> Config { + let mut config = Config::default(); + config.listen.bind = Some(format!("127.0.0.1:{port}")); + config.listen.tls.cert = vec![cert.to_path_buf()]; + config.listen.tls.key = vec![key.to_path_buf()]; + config.runtime.workers = Some(workers); + config.runtime.pin = Some(false); + #[allow(deprecated)] + let public = PublicConfig::Simple(vec![String::new()]); + config.auth.public = Some(public); + config +} + fn client() -> moq_tokio::Client { let mut config = moq_tokio::connect::Config::default(); config.tls.insecure = Some(true); @@ -68,27 +85,19 @@ async fn workers_serve_quic_and_share_one_origin() { let (cert, key) = certificate(dir.path()); let port = free_udp_port(); - let mut config = Config::default(); - config.listen.bind = Some(format!("127.0.0.1:{port}")); - config.listen.tls.cert = vec![cert]; - config.listen.tls.key = vec![key]; - config.runtime.workers = Some(WORKERS); - // A CI container may restrict which cores it may run on, and pinning is not - // what this test is about. - config.runtime.pin = Some(false); - #[allow(deprecated)] - let public = PublicConfig::Simple(vec![String::new()]); - config.auth.public = Some(public); + let config = worker_config(&cert, &key, port, WORKERS); let relay = Relay::load(config).await.expect("load relay"); let expected: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap(); assert_eq!(relay.addr, Some(expected), "workers bound a different address"); - let workers = tokio::spawn( - relay - .workers - .serve(relay.cluster.clone(), relay.auth.clone(), relay.shutdown.clone()), - ); + let cluster = relay.cluster.clone(); + let auth = relay.auth.clone(); + let shutdown = relay.shutdown.clone(); + // Moved into the task rather than borrowed: the pool has to outlive the accept + // loops it owns, and dropping it is what stops them. + let mut pool = relay.workers; + let workers = tokio::spawn(async move { pool.serve(cluster, auth, shutdown).await }); let url: url::Url = format!("https://127.0.0.1:{port}/workers").parse().expect("parse url"); @@ -150,3 +159,69 @@ async fn workers_serve_quic_and_share_one_origin() { drop(subscribers); workers.abort(); } + +/// A worker pool that outlives its owner is a leak with teeth: the threads keep +/// accepting into a cluster nobody is driving, and a replacement pool joins the +/// same reuseport group as the orphans and loses a share of its traffic to them. +/// Dropping the pool has to actually release the port. +#[tokio::test] +async fn dropping_the_pool_releases_the_port() { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let dir = tempfile::tempdir().expect("tempdir"); + let (cert, key) = certificate(dir.path()); + let port = free_udp_port(); + let config = worker_config(&cert, &key, port, WORKERS); + + let relay = Relay::load(config).await.expect("load relay"); + let cluster = relay.cluster.clone(); + let auth = relay.auth.clone(); + let shutdown = relay.shutdown.clone(); + let mut pool = relay.workers; + + // Serving before dropping is the case that used to strand the threads: `serve` + // consumed the pool, so nothing was left that could stop them. + let serving = tokio::spawn(async move { + let _ = pool.serve(cluster, auth, shutdown).await; + }); + tokio::time::sleep(Duration::from_millis(50)).await; + serving.abort(); + let _ = serving.await; + + // A plain bind refuses a port any socket still holds, reuseport or not, so this + // succeeds only if every worker's socket is really gone. + let addr: SocketAddr = format!("127.0.0.1:{port}").parse().unwrap(); + moq_tokio::bind::udp(moq_tokio::bind::Udp::new(addr)).expect("workers left the port bound"); +} + +/// An ephemeral bind gives every worker a port of its own instead of a shared +/// one, leaving all but the first unreachable behind an address that reads as +/// bound. That has to fail at startup rather than come up looking healthy. +#[tokio::test] +async fn an_ephemeral_port_is_refused() { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let dir = tempfile::tempdir().expect("tempdir"); + let (cert, key) = certificate(dir.path()); + let config = worker_config(&cert, &key, 0, WORKERS); + + let err = Workers::bind(&config).expect_err("an ephemeral port cannot be shared"); + let err = format!("{err:#}"); + assert!( + err.contains("every worker must share one port"), + "unexpected error: {err}" + ); +} + +/// One worker has no group to disagree with, so an ephemeral port is fine there. +#[tokio::test] +async fn a_single_worker_may_use_an_ephemeral_port() { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let dir = tempfile::tempdir().expect("tempdir"); + let (cert, key) = certificate(dir.path()); + let config = worker_config(&cert, &key, 0, 1); + + let workers = Workers::bind(&config).expect("a lone worker may take any port"); + assert!(workers.local_addr().is_some_and(|addr| addr.port() != 0)); +} diff --git a/rs/moq-tokio/src/lib.rs b/rs/moq-tokio/src/lib.rs index d04d8969c..c85b5360a 100644 --- a/rs/moq-tokio/src/lib.rs +++ b/rs/moq-tokio/src/lib.rs @@ -35,6 +35,7 @@ pub mod quinn; #[cfg(any(feature = "quinn", feature = "noq", feature = "quiche", feature = "tcp"))] mod resolve; mod server; +mod steer; #[cfg(feature = "tcp")] pub mod tcp; pub mod tls; diff --git a/rs/moq-tokio/src/noq.rs b/rs/moq-tokio/src/noq.rs index ac48b4966..78c01cbc7 100644 --- a/rs/moq-tokio/src/noq.rs +++ b/rs/moq-tokio/src/noq.rs @@ -552,8 +552,7 @@ impl NoqServer { })); } - let socket = crate::bind::udp(crate::bind::Udp::new(listen).with_reuse_port(config.shard.is_some())) - .map_err(Error::BindSocket)?; + let socket = crate::steer::bind(listen, config.shard).map_err(Error::BindSocket)?; // Create the generic QUIC endpoint. let quic = noq::Endpoint::new(endpoint_config, Some(tls), socket, runtime).map_err(Error::CreateEndpoint)?; @@ -652,6 +651,41 @@ pub(crate) async fn accept( } } +// ── ShardIdGenerator ──────────────────────────────────────────────── + +/// Connection IDs whose first byte names the reuseport member that owns them, +/// so the group's steering filter can route later packets back to it. +struct ShardIdGenerator { + shard: crate::listen::Shard, +} + +impl ShardIdGenerator { + /// Only the first byte is spoken for; the rest stays random. + const LEN: usize = 8; + + fn new(shard: crate::listen::Shard) -> Self { + Self { shard } + } +} + +impl noq::ConnectionIdGenerator for ShardIdGenerator { + fn generate_cid(&mut self) -> noq::ConnectionId { + use rand::RngExt; + let mut cid = Vec::with_capacity(Self::LEN); + cid.push(crate::steer::cid_prefix(self.shard)); + cid.extend(rand::rng().random_iter::().take(Self::LEN - 1)); + noq::ConnectionId::new(cid.as_slice()) + } + + fn cid_len(&self) -> usize { + Self::LEN + } + + fn cid_lifetime(&self) -> Option { + None + } +} + // ── ServerIdGenerator ─────────────────────────────────────────────── struct ServerIdGenerator { diff --git a/rs/moq-tokio/src/quiche.rs b/rs/moq-tokio/src/quiche.rs index fb394a792..67396db3f 100644 --- a/rs/moq-tokio/src/quiche.rs +++ b/rs/moq-tokio/src/quiche.rs @@ -96,6 +96,11 @@ pub enum Error { #[error("failed to get local address")] NoLocalAddr, + /// This backend cannot encode the owning worker in its connection IDs, so a + /// reuseport group built on it would have no way to steer packets back. + #[error("the quiche backend cannot serve per-core workers; use the quinn backend")] + ShardUnsupported, + /// The server was given neither a certificate pair nor hostnames to generate one from. #[error("--tls-cert and --tls-key are required with the quiche backend")] CertRequired, @@ -562,11 +567,19 @@ impl QuicheServer { tracing::warn!("QUIC-LB is not supported with the quiche backend; ignoring server ID"); } + // Refused rather than degraded: without connection-ID steering the kernel + // falls back to hashing addresses, and a client that changes address lands + // on a worker that has never seen its connection. `web-transport-quiche` + // exposes no connection ID hook to encode the owner with. + if config.shard.is_some() { + return Err(Error::ShardUnsupported); + } + let quic = quic.resolve(); let listen = crate::util::resolve(config.bind.as_deref(), crate::server::DEFAULT_BIND).map_err(Error::ResolveBind)?; - let socket = crate::bind::udp(crate::bind::Udp::new(listen).with_reuse_port(config.shard.is_some()))?; + let socket = crate::bind::udp(crate::bind::Udp::new(listen))?; let (chain, key) = if !config.tls.generate.is_empty() { generate_quiche_cert(&config.tls.generate)? diff --git a/rs/moq-tokio/src/quinn.rs b/rs/moq-tokio/src/quinn.rs index bceee64b5..a0a56db8f 100644 --- a/rs/moq-tokio/src/quinn.rs +++ b/rs/moq-tokio/src/quinn.rs @@ -182,6 +182,11 @@ pub enum Error { #[error("connection ID length ({0}) exceeds maximum of 20")] QuicLbCidTooLong(usize), + /// Both QUIC-LB and reuseport steering want to own the connection ID, and + /// each reads back only its own encoding. + #[error("QUIC-LB connection IDs cannot be combined with per-core workers")] + ShardWithQuicLb, + /// The mTLS client verifier couldn't be built from the configured roots. #[error("failed to build client certificate verifier")] ClientVerifier(#[source] rustls::server::VerifierBuilderError), @@ -544,7 +549,17 @@ impl QuinnServer { // Configure connection ID generator with server ID if provided let mut endpoint_config = quinn::EndpointConfig::default(); - if let Some(server_id) = config.lb_id { + if let Some(shard) = config.shard { + if config.lb_id.is_some() { + return Err(Error::ShardWithQuicLb); + } + tracing::debug!( + index = shard.index(), + count = shard.count(), + "encoding the shard in connection IDs" + ); + endpoint_config.cid_generator(move || Box::new(ShardIdGenerator::new(shard))); + } else if let Some(server_id) = config.lb_id { let nonce_len = config.lb_nonce.unwrap_or(8); if nonce_len < 4 { return Err(Error::QuicLbNonceTooSmall); @@ -563,8 +578,7 @@ impl QuinnServer { endpoint_config.cid_generator(move || Box::new(ServerIdGenerator::new(server_id.clone(), nonce_len))); } - let socket = crate::bind::udp(crate::bind::Udp::new(listen).with_reuse_port(config.shard.is_some())) - .map_err(Error::BindSocket)?; + let socket = crate::steer::bind(listen, config.shard).map_err(Error::BindSocket)?; // Create the generic QUIC endpoint. let quic = quinn::Endpoint::new(endpoint_config, Some(tls), socket, runtime).map_err(Error::CreateEndpoint)?; @@ -664,6 +678,41 @@ pub(crate) async fn accept( } } +// ── ShardIdGenerator ──────────────────────────────────────────────── + +/// Connection IDs whose first byte names the reuseport member that owns them, +/// so the group's steering filter can route later packets back to it. +struct ShardIdGenerator { + shard: crate::listen::Shard, +} + +impl ShardIdGenerator { + /// Matches quinn's own default length. Only the first byte is spoken for. + const LEN: usize = 8; + + fn new(shard: crate::listen::Shard) -> Self { + Self { shard } + } +} + +impl quinn::ConnectionIdGenerator for ShardIdGenerator { + fn generate_cid(&mut self) -> quinn::ConnectionId { + use rand::RngExt; + let mut cid = Vec::with_capacity(Self::LEN); + cid.push(crate::steer::cid_prefix(self.shard)); + cid.extend(rand::rng().random_iter::().take(Self::LEN - 1)); + quinn::ConnectionId::new(&cid) + } + + fn cid_len(&self) -> usize { + Self::LEN + } + + fn cid_lifetime(&self) -> Option { + None + } +} + // ── ServerIdGenerator ─────────────────────────────────────────────── struct ServerIdGenerator { diff --git a/rs/moq-tokio/src/steer.rs b/rs/moq-tokio/src/steer.rs new file mode 100644 index 000000000..12943a6e0 --- /dev/null +++ b/rs/moq-tokio/src/steer.rs @@ -0,0 +1,284 @@ +//! Steering a `SO_REUSEPORT` group by QUIC connection ID. +//! +//! A reuseport group left to itself is picked by hashing the packet's 4-tuple, +//! which is wrong for QUIC: a connection is identified by its connection ID, not +//! its address, so a client that changes address (a NAT rebinding, a network +//! change, or plain connection migration) hashes to a different member and its +//! packets arrive at a socket that has never heard of it. The connection dies +//! instead of migrating. +//! +//! The fix is the standard one: the server chooses connection IDs that say which +//! member owns them, and a filter on the group reads that back. +//! +//! # The encoding +//! +//! [`cid_prefix`] makes the first byte of every connection ID this member issues +//! congruent to its index modulo the group size, keeping the remaining +//! `256 / count` values of that byte (and every later byte) random. The kernel +//! selects `socks[A % count]` for a returned accumulator `A`, so feeding it that +//! byte lands on the owner. +//! +//! A client's first packets carry a connection ID the *client* invented, which +//! encodes nothing. Those hash to an arbitrary member, and that is fine: whoever +//! receives the Initial owns the connection and issues IDs carrying its own +//! index, so every later packet steers to it. Retransmitted Initials repeat the +//! same client-chosen ID and so keep reaching the same member, which is what +//! stops a retry from starting a second handshake elsewhere. +//! +//! # Why classic BPF +//! +//! `SO_ATTACH_REUSEPORT_CBPF` runs the program with the UDP payload at offset 0 +//! (the kernel pulls the header off first) and uses the return value as the +//! index. That is the whole rule in seven instructions, with no `CAP_BPF`, no +//! BPF toolchain in the build, and no map to keep alive across restarts. The +//! cost is that the kernel selects by *position*, so the group has to be built +//! once, in order, and never resized. [`bind`] is what guarantees that. + +use std::io; +use std::net::{SocketAddr, UdpSocket}; + +use crate::listen::Shard; + +/// Bind this member's socket into the group, steering the whole group once the +/// last member has joined. +/// +/// Call it for every member in index order and nothing else in between: the +/// kernel identifies a member by its position in the group, which is the order +/// the sockets bound. An unsharded bind is the plain one. +pub(crate) fn bind(addr: SocketAddr, shard: Option) -> io::Result { + let options = crate::bind::Udp::new(addr).with_reuse_port(shard.is_some()); + let socket = crate::bind::udp(options)?; + + // The filter covers the group, not the socket, so it goes on once everyone is + // in. Attaching earlier would steer by an index range that is still growing, + // and the members that joined later would be unreachable until it was redone. + if let Some(shard) = shard.filter(|shard| shard.index() + 1 == shard.count()) { + attach(&socket, shard.count())?; + } + + Ok(socket) +} + +/// The first byte of a connection ID issued by `shard`. +/// +/// `count` values of the byte are reserved per member, so this keeps +/// `256 / count` of its randomness: 32 values across 8 members, and the other 19 +/// bytes of the ID stay fully random either way. Connection IDs are not secrets, +/// but they are unlinkable only while they look random, which is why this spends +/// as little of that byte as it can. +pub(crate) fn cid_prefix(shard: Shard) -> u8 { + use rand::RngExt; + + let count = u32::from(shard.count()); + // The number of whole strides of `count` that fit in a byte. At least one, + // since a group cannot have more than 256 members. + let strides = 256 / count; + let stride = rand::rng().random_range(0..strides); + + (stride * count + u32::from(shard.index())) as u8 +} + +/// Attach the steering filter to a socket in a reuseport group of `count`. +/// +/// Applies to the group, so once is enough, and the last member to bind is the +/// first point where that is true. +#[cfg(target_os = "linux")] +fn attach(socket: &UdpSocket, count: u16) -> io::Result<()> { + use std::os::fd::AsRawFd; + + let program = program(count); + let fprog = libc::sock_fprog { + len: program.len() as u16, + filter: program.as_ptr() as *mut libc::sock_filter, + }; + + // SAFETY: `fprog` points at `program`, which outlives the call, and its `len` + // is that slice's length. + let res = unsafe { + libc::setsockopt( + socket.as_raw_fd(), + libc::SOL_SOCKET, + libc::SO_ATTACH_REUSEPORT_CBPF, + std::ptr::from_ref(&fprog).cast(), + size_of::() as libc::socklen_t, + ) + }; + + if res != 0 { + return Err(io::Error::last_os_error()); + } + + tracing::debug!(count, "steering the reuseport group by connection ID"); + Ok(()) +} + +#[cfg(not(target_os = "linux"))] +fn attach(_socket: &UdpSocket, _count: u16) -> io::Result<()> { + // Unreachable in practice: a shard only exists once `bind::udp` accepted + // `SO_REUSEPORT`, which is Linux-only. + Err(io::Error::new( + io::ErrorKind::Unsupported, + "reuseport steering is Linux-only", + )) +} + +/// The classic-BPF program that maps a QUIC packet to its owner's index. +/// +/// Reads the first byte of the destination connection ID, whose offset depends +/// on the header form, and reduces it modulo the group size. The packet starts +/// at offset 0 because the kernel pulls the UDP header off before running this. +/// +/// A long header carries an explicit ID length, and a zero there would make +/// offset 6 the *source* ID's length byte instead. Not worth a branch: only a +/// malformed or hostile packet gets there, since a client's Initial must carry a +/// destination ID of at least 8 bytes, and the worst outcome is that one junk +/// packet reaches a member that discards it. +#[cfg(target_os = "linux")] +fn program(count: u16) -> [libc::sock_filter; 7] { + /// Long-header form bit in the first byte of a QUIC packet. + const LONG_HEADER: u32 = 0x80; + /// 1 first byte + 4 version bytes + 1 length byte precede a long header's + /// destination connection ID. + const LONG_DCID: u32 = 6; + /// A short header's destination connection ID follows the first byte. + const SHORT_DCID: u32 = 1; + + fn insn(code: u32, jt: u8, jf: u8, k: u32) -> libc::sock_filter { + libc::sock_filter { + code: code as u16, + jt, + jf, + k, + } + } + + [ + // A = packet[0] + insn(libc::BPF_LD | libc::BPF_B | libc::BPF_ABS, 0, 0, 0), + // Long header? Skip the short-header load. + insn(libc::BPF_JMP | libc::BPF_JSET | libc::BPF_K, 2, 0, LONG_HEADER), + // A = packet[1] + insn(libc::BPF_LD | libc::BPF_B | libc::BPF_ABS, 0, 0, SHORT_DCID), + // Skip the long-header load. + insn(libc::BPF_JMP | libc::BPF_JA | libc::BPF_K, 0, 0, 1), + // A = packet[6] + insn(libc::BPF_LD | libc::BPF_B | libc::BPF_ABS, 0, 0, LONG_DCID), + // A %= count. The kernel would reduce out-of-range indices to a 4-tuple + // hash rather than an index, so the program has to land in range itself. + insn(libc::BPF_ALU | libc::BPF_MOD | libc::BPF_K, 0, 0, u32::from(count)), + insn(libc::BPF_RET | libc::BPF_A, 0, 0, 0), + ] +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The encoding's only real requirement: whatever byte a member issues has to + /// reduce back to that member. + #[test] + fn a_prefix_reduces_to_its_own_shard() { + for count in [1u16, 2, 3, 4, 7, 8, 16, 64, 255] { + for index in 0..count { + let shard = Shard::new(index, count).unwrap(); + // Sampled, since the prefix is randomized within the member's stride. + for _ in 0..64 { + let prefix = cid_prefix(shard); + assert_eq!( + u16::from(prefix) % count, + index, + "prefix {prefix} of shard {index}/{count} steers elsewhere" + ); + } + } + } + } + + /// Every member of a group has to be reachable, or the ones that aren't sit + /// idle while their share of the traffic goes to a sibling. + #[test] + fn prefixes_cover_every_shard() { + const COUNT: u16 = 8; + let mut seen = std::collections::HashSet::new(); + for index in 0..COUNT { + let shard = Shard::new(index, COUNT).unwrap(); + for _ in 0..256 { + seen.insert(u16::from(cid_prefix(shard)) % COUNT); + } + } + assert_eq!(seen.len(), usize::from(COUNT)); + } + + /// The whole point, end to end against a real kernel: a packet carrying a + /// member's connection ID has to reach that member and no other. + /// + /// This is what the 4-tuple hash cannot do. Every packet here is sent from a + /// different ephemeral source port, so under the default hashing the landing + /// socket would be unrelated to the connection ID. + #[test] + #[cfg(target_os = "linux")] + fn a_connection_id_reaches_its_own_member() { + const COUNT: u16 = 4; + + // Bound in index order, because the kernel identifies a member by its + // position in the group. + let mut group = Vec::new(); + let mut addr: SocketAddr = "127.0.0.1:0".parse().unwrap(); + for index in 0..COUNT { + let shard = Shard::new(index, COUNT).unwrap(); + let socket = bind(addr, Some(shard)).unwrap(); + addr = socket.local_addr().unwrap(); + socket.set_nonblocking(true).unwrap(); + group.push(socket); + } + + /// Which member received something, given every send targets exactly one. + fn receiver(group: &[UdpSocket]) -> Option { + let mut buf = [0u8; 64]; + group.iter().position(|socket| socket.recv_from(&mut buf).is_ok()) + } + + for index in 0..COUNT { + let prefix = cid_prefix(Shard::new(index, COUNT).unwrap()); + + // A short header: form bit clear, connection ID straight after. + let short = [0x40, prefix, 1, 2, 3, 4, 5, 6, 7, 8]; + // A long header: form bit set, then version, then the ID's length. + let long = [0xc0, 0, 0, 0, 1, 8, prefix, 1, 2, 3, 4, 5, 6, 7]; + + for (form, packet) in [("short", &short[..]), ("long", &long[..])] { + // A fresh source port per packet, so nothing but the connection ID + // can be steering this. + let sender = crate::bind::udp(crate::bind::Udp::new("127.0.0.1:0".parse().unwrap())).unwrap(); + sender.send_to(packet, addr).unwrap(); + + // The send is local and the filter is synchronous, but the receive + // still has to be given a moment to land. + std::thread::sleep(std::time::Duration::from_millis(50)); + assert_eq!( + receiver(&group), + Some(usize::from(index)), + "{form} header for member {index} landed on the wrong socket" + ); + } + } + } + + /// The jumps are offsets from the *following* instruction, so an off-by-one + /// silently reads the wrong byte rather than failing to load. + #[test] + #[cfg(target_os = "linux")] + fn the_program_branches_to_the_right_loads() { + let program = program(4); + assert_eq!(program.len(), 7); + + // Instruction 1 falls through to the short-header load at 2, and jumps + // over it (and its trailing jump) to the long-header load at 4. + assert_eq!(program[1].jt, 2, "long header must land on the long load"); + assert_eq!(program[1].jf, 0, "short header must fall through"); + assert_eq!(program[2].k, 1, "short header reads the byte after the first"); + assert_eq!(program[3].k, 1, "the short path must skip the long load"); + assert_eq!(program[4].k, 6, "long header reads past version and length"); + assert_eq!(program[5].k, 4, "the modulus is the group size"); + } +}