diff --git a/Cargo.toml b/Cargo.toml index d2b2ed8..b753544 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "endpoint-libs" -version = "3.3.0" +version = "3.3.1" edition = "2024" authors = ["Veon "] description = "Launch MCP services fast: describe endpoints once in RON, and endpoint-gen generates the Rust models, docs and MCP tool schemas that this crate serves over WebSocket RPC, with roles and typed errors built in." @@ -203,7 +203,7 @@ alloy-primitives = { version = "1.3.0", features = [ # 0.1.11 is the release that carries `signal::Signal`, `resolve` and # `connect_any`, which is what the tokio removal needed and what 0.1.9 did not # have. A caret so the family moves on its own: 0.1.11 is the floor, not a pin. -nagoya = { version = "^0.1.11", default-features = false, features = ["reactor"], optional = true } +nagoya = { version = "^0.1.13", default-features = false, features = ["reactor"], optional = true } # The WebSocket protocol itself: the RFC 6455 handshake and framing over a # nagoya socket, replacing hyper on the serving side and tokio-tungstenite on # the connecting side. `default-features = false` is load bearing: the default diff --git a/src/libs/ws/listener.rs b/src/libs/ws/listener.rs index c8f33d7..e136f16 100644 --- a/src/libs/ws/listener.rs +++ b/src/libs/ws/listener.rs @@ -56,6 +56,13 @@ impl TcpListener { Ok(Self { listener }) } + /// [`Self::bind`] with `SO_REUSEPORT`, so each shard's reactor can hold its own + /// listener on the same port. Every listener on the port has to use this. + pub fn bind_shared(addr: Addr, handle: &Handle) -> Result { + let listener = nagoya::net::TcpListener::bind_shared(addr, handle)?; + Ok(Self { listener }) + } + /// Bind the first address in `addrs` that the kernel accepts. /// /// [`nagoya::reactor::resolve`] returns *every* address a name has, and a host @@ -75,9 +82,22 @@ impl TcpListener { /// attached: unlike [`nagoya::resolve::connect_any`], every error here is about /// this machine, so there is no remote-versus-local distinction to preserve. pub fn bind_any(addrs: &[Addr], handle: &Handle) -> Result { + Self::bind_any_with(addrs, handle, Self::bind) + } + + /// [`Self::bind_any`], binding each candidate with [`Self::bind_shared`]. + pub fn bind_any_shared(addrs: &[Addr], handle: &Handle) -> Result { + Self::bind_any_with(addrs, handle, Self::bind_shared) + } + + fn bind_any_with( + addrs: &[Addr], + handle: &Handle, + bind: fn(Addr, &Handle) -> Result, + ) -> Result { let mut last = None; for addr in addrs { - match Self::bind(*addr, handle) { + match bind(*addr, handle) { Ok(listener) => { tracing::debug!( ws_server = true, diff --git a/src/libs/ws/server.rs b/src/libs/ws/server.rs index 5c0fed4..493e1be 100644 --- a/src/libs/ws/server.rs +++ b/src/libs/ws/server.rs @@ -18,6 +18,7 @@ use tracing::*; use crate::libs::error_code::ErrorCode; use crate::libs::handler::{RequestHandler, RequestHandlerErased}; use crate::libs::peer::{Extensions, PeerIdentity}; +use crate::libs::signal::Shutdown; use crate::libs::toolbox::{ArcToolbox, RequestContext, TOOLBOX, Toolbox}; use crate::libs::utils::{get_conn_id, get_log_id}; #[cfg(feature = "ws")] @@ -416,19 +417,18 @@ impl WebsocketServer { /// registered with one specific reactor at birth. So the order is: resolve, /// create the reactor, bind, and only then start polling. /// - /// # One thread, for now + /// # Shards /// - /// The shard-per-core fan-out is gone, and this is the honest state of the port - /// rather than a simplification. A socket accepted by a nagoya listener is - /// registered with that listener's reactor, and only that reactor will ever - /// report its readiness, so the connection must be driven on the thread polling - /// it. Handing accepted sockets to other threads would put every connection's - /// readiness through one poller and add a cross-thread wake per message — - /// exactly the cost `Reactor::sharded` exists to avoid. The shape that does work - /// is a listening socket per reactor, and nagoya cannot yet produce one: its - /// `TcpListener` has no `try_clone` and no `SO_REUSEPORT` bind, and a bound - /// listener cannot be adopted by a second reactor. `shard_count` still reads the - /// operator's intent, and the server says so when it cannot honour it. + /// A socket accepted by a nagoya listener is registered with that listener's + /// reactor, and only that reactor will ever report its readiness, so the + /// connection must be driven on the thread polling it. Handing accepted sockets + /// to other threads would put every connection's readiness through one poller + /// and add a cross-thread wake per message. So each shard is a thread with its + /// own reactor and its own listening socket, all bound to one port with + /// `SO_REUSEPORT`, and the kernel picks the shard per connection. This thread + /// is the first shard. [`shard_count`] reads `WS_SHARDS`, then the cgroup CPU + /// quota, then the CPU count. Linux balances across the sockets; macOS accepts + /// the binds without balancing. /// /// # Signals /// @@ -470,21 +470,81 @@ impl WebsocketServer { let addrs = resolve(host, port) .wrap_err_with(|| format!("Failed to lookup host to bind: {}", self.config.address))?; - // No thread of its own: `block_on_with` below waits for readiness and polls - // the futures on this thread, so the accept loop, the signal waiters and - // every connection share one reactor. The signal waiters in particular only - // fire while their reactor is polled, which is why they are created inside - // this reactor's future rather than anywhere else. + // This thread is the first shard: `block_on_with` below waits for readiness + // and polls the futures on this thread, so its accept loop, the signal + // waiters and its connections share one reactor. The signal waiters in + // particular only fire while their reactor is polled, which is why they are + // created inside this reactor's future rather than anywhere else. let reactor = Reactor::local()?; let handle = reactor.handle(); - let listener = TcpListener::bind_any(&addrs, &handle)?; + let shards = shard_count(); + if shards <= 1 { + let listener = TcpListener::bind_any(&addrs, &handle)?; + return block_on_with(&reactor, async move { + self.listen_impl(Arc::new(listener), &handle, stop).await + }); + } + + // One listening socket per shard, each on its own reactor and thread, all + // bound to one port with SO_REUSEPORT. The kernel gives each connection to + // one socket, so a connection is accepted, registered and served by the same + // reactor and never crosses a thread. The first bind picks the address; the + // others take it exactly, which matters when the port asked for was zero. + let listener = TcpListener::bind_any_shared(&addrs, &handle)?; + let addr = listener.local_addr()?; + let (this, states) = self.prepare(); + let shutdown = Arc::new(Shutdown::default()); + let mut threads = Vec::with_capacity(shards - 1); + for shard in 1..shards { + let this = Arc::clone(&this); + let states = Arc::clone(&states); + let stop = Arc::clone(&shutdown); + let spawned = std::thread::Builder::new() + .name(format!("ws-shard-{shard}")) + .spawn(move || { + let served = (|| { + let reactor = Reactor::local()?; + let handle = reactor.handle(); + let listener = TcpListener::bind_shared(addr, &handle)?; + let stop = async move { stop.cancelled().await }; + block_on_with( + &reactor, + Self::serve_shard(this, states, Arc::new(listener), &handle, Some(stop)), + ) + })(); + if let Err(err) = &served { + error!(ws_server = true, shard, error = %err, "shard stopped with an error"); + } + served + }); + match spawned { + Ok(thread) => threads.push(thread), + Err(err) => { + shutdown.cancel(); + for thread in threads { + let _ = thread.join(); + } + return Err(eyre::eyre!("failed to start shard {shard}: {err}")); + } + } + } + info!( + ws_server = true, + shards, "serving on {shards} reactor threads" + ); - // One path. `listen_tls`, the `TlsListener` it built and the - // `config.insecure` branch that chose between them are gone: this server - // serves plain `ws://` and nothing else, with TLS terminated at the edge. - block_on_with(&reactor, async move { - self.listen_impl(Arc::new(listener), &handle, stop).await - }) + let served = block_on_with(&reactor, async { + let served = Self::serve_shard(this, states, Arc::new(listener), &handle, stop).await; + // Whatever stopped this shard, a signal or an error, stops them all. + shutdown.cancel(); + served + }); + for thread in threads { + if thread.join().is_err() { + error!(ws_server = true, "a shard thread panicked"); + } + } + served } /// Refuse to start when the configuration expects TLS this server cannot serve. @@ -510,9 +570,9 @@ impl WebsocketServer { /// Accept, handshake and serve, all on the reactor `handle` names. /// - /// One loop, three things to wait on, and no channel between them: the accept - /// fan-out and its per-shard `mpsc` are gone with the shard threads, because - /// there is nowhere to hand a connection to. See [`Self::listen`] for why. + /// One loop, three things to wait on, and no channel between them: a shard + /// serves what its own listener accepts, so nothing is handed between threads. + /// See [`Self::listen`] for how shards share a port. /// /// Connections are held in a `FuturesUnordered` and polled in place, the same /// model `serve_with` uses. nagoya's `TaskSet` is the purpose-built runner for @@ -532,10 +592,15 @@ impl WebsocketServer { T: ConnectionListener + 'static, F: std::future::Future + 'static, { - use futures::StreamExt; - use futures::future::{Either, FutureExt, select}; - use futures::stream::FuturesUnordered; + let (this, states) = self.prepare(); + Self::serve_shard(this, states, listener, handle, stop).await + } + /// Share the server and one connection table across every shard. + /// + /// Once per server, not per shard: the toolbox reaches a connection through + /// these states whichever reactor it lives on. + fn prepare(self) -> (Arc, Arc) { let states = Arc::new(WebsocketStates::new()); let this = Arc::new(self); this.toolbox.set_ws_states( @@ -543,16 +608,24 @@ impl WebsocketServer { this.config.header_only, this.config.drop_conn_on_buffer_full, ); + (this, states) + } - let num_shards = shard_count(); - if num_shards > 1 { - warn!( - ws_server = true, - "{} shards were asked for; serving on one thread because a nagoya \ - listener cannot be shared between reactors yet", - num_shards - ); - } + /// The accept loop of one shard, on the reactor `handle` names. + async fn serve_shard( + this: Arc, + states: Arc, + listener: Arc, + handle: &Handle, + stop: Option, + ) -> Result<()> + where + T: ConnectionListener + 'static, + F: std::future::Future + 'static, + { + use futures::StreamExt; + use futures::future::{Either, FutureExt, select}; + use futures::stream::FuturesUnordered; // The caller's stop, or the process signals when it gave none. Signals are // registered on this reactor, and awaited on it below: a `Signal` only diff --git a/tests/sharded_listen.rs b/tests/sharded_listen.rs new file mode 100644 index 0000000..96e41b8 --- /dev/null +++ b/tests/sharded_listen.rs @@ -0,0 +1,160 @@ +//! `listen_until` with several shards, over real TCP. +//! +//! Each shard is a thread with its own reactor and its own `SO_REUSEPORT` +//! listener on one port. This drives a server with three of them: every client +//! must complete a handshake and a request whichever shard the kernel picks, +//! and one stop must bring every shard down so `listen_until` returns. + +#![cfg(feature = "ws-client")] + +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use endpoint_libs::libs::handler::{RequestHandler, Response}; +use endpoint_libs::libs::toolbox::{ArcToolbox, CustomError, RequestContext}; +use endpoint_libs::libs::ws::{ + AuthController, WebsocketServer, WsClientBuilder, WsConnection, WsRequest, WsResponse, + WsServerConfig, +}; +use eyre::Result; +use futures::FutureExt; +use futures::future::LocalBoxFuture; +use nagoya::reactor::{Reactor, block_on_with}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct EchoRequest { + pub message: String, +} + +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct EchoResponse { + pub message: String, +} + +impl WsRequest for EchoRequest { + type Response = EchoResponse; + const METHOD_ID: u32 = 1; + const ROLES: &'static [u32] = &[1]; + const SCHEMA: &'static str = r#"{ + "name": "Echo", + "code": 1, + "parameters": [{"name": "message", "ty": "String"}], + "returns": [{"name": "message", "ty": "String"}], + "description": "Echoes the message back.", + "roles": [] + }"#; +} + +impl WsResponse for EchoResponse { + type Request = EchoRequest; +} + +struct MethodEcho; + +#[async_trait(?Send)] +impl RequestHandler for MethodEcho { + type Request = EchoRequest; + type Error = CustomError; + + async fn handle(&self, _ctx: RequestContext, req: EchoRequest) -> Response { + Ok(EchoResponse { + message: format!("echo[{}]: {}", thread_name(), req.message), + }) + } +} + +fn thread_name() -> String { + std::thread::current() + .name() + .unwrap_or("unnamed") + .to_string() +} + +struct AllowAll; + +impl AuthController for AllowAll { + fn auth( + self: Arc, + _toolbox: &ArcToolbox, + _header: String, + conn: Arc, + ) -> LocalBoxFuture<'static, Result<()>> { + async move { + conn.set_roles(Arc::new(vec![1])); + Ok(()) + } + .boxed_local() + } +} + +/// A port nothing holds, found by binding zero and letting it go. +fn free_port() -> u16 { + std::net::TcpListener::bind("127.0.0.1:0") + .unwrap() + .local_addr() + .unwrap() + .port() +} + +#[test] +fn three_shards_serve_one_port_and_stop_together() { + // SAFETY: this test binary has one test, so nothing else reads the + // environment concurrently. + unsafe { std::env::set_var("WS_SHARDS", "3") }; + let port = free_port(); + + let (stop_tx, stop_rx) = futures::channel::oneshot::channel::<()>(); + let server = std::thread::spawn(move || { + let config = WsServerConfig { + address: format!("127.0.0.1:{port}"), + ..WsServerConfig::default() + }; + let mut server = WebsocketServer::new(config); + server.set_auth_controller(AllowAll); + server.add_handler(MethodEcho); + server.listen_until(async move { + let _ = stop_rx.await; + }) + }); + + let reactor = Reactor::local().unwrap(); + let handle = reactor.handle(); + let url = format!("ws://127.0.0.1:{port}/"); + let replies = block_on_with(&reactor, async { + // The shards bind as their threads start; give them a moment. + let mut replies = Vec::new(); + for attempt in 0..12 { + let connected = WsClientBuilder::new() + .protocol_header("0echo") + .build(&url, &handle) + .await; + let (mut client, _) = match connected { + Ok(client) => client, + Err(_) if replies.is_empty() && attempt < 11 => { + nagoya::sleep(Duration::from_millis(100)).await; + continue; + } + Err(err) => panic!("connect failed: {err}"), + }; + let reply = client + .request(EchoRequest { + message: format!("hello {attempt}"), + }) + .await + .expect("request"); + replies.push(reply.message); + } + replies + }); + assert!(!replies.is_empty()); + for reply in &replies { + assert!(reply.contains(": hello "), "{reply}"); + eprintln!("{reply}"); + } + + stop_tx.send(()).unwrap(); + let served = server.join().expect("server thread panicked"); + assert!(served.is_ok(), "{served:?}"); +}