From e3d26666c1efc47173a1da0a3a05fe5a29a1df39 Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 24 Sep 2026 03:27:26 +0700 Subject: [PATCH 1/2] Serve on every shard: one reactor and SO_REUSEPORT listener each The server ran on one thread whatever WS_SHARDS said, and warned about it at every boot, because a nagoya listener could not be shared between reactors. Consumers set WS_SHARDS=1 to hide the warning. nagoya 0.1.13 adds bind_shared (SO_REUSEPORT). listen now starts a thread per extra shard, each with its own local reactor and its own listener on the one port, so a connection is accepted, registered and served by the same reactor and never crosses a thread. The first bind picks the address and the others take it exactly, which matters when the port was zero. The server and the connection table are prepared once and shared, so the toolbox reaches a connection on any shard. The first shard owns the stop, whether a signal or the caller's, and cancels the rest; listen returns once all have joined. With one shard nothing changes. tests/sharded_listen.rs drives three shards over TCP: every client completes a request, and one stop ends them all. Release as 3.3.1, needing nagoya ^0.1.13. --- Cargo.toml | 4 +- src/libs/ws/listener.rs | 22 +++++- src/libs/ws/server.rs | 151 ++++++++++++++++++++++++++++---------- tests/sharded_listen.rs | 157 ++++++++++++++++++++++++++++++++++++++++ 4 files changed, 292 insertions(+), 42 deletions(-) create mode 100644 tests/sharded_listen.rs 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..c7522a9 --- /dev/null +++ b/tests/sharded_listen.rs @@ -0,0 +1,157 @@ +//! `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:?}"); +} From 28b7f627fb5feac38eacb073d0d2722eb1ceacfe Mon Sep 17 00:00:00 2001 From: meh Date: Thu, 24 Sep 2026 05:44:52 +0700 Subject: [PATCH 2/2] Format the sharded listen test --- tests/sharded_listen.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/sharded_listen.rs b/tests/sharded_listen.rs index c7522a9..96e41b8 100644 --- a/tests/sharded_listen.rs +++ b/tests/sharded_listen.rs @@ -66,7 +66,10 @@ impl RequestHandler for MethodEcho { } fn thread_name() -> String { - std::thread::current().name().unwrap_or("unnamed").to_string() + std::thread::current() + .name() + .unwrap_or("unnamed") + .to_string() } struct AllowAll;