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..6b5f16ff2 100644 --- a/doc/bin/relay/config.md +++ b/doc/bin/relay/config.md @@ -45,6 +45,47 @@ 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. +# +# 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. +# +# 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. +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..ef0dd698f 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,12 @@ 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 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`]. @@ -73,6 +79,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 +100,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 +163,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 +195,7 @@ impl Relay { addr, shutdown, shutdown_trigger, + workers, }) } @@ -187,6 +213,7 @@ impl Relay { web, shutdown, shutdown_trigger, + mut workers, .. } = self; @@ -209,7 +236,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..e78134307 --- /dev/null +++ b/rs/moq-relay/src/runtime.rs @@ -0,0 +1,432 @@ +//! 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) 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 +//! +//! 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; + +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. + /// 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, + + /// 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 { + workers: Vec, + certificates: Option, + addr: Option, +} + +/// 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, + + /// 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. +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 workers = Vec::with_capacity(count as usize); + let mut certificates = None; + let mut addr: Option = 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 (stop_tx, stop_rx) = tokio::sync::oneshot::channel(); + let (done_tx, done_rx) = tokio::sync::oneshot::channel(); + + // 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, 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 + // 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"))?; + + // 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); + + workers.push(Worker { + shard, + 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 { + workers, + // 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 { + workers: Vec::new(), + certificates: None, + addr: None, + } + } + + /// Whether QUIC is being served by workers rather than the shared runtime. + pub fn enabled(&self) -> bool { + !self.workers.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. + /// + /// 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 &mut self.workers { + 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. + if let Some(sender) = worker.start.take() { + let _ = sender.send(start); + } + let done = worker.done.take(); + running.push(async move { + 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, + } + }); + } + + use futures::StreamExt; + running + .next() + .await + .unwrap_or_else(|| Err(anyhow::anyhow!("no QUIC 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, + 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. +#[allow(clippy::too_many_arguments)] +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, + stop: 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; + }; + + // `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"), + } + }); +} + +/// 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..845cd0826 --- /dev/null +++ b/rs/moq-relay/tests/runtime_workers.rs @@ -0,0 +1,227 @@ +//! 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, Workers}; +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) +} + +/// 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); + 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 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 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"); + + // ── 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(); +} + +/// 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/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/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/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..78c01cbc7 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,7 @@ impl NoqServer { })); } - let socket = crate::bind::udp(listen).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)?; @@ -651,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 4d2b4dd45..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, @@ -322,7 +327,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 +364,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 @@ -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(listen)?; + 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 0a0dad25b..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), @@ -269,7 +274,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(); @@ -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,7 +578,7 @@ 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::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)?; @@ -663,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/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. 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"); + } +}