From 6dd0a69dc0eb5ca33d5beed2c0470ff43ec61543 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Tue, 4 Aug 2026 17:09:57 +0800 Subject: [PATCH 1/4] bench: env-gated runtime and client-locality experiment knobs Add three bench-only knobs, all default-off, used by the Linux performance verification of the onthebench throughput gap: - BENCH_RT_WORKERS=N pins the tokio multi-thread worker count. - BENCH_RT_TPC=N short-circuits the plain-HTTP proxy listener into N thread-per-core workers (one OS thread + current_thread runtime + own listener each). BENCH_RT_TPC_MODE=reuseport|stride picks SO_REUSEPORT on one shared port vs a port-per-worker stride. - BENCH_THREAD_LOCAL_CLIENT=1 hands every OS thread its own upstream reqwest client, so a thread-per-core worker's upstream connections are polled by its own runtime and a response arrival never needs a cross-thread wakeup. Production behavior is unchanged when the variables are unset. Measured on 4 pinned vCPUs proxying to a local mock upstream (c=128 saturation, 25s windows, fail=0): mt work-stealing 24.0k rps 152us CPU/req p99 9.1ms TPC + SO_REUSEPORT 39.0k rps 100us CPU/req p99 5.9ms TPC + thread-local client 44.6k rps 89us CPU/req p99 4.6ms With the 10ms-TTFT mock (leaderboard methodology) the same A/B is 21.4k vs 39.4k rps (+84%). The whole win is kernel-side: the futex park/unpark storm (232k context switches/s down to 0.2k) and the per-response cross-thread eventfd wakeup (1.77/req down to 0). Thread-local clients under work-stealing show zero gain, confirming the mechanism is connection/task thread locality, not pool sizing. --- Cargo.lock | 1 + crates/aisix-gateway/src/upstream_tls.rs | 32 ++++++ crates/aisix-server/Cargo.toml | 3 + crates/aisix-server/src/main.rs | 124 ++++++++++++++++++++++- 4 files changed, 158 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b89e4d40..dc7a0314 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -446,6 +446,7 @@ dependencies = [ "rustls 0.23.38", "serde", "serde_json", + "socket2 0.5.10", "tempfile", "tokio", "tracing", diff --git a/crates/aisix-gateway/src/upstream_tls.rs b/crates/aisix-gateway/src/upstream_tls.rs index 494bd977..ca739950 100644 --- a/crates/aisix-gateway/src/upstream_tls.rs +++ b/crates/aisix-gateway/src/upstream_tls.rs @@ -210,6 +210,30 @@ pub fn reqwest_material() -> &'static ReqwestTlsMaterial { /// idle connection pool. static PK_CLIENTS: OnceLock> = OnceLock::new(); +/// Bench-only (BENCH_THREAD_LOCAL_CLIENT): see `client_for_provider_key`. +fn bench_thread_local_client() -> bool { + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| std::env::var("BENCH_THREAD_LOCAL_CLIENT").is_ok_and(|v| v == "1")) +} + +/// One client per OS thread, built with the deployment's connection and +/// trust settings. Bench-only. +fn thread_local_bench_client() -> reqwest::Client { + thread_local! { + static CLIENT: std::cell::OnceCell = + const { std::cell::OnceCell::new() }; + } + CLIENT.with(|c| { + c.get_or_init(|| { + crate::upstream_http::client_builder() + .user_agent("aisix-bench-thread-local") + .build() + .expect("bench thread-local client build failed") + }) + .clone() + }) +} + /// The client to dispatch this Provider Key's request on. /// /// Returns `shared` unchanged whenever the key sets no override, which @@ -226,6 +250,14 @@ pub fn client_for_provider_key( tls: Option<&ProviderKeyTls>, ) -> reqwest::Client { let Some(tls) = tls.filter(|t| !t.is_noop()) else { + // Bench-only: BENCH_THREAD_LOCAL_CLIENT=1 hands every OS thread its + // own client/pool so a thread-per-core worker's upstream + // connections are driven by its own runtime — no cross-thread + // wakeups on response arrival. Keys with a TLS override keep the + // normal path. + if bench_thread_local_client() { + return thread_local_bench_client(); + } return shared.clone(); }; let cache = PK_CLIENTS.get_or_init(dashmap::DashMap::new); diff --git a/crates/aisix-server/Cargo.toml b/crates/aisix-server/Cargo.toml index 3b3ad9ea..06eb157f 100644 --- a/crates/aisix-server/Cargo.toml +++ b/crates/aisix-server/Cargo.toml @@ -46,6 +46,9 @@ rcgen.workspace = true uuid.workspace = true hostname.workspace = true x509-parser = "0.16" +# Bench-only (BENCH_RT_TPC): SO_REUSEPORT listener setup for the +# thread-per-core A/B path in main.rs. Not used in production serving. +socket2 = "0.5" # Only for the hyper connection builder `axum_server::Server::http_builder()` # hands back — where `downstream.idle_timeout_secs` is applied. Same 0.1.x # axum-server itself depends on, so the builder types unify. diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 322230c0..774d9cd7 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -101,8 +101,19 @@ enum CliCommand { }, } -#[tokio::main] -async fn main() -> anyhow::Result<()> { +fn main() -> anyhow::Result<()> { + // Bench-only: BENCH_RT_WORKERS=N pins the tokio worker-thread count so + // saturation A/B runs control parallelism explicitly. Named without the + // AISIX_ prefix because that namespace is claimed by the config env + // override layer. + let mut rt = tokio::runtime::Builder::new_multi_thread(); + if let Ok(n) = std::env::var("BENCH_RT_WORKERS") { + rt.worker_threads(n.parse().expect("BENCH_RT_WORKERS must be a number")); + } + rt.enable_all().build()?.block_on(async_main()) +} + +async fn async_main() -> anyhow::Result<()> { // Install the process-level rustls CryptoProvider before anything // else touches TLS. rustls 0.23 dropped implicit provider selection // and panics at first use when both `aws-lc-rs` and `ring` features @@ -1392,6 +1403,20 @@ async fn serve_http( cancel: watch::Receiver, label: &'static str, ) -> anyhow::Result<()> { + // Bench-only: BENCH_RT_TPC=N short-circuits the plain-HTTP proxy + // listener into N thread-per-core workers (one OS thread + one + // current_thread runtime + one listener each). BENCH_RT_TPC_MODE + // picks the listener layout: "reuseport" (default) binds every worker + // to the same port with SO_REUSEPORT so the kernel spreads accepted + // connections by 4-tuple hash; "stride" binds port, port+1, ... like + // the macOS phase did. + if label == "proxy" && tls.is_none() { + if let Ok(n) = std::env::var("BENCH_RT_TPC") { + let workers: usize = n.parse().expect("BENCH_RT_TPC must be a number"); + return serve_http_tpc(addr, router, idle_timeout, cancel, label, workers).await; + } + } + // Resolved before binding so a bad cert path still fails with the // same error it always did, before a port is taken. let tls_config = match tls { @@ -1448,6 +1473,101 @@ async fn serve_http( Ok(()) } +/// Bench-only (BENCH_RT_TPC): serve the proxy from N OS threads, each with +/// its own current_thread runtime and its own listener. The router (and so +/// the shared AppState / reqwest client) is cloned per thread, matching the +/// macOS-phase experiment semantics. +async fn serve_http_tpc( + addr: std::net::SocketAddr, + router: axum::Router, + idle_timeout: Option, + cancel: watch::Receiver, + label: &'static str, + workers: usize, +) -> anyhow::Result<()> { + let mode = std::env::var("BENCH_RT_TPC_MODE").unwrap_or_else(|_| "reuseport".to_string()); + let reuseport = match mode.as_str() { + "reuseport" => true, + "stride" => false, + other => anyhow::bail!("BENCH_RT_TPC_MODE must be reuseport|stride, got {other}"), + }; + let mut threads = Vec::with_capacity(workers); + for i in 0..workers { + let bind_addr = if reuseport { + addr + } else { + std::net::SocketAddr::new(addr.ip(), addr.port() + u16::try_from(i)?) + }; + // Bound on the parent thread so a bind failure aborts startup + // before any worker begins serving. + let listener = bind_tpc_listener(bind_addr, reuseport) + .map_err(|e| anyhow::anyhow!("{label} tpc listener bind {bind_addr} failed: {e}"))?; + listener.set_nonblocking(true)?; + let router = router.clone(); + let cancel = cancel.clone(); + let thread = std::thread::Builder::new().name(format!("tpc-{i}")).spawn( + move || -> anyhow::Result<()> { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + rt.block_on(async move { + let handle = axum_server::Handle::new(); + tokio::spawn({ + let handle = handle.clone(); + async move { + shutdown_signal(cancel, label).await; + handle.graceful_shutdown(None); + } + }); + tracing::info!( + %bind_addr, + label, + worker = i, + reuseport, + "aisix listening (http, tpc)" + ); + let make_service = + router.into_make_service_with_connect_info::(); + let mut server = axum_server::from_tcp(listener).handle(handle); + apply_idle_timeout(server.http_builder(), idle_timeout); + server.serve(make_service).await?; + Ok(()) + }) + }, + )?; + threads.push(thread); + } + tokio::task::spawn_blocking(move || { + for t in threads { + t.join() + .map_err(|_| anyhow::anyhow!("tpc worker thread panicked"))??; + } + Ok(()) + }) + .await? +} + +/// Listener for the bench-only TPC path: plain bind for "stride" mode, +/// SO_REUSEPORT (set before bind) for same-port kernel load balancing. +fn bind_tpc_listener( + addr: std::net::SocketAddr, + reuseport: bool, +) -> std::io::Result { + let domain = if addr.is_ipv4() { + socket2::Domain::IPV4 + } else { + socket2::Domain::IPV6 + }; + let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))?; + socket.set_reuse_address(true)?; + if reuseport { + socket.set_reuse_port(true)?; + } + socket.bind(&addr.into())?; + socket.listen(1024)?; + Ok(socket.into()) +} + /// Close an accepted HTTP/1.1 connection that sits idle for `idle_timeout`. /// /// hyper arms this timer only when it is waiting for a request head, and From 93d61d9f97640d077f7e614e812b93578cf370f4 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Tue, 4 Aug 2026 23:01:34 +0800 Subject: [PATCH 2/4] feat(proxy): serve from thread-per-core workers with per-worker upstream pools The proxy hands a request between threads about twice per request: once when a work-stealing worker picks up the task, again when the upstream response lands on whichever thread happens to own that connection. Each handoff costs a wakeup and a context switch, and on a 4-core saturation run that is where most of the CPU goes -- 89% of it is kernel time, at 232k context switches per second. Serve instead from N independent workers, each with its own runtime, its own SO_REUSEPORT listener on proxy.addr, and its own upstream connection pool, so a request is accepted, dispatched, answered, and its upstream call polled all on one thread. Two bootstrap knobs, applied at startup: - proxy.thread_per_core: on for Linux when omitted, off elsewhere, since the kernel spreading this relies on is a Linux behavior. Set false to serve from one shared runtime on any platform. - proxy.workers: defaults to the parallelism available to the process, which follows a cgroup CPU limit or a taskset affinity mask. Rejected at load when zero. Per-worker pools hang off the one chokepoint every handler family already dispatches through (client_for_provider_key), keyed by a per-thread marker rather than a process-wide flag -- the playground runs proxy handlers on the shared runtime, and those correctly keep using the process-wide pool. Provider keys carrying a TLS override keep their dedicated client. A scan test holds every dispatch client to one user agent, since one per-worker pool now stands in for all of them. The listeners co-bind, which would turn a second gateway on the same address from a startup failure into a silent traffic split, so the mode probes the address with a non-reuseport bind first and keeps failing loudly. A worker leaving for any reason -- accept loop error, panic unwinding -- brings the process down rather than leaving it serving on fewer listeners than it reported binding. Measured on 4 pinned vCPUs against a local mock upstream, 25s windows, fail=0, config knobs only: c=128 thread-per-core 45,447 rps 398% CPU 88us/req p99 4.38ms work-stealing 24,297 rps 365% CPU 150us/req p99 8.96ms c=768, 10ms TTFT mock (leaderboard methodology) thread-per-core 39,873 rps 398% CPU 100us/req p99 24.9ms work-stealing 21,897 rps 358% CPU 164us/req p99 55.4ms Below about four client connections per worker the kernel's per-connection spreading leaves workers uneven: c=8 measures -26.8%. That is documented on the knob. The e2e suite passes in full in both modes (175 files, 463 tests); CI runs it as a two-leg matrix. --- .github/workflows/ci.yml | 30 +- config.example.yaml | 10 + config.managed.yaml | 5 + crates/aisix-admin/src/playground_handler.rs | 2 + crates/aisix-core/src/config.rs | 132 +++++++ crates/aisix-gateway/src/upstream_http.rs | 42 +++ crates/aisix-gateway/src/upstream_tls.rs | 87 +++-- crates/aisix-proxy/src/a2a.rs | 2 + crates/aisix-proxy/src/audio.rs | 2 + crates/aisix-proxy/src/completions.rs | 2 + crates/aisix-proxy/src/count_tokens.rs | 2 + crates/aisix-proxy/src/embeddings.rs | 2 + crates/aisix-proxy/src/images.rs | 2 + crates/aisix-proxy/src/jobs.rs | 2 + crates/aisix-proxy/src/lib.rs | 4 + crates/aisix-proxy/src/mcp.rs | 2 + crates/aisix-proxy/src/messages.rs | 2 + crates/aisix-proxy/src/models.rs | 2 + crates/aisix-proxy/src/passthrough.rs | 2 + crates/aisix-proxy/src/realtime.rs | 2 + crates/aisix-proxy/src/rerank.rs | 2 + crates/aisix-proxy/src/responses.rs | 2 + crates/aisix-proxy/src/state.rs | 2 + crates/aisix-proxy/src/videos.rs | 2 + crates/aisix-server/Cargo.toml | 8 +- crates/aisix-server/src/main.rs | 334 ++++++++++++------ tests/e2e/src/cases/listener-tls-e2e.test.ts | 9 +- .../src/cases/upstream-pool-idle-e2e.test.ts | 15 +- tests/e2e/src/harness/app.ts | 29 ++ tests/e2e/src/harness/index.ts | 2 +- 30 files changed, 600 insertions(+), 141 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1023f944..64ada348 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -218,9 +218,25 @@ jobs: retention-days: 1 e2e: - name: e2e (vitest) + coverage + name: e2e (vitest, ${{ matrix.serving }}) + coverage needs: [build-bin] runs-on: ubuntu-latest + strategy: + # The two serving modes are separate products of the same binary: + # thread-per-core answers a request on the thread that accepted it + # and dispatches on that thread's own upstream pool, work-stealing + # shares one runtime and one pool. A failure in one says nothing + # about the other, so neither leg cancels the other. + fail-fast: false + matrix: + include: + # The Linux default, and so the one production runs. + - serving: thread-per-core + thread_per_core: "true" + # The documented fallback. Kept green so it stays a real + # option when a deployment needs it. + - serving: work-stealing + thread_per_core: "false" services: etcd: image: quay.io/coreos/etcd:v3.5.15 @@ -270,11 +286,19 @@ jobs: - name: run e2e if: steps.harness.outputs.present == 'true' working-directory: tests/e2e + # Read by the harness, which writes it into every spawned + # gateway's `proxy.thread_per_core`. The binary's own default is + # platform-derived, so without this both legs would run the same + # mode and the second would prove nothing. + env: + E2E_THREAD_PER_CORE: ${{ matrix.thread_per_core }} run: pnpm test - uses: actions/upload-artifact@v7 if: always() with: - name: coverage-e2e + # Per leg: two uploads under one name is a hard conflict the + # moment the suite starts emitting coverage. + name: coverage-e2e-${{ matrix.serving }} path: tests/e2e/coverage/lcov.info if-no-files-found: ignore retention-days: 7 @@ -296,7 +320,7 @@ jobs: # follow-up), this download must remain soft — otherwise the # coverage-gate hard-fails every run on `Artifact not found`. continue-on-error: true - with: { name: coverage-e2e, path: cov/e2e } + with: { pattern: coverage-e2e-*, path: cov/e2e, merge-multiple: true } - name: merge + threshold run: | npm i -g lcov-result-merger @lcov-viewer/cli diff --git a/config.example.yaml b/config.example.yaml index c66ca767..3ce35145 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -51,6 +51,16 @@ proxy: # trusted_proxies: ["10.0.0.0/8", "127.0.0.1/32"] # recursive: true # header: x-forwarded-for + # Serve from independent worker threads, each with its own runtime, + # its own SO_REUSEPORT listener on addr, and its own upstream + # connection pool, so a request is handled end to end on the thread + # that accepted it. Omitted, it is on for Linux and off elsewhere. + # Set false to serve from one shared runtime. Applied at startup. + # thread_per_core: true + # Number of proxy worker threads. Omitted, it follows the parallelism + # available to the process, so a container CPU limit or a taskset + # affinity mask sizes it. Applied at startup. + # workers: 4 admin: addr: "127.0.0.1:3001" diff --git a/config.managed.yaml b/config.managed.yaml index c3cf4deb..253c78a2 100644 --- a/config.managed.yaml +++ b/config.managed.yaml @@ -44,6 +44,11 @@ proxy: # 0 = no request-body cap (the default); set a value to bound # per-request memory. # request_body_limit_bytes: 0 + # Serving topology; on for Linux by default. A managed deployment + # reaches these through the environment rather than by editing this + # file: AISIX_PROXY__THREAD_PER_CORE, AISIX_PROXY__WORKERS. + # thread_per_core: true + # workers: 4 admin: # Bind to an unbindable port so even if managed mode somehow diff --git a/crates/aisix-admin/src/playground_handler.rs b/crates/aisix-admin/src/playground_handler.rs index cc0e7043..4f88fcab 100644 --- a/crates/aisix-admin/src/playground_handler.rs +++ b/crates/aisix-admin/src/playground_handler.rs @@ -80,6 +80,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-core/src/config.rs b/crates/aisix-core/src/config.rs index a4a97315..92fbf818 100644 --- a/crates/aisix-core/src/config.rs +++ b/crates/aisix-core/src/config.rs @@ -393,12 +393,58 @@ pub struct ProxyConfig { /// an L7 LB / ingress that sets `x-forwarded-for`. #[serde(default)] pub real_ip: RealIpConfig, + /// Serve the proxy from independent worker threads — each with its + /// own runtime, its own `SO_REUSEPORT` listener on `addr`, and its + /// own upstream connection pool — instead of one shared runtime + /// whose threads hand work to each other. + /// + /// Omitted, the default, enables it on Linux and disables it + /// elsewhere: the kernel spreads incoming connections across + /// same-port listeners on Linux, and other platforms do not. + /// Set `false` to serve from one shared runtime on any platform. + /// + /// A request is handled end to end on the thread that accepted it, + /// which removes a cross-thread handoff per request. On a small + /// number of client connections (fewer than about four per worker) + /// the kernel's per-connection spreading can leave workers unevenly + /// loaded; throughput at that size may be lower than with a shared + /// runtime. + /// + /// Applied at startup. Changing it requires a restart. + #[serde(default)] + pub thread_per_core: Option, + /// Number of proxy worker threads. + /// + /// Omitted, the default, uses the parallelism available to the + /// process, which follows the CPU limits applied by a container + /// runtime, cgroup, or `taskset`. Must be at least 1. + /// + /// Applied at startup. Changing it requires a restart. + #[serde(default)] + pub workers: Option, } impl ProxyConfig { const fn default_body_limit() -> usize { 0 } + + /// Whether the proxy serves from thread-per-core workers, resolving + /// the platform default when unset. + pub fn thread_per_core_enabled(&self) -> bool { + self.thread_per_core.unwrap_or(cfg!(target_os = "linux")) + } + + /// Proxy worker-thread count, resolving the default when unset. + /// + /// `available_parallelism` reports the CPUs this process may actually + /// run on, so a cgroup CPU limit or a `taskset` affinity mask sizes + /// the pool correctly without the operator restating it here. Falls + /// back to 1 on the platforms that cannot report it. + pub fn worker_threads(&self) -> usize { + self.workers + .unwrap_or_else(|| std::thread::available_parallelism().map_or(1, |n| n.get())) + } } /// nginx `set_real_ip_from` + `real_ip_recursive` equivalent. Resolves @@ -1155,6 +1201,15 @@ impl Config { "proxy.real_ip.trusted_proxies invalid CIDR/IP: {bad}" ))); } + // Zero workers would bind no listener at all: the proxy would + // boot, report healthy, and refuse every connection. + if self.proxy.workers == Some(0) { + return Err(BootstrapError::Config( + "proxy.workers must be at least 1 (omit it to use the \ + parallelism available to the process)" + .into(), + )); + } // The dedicated metrics listener address must be a bindable // socket address — it is always bound when prometheus is enabled. let metrics_addr = &self.observability.metrics.prometheus.addr; @@ -2191,4 +2246,81 @@ admin: assert_eq!(tls.client_key_file, "/c.key"); assert_eq!(tls.domain_name.as_deref(), Some("etcd.aisix.cloud")); } + + /// Serving topology is a startup decision, so every existing config + /// — none of which names it — has to keep loading and resolve to the + /// platform's answer. + #[test] + fn serving_topology_defaults_to_the_platform_answer() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let cfg = Config::load_from_path(Some(f.path())).unwrap(); + assert_eq!(cfg.proxy.thread_per_core, None); + assert_eq!(cfg.proxy.workers, None); + assert_eq!( + cfg.proxy.thread_per_core_enabled(), + cfg!(target_os = "linux"), + "thread-per-core is the default where the kernel spreads \ + connections across same-port listeners, and only there" + ); + assert_eq!( + cfg.proxy.worker_threads(), + std::thread::available_parallelism().map_or(1, |n| n.get()), + ); + } + + /// The fallback an operator reaches for when thread-per-core is the + /// wrong shape for their traffic. It has to win on every platform, + /// including the one where it is also the default. + #[test] + fn explicit_serving_topology_overrides_the_platform_default() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" + thread_per_core: false + workers: 3 +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let cfg = Config::load_from_path(Some(f.path())).unwrap(); + assert_eq!(cfg.proxy.thread_per_core, Some(false)); + assert!(!cfg.proxy.thread_per_core_enabled()); + assert_eq!(cfg.proxy.worker_threads(), 3); + } + + /// Zero workers would bind no listener and still report a healthy + /// boot, so it has to fail at load naming the field to fix. + #[test] + fn rejects_zero_proxy_workers() { + let f = write_yaml( + r#" +etcd: + endpoints: ["http://localhost:2379"] +proxy: + addr: "0.0.0.0:3000" + workers: 0 +admin: + addr: "127.0.0.1:3001" + admin_keys: ["k1"] +"#, + ); + let err = Config::load_from_path(Some(f.path())) + .unwrap_err() + .to_string(); + assert!(err.contains("proxy.workers"), "unexpected error: {err}"); + } } diff --git a/crates/aisix-gateway/src/upstream_http.rs b/crates/aisix-gateway/src/upstream_http.rs index df6e45b9..3fcc1286 100644 --- a/crates/aisix-gateway/src/upstream_http.rs +++ b/crates/aisix-gateway/src/upstream_http.rs @@ -418,6 +418,48 @@ mod tests { assert!(!production.contains("fn t()"), "{production}"); } + /// On a thread-per-core worker every dispatch runs on that worker's + /// own pool, and that one pool stands in for all of the clients + /// below — so it is built with a single user agent + /// (`upstream_tls::DISPATCH_USER_AGENT`). + /// + /// That substitution is only invisible while the clients it replaces + /// agree on the user agent. Give one bridge its own and the header + /// it sends changes depending on which serving mode the deployment + /// runs, which is not something an upstream-facing identity is + /// allowed to do. Whoever wants a distinct user agent has to give + /// that client a distinct pool as well. + #[test] + fn every_dispatch_client_presents_the_same_user_agent() { + let crates_dir = std::path::Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/..")); + // Only literals: the telemetry, heartbeat, and OTLP clients build + // theirs from a version string or a named constant, and none of + // them talks to a model provider or reaches the per-worker pool. + const NEEDLE: &str = ".user_agent(\""; + let mut offenders = Vec::new(); + for file in rust_sources(crates_dir) { + let src = std::fs::read_to_string(&file).expect("read source"); + for (n, line) in production_half(&src).lines().enumerate() { + let Some(rest) = line.split_once(NEEDLE).map(|(_, r)| r) else { + continue; + }; + let Some((agent, _)) = rest.split_once('"') else { + continue; + }; + if agent != crate::upstream_tls::DISPATCH_USER_AGENT { + offenders.push(format!("{}:{}: {agent}", file.display(), n + 1)); + } + } + } + assert!( + offenders.is_empty(), + "these dispatch clients present a user agent the per-worker \ + pool would replace with `{}`:\n{}", + crate::upstream_tls::DISPATCH_USER_AGENT, + offenders.join("\n"), + ); + } + fn rust_sources(dir: &std::path::Path) -> Vec { let mut out = Vec::new(); let Ok(entries) = std::fs::read_dir(dir) else { diff --git a/crates/aisix-gateway/src/upstream_tls.rs b/crates/aisix-gateway/src/upstream_tls.rs index ca739950..ed477a98 100644 --- a/crates/aisix-gateway/src/upstream_tls.rs +++ b/crates/aisix-gateway/src/upstream_tls.rs @@ -210,25 +210,67 @@ pub fn reqwest_material() -> &'static ReqwestTlsMaterial { /// idle connection pool. static PK_CLIENTS: OnceLock> = OnceLock::new(); -/// Bench-only (BENCH_THREAD_LOCAL_CLIENT): see `client_for_provider_key`. -fn bench_thread_local_client() -> bool { - static ON: OnceLock = OnceLock::new(); - *ON.get_or_init(|| std::env::var("BENCH_THREAD_LOCAL_CLIENT").is_ok_and(|v| v == "1")) +// ─── per-worker pools ──────────────────────────────────────────────── + +/// The user agent every dispatch-path client is built with. +/// +/// A worker's pool stands in for those clients, so it has to present the +/// same identity upstream. `every_dispatch_client_presents_the_same_user_agent` +/// holds them in step. +pub(crate) const DISPATCH_USER_AGENT: &str = "aisix/0.1"; + +thread_local! { + /// Whether this thread serves proxy traffic on its own runtime. + static IS_WORKER_THREAD: std::cell::Cell = const { std::cell::Cell::new(false) }; + + /// This worker's upstream pool, built on first dispatch. `None` once + /// a build failure has been reported for this thread. + static WORKER_CLIENT: std::cell::OnceCell> = + const { std::cell::OnceCell::new() }; +} + +/// Declares the calling thread a proxy worker with its own runtime, so +/// its dispatches use a pool this thread alone polls. +/// +/// Called once per worker in thread-per-core serving. Left unset +/// everywhere else — the shared runtime's threads, the blocking pool, +/// background tasks — so those keep dispatching on the process-wide +/// pools exactly as before. +pub fn mark_worker_thread() { + IS_WORKER_THREAD.set(true); } -/// One client per OS thread, built with the deployment's connection and -/// trust settings. Bench-only. -fn thread_local_bench_client() -> reqwest::Client { - thread_local! { - static CLIENT: std::cell::OnceCell = - const { std::cell::OnceCell::new() }; +/// This worker's pool, or `None` when the thread is not a worker or the +/// pool could not be built. +/// +/// One pool per worker rather than per client: the dispatch clients are +/// all built from the same recipe, so merging them costs nothing but the +/// per-owner split of `pool_max_idle_per_host`, and keeping them split +/// would multiply idle connections by the worker count for no gain. +/// +/// A build failure falls back to the shared pool — which carries the +/// deployment's trust settings — and is cached so one broken +/// configuration cannot log per request. +fn worker_client() -> Option { + if !IS_WORKER_THREAD.get() { + return None; } - CLIENT.with(|c| { - c.get_or_init(|| { - crate::upstream_http::client_builder() - .user_agent("aisix-bench-thread-local") + WORKER_CLIENT.with(|cell| { + cell.get_or_init(|| { + match crate::upstream_http::client_builder() + .user_agent(DISPATCH_USER_AGENT) .build() - .expect("bench thread-local client build failed") + { + Ok(client) => Some(client), + Err(e) => { + tracing::error!( + error = %e, + "per-worker upstream pool could not be built; this worker \ + dispatches on the shared pool" + ); + None + } + } }) .clone() }) @@ -250,15 +292,12 @@ pub fn client_for_provider_key( tls: Option<&ProviderKeyTls>, ) -> reqwest::Client { let Some(tls) = tls.filter(|t| !t.is_noop()) else { - // Bench-only: BENCH_THREAD_LOCAL_CLIENT=1 hands every OS thread its - // own client/pool so a thread-per-core worker's upstream - // connections are driven by its own runtime — no cross-thread - // wakeups on response arrival. Keys with a TLS override keep the - // normal path. - if bench_thread_local_client() { - return thread_local_bench_client(); - } - return shared.clone(); + // On a thread-per-core worker, dispatch on that worker's own + // pool: the upstream connection is then read by the same runtime + // that is waiting for the response, instead of waking a thread + // that has to hand it back. Everywhere else this is `None` and + // the shared pool is used, as it always was. + return worker_client().unwrap_or_else(|| shared.clone()); }; let cache = PK_CLIENTS.get_or_init(dashmap::DashMap::new); if let Some(existing) = cache.get(tls) { diff --git a/crates/aisix-proxy/src/a2a.rs b/crates/aisix-proxy/src/a2a.rs index e722b36a..365e6ab0 100644 --- a/crates/aisix-proxy/src/a2a.rs +++ b/crates/aisix-proxy/src/a2a.rs @@ -376,6 +376,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/audio.rs b/crates/aisix-proxy/src/audio.rs index bcc6fb34..ea1a7fab 100644 --- a/crates/aisix-proxy/src/audio.rs +++ b/crates/aisix-proxy/src/audio.rs @@ -1440,6 +1440,8 @@ mod tests { request_body_limit_bytes: 10_485_760, // 10 MB for audio real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/completions.rs b/crates/aisix-proxy/src/completions.rs index 5ba0fa18..afe8921f 100644 --- a/crates/aisix-proxy/src/completions.rs +++ b/crates/aisix-proxy/src/completions.rs @@ -745,6 +745,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/count_tokens.rs b/crates/aisix-proxy/src/count_tokens.rs index ea0f8bf8..e3ec09ee 100644 --- a/crates/aisix-proxy/src/count_tokens.rs +++ b/crates/aisix-proxy/src/count_tokens.rs @@ -545,6 +545,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/embeddings.rs b/crates/aisix-proxy/src/embeddings.rs index 738e4c82..59737635 100644 --- a/crates/aisix-proxy/src/embeddings.rs +++ b/crates/aisix-proxy/src/embeddings.rs @@ -670,6 +670,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/images.rs b/crates/aisix-proxy/src/images.rs index 5019831e..2f988b99 100644 --- a/crates/aisix-proxy/src/images.rs +++ b/crates/aisix-proxy/src/images.rs @@ -542,6 +542,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/jobs.rs b/crates/aisix-proxy/src/jobs.rs index 7519b399..605729b8 100644 --- a/crates/aisix-proxy/src/jobs.rs +++ b/crates/aisix-proxy/src/jobs.rs @@ -1803,6 +1803,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/lib.rs b/crates/aisix-proxy/src/lib.rs index d9b32898..2c17a9af 100644 --- a/crates/aisix-proxy/src/lib.rs +++ b/crates/aisix-proxy/src/lib.rs @@ -640,6 +640,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } @@ -1613,6 +1615,8 @@ mod tests { request_body_limit_bytes: limit, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, }; ProxyState::new(handle, hub, &cfg).without_cache() } diff --git a/crates/aisix-proxy/src/mcp.rs b/crates/aisix-proxy/src/mcp.rs index ee2452d1..c7de7f1c 100644 --- a/crates/aisix-proxy/src/mcp.rs +++ b/crates/aisix-proxy/src/mcp.rs @@ -464,6 +464,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/messages.rs b/crates/aisix-proxy/src/messages.rs index b6caeee5..1a7fd2a9 100644 --- a/crates/aisix-proxy/src/messages.rs +++ b/crates/aisix-proxy/src/messages.rs @@ -3579,6 +3579,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/models.rs b/crates/aisix-proxy/src/models.rs index 4583add0..28405063 100644 --- a/crates/aisix-proxy/src/models.rs +++ b/crates/aisix-proxy/src/models.rs @@ -120,6 +120,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/passthrough.rs b/crates/aisix-proxy/src/passthrough.rs index 805449fe..25ae35bb 100644 --- a/crates/aisix-proxy/src/passthrough.rs +++ b/crates/aisix-proxy/src/passthrough.rs @@ -895,6 +895,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/realtime.rs b/crates/aisix-proxy/src/realtime.rs index 3a975feb..04028eff 100644 --- a/crates/aisix-proxy/src/realtime.rs +++ b/crates/aisix-proxy/src/realtime.rs @@ -800,6 +800,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/rerank.rs b/crates/aisix-proxy/src/rerank.rs index 4be1aef9..1f0f7224 100644 --- a/crates/aisix-proxy/src/rerank.rs +++ b/crates/aisix-proxy/src/rerank.rs @@ -736,6 +736,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/responses.rs b/crates/aisix-proxy/src/responses.rs index be9543e6..2dacc5a5 100644 --- a/crates/aisix-proxy/src/responses.rs +++ b/crates/aisix-proxy/src/responses.rs @@ -3066,6 +3066,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-proxy/src/state.rs b/crates/aisix-proxy/src/state.rs index e278573a..3cd8c78b 100644 --- a/crates/aisix-proxy/src/state.rs +++ b/crates/aisix-proxy/src/state.rs @@ -388,6 +388,8 @@ mod tests { request_body_limit_bytes: 1_048_576, tls: None, real_ip: Default::default(), + thread_per_core: None, + workers: None, }, ) } diff --git a/crates/aisix-proxy/src/videos.rs b/crates/aisix-proxy/src/videos.rs index ecb0968b..0a871064 100644 --- a/crates/aisix-proxy/src/videos.rs +++ b/crates/aisix-proxy/src/videos.rs @@ -2328,6 +2328,8 @@ mod tests { request_body_limit_bytes: 1_048_576, real_ip: Default::default(), tls: None, + thread_per_core: None, + workers: None, } } diff --git a/crates/aisix-server/Cargo.toml b/crates/aisix-server/Cargo.toml index 06eb157f..99472d9b 100644 --- a/crates/aisix-server/Cargo.toml +++ b/crates/aisix-server/Cargo.toml @@ -46,9 +46,11 @@ rcgen.workspace = true uuid.workspace = true hostname.workspace = true x509-parser = "0.16" -# Bench-only (BENCH_RT_TPC): SO_REUSEPORT listener setup for the -# thread-per-core A/B path in main.rs. Not used in production serving. -socket2 = "0.5" +# `SO_REUSEPORT` for the thread-per-core proxy listeners. `all` is +# explicit because `set_reuse_port` lives behind it — it resolves +# without the feature today only because other dependencies happen to +# enable it on the same socket2, which a version bump would end. +socket2 = { version = "0.5", features = ["all"] } # Only for the hyper connection builder `axum_server::Server::http_builder()` # hands back — where `downstream.idle_timeout_secs` is applied. Same 0.1.x # axum-server itself depends on, so the builder types unify. diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 774d9cd7..7b819821 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -101,19 +101,17 @@ enum CliCommand { }, } -fn main() -> anyhow::Result<()> { - // Bench-only: BENCH_RT_WORKERS=N pins the tokio worker-thread count so - // saturation A/B runs control parallelism explicitly. Named without the - // AISIX_ prefix because that namespace is claimed by the config env - // override layer. - let mut rt = tokio::runtime::Builder::new_multi_thread(); - if let Ok(n) = std::env::var("BENCH_RT_WORKERS") { - rt.worker_threads(n.parse().expect("BENCH_RT_WORKERS must be a number")); - } - rt.enable_all().build()?.block_on(async_main()) -} +/// Threads left to the control surfaces when the proxy serves from +/// thread-per-core workers. +/// +/// The proxy's own runtimes belong to its workers in that mode, so this +/// runtime is only running the etcd watch, the admin and metrics +/// listeners, signal handling, and the background exporters. Two threads +/// keep a config reload from stalling behind a telemetry flush without +/// taking a core back from the workers. +const CONTROL_RUNTIME_THREADS: usize = 2; -async fn async_main() -> anyhow::Result<()> { +fn main() -> anyhow::Result<()> { // Install the process-level rustls CryptoProvider before anything // else touches TLS. rustls 0.23 dropped implicit provider selection // and panics at first use when both `aws-lc-rs` and `ring` features @@ -138,13 +136,15 @@ async fn async_main() -> anyhow::Result<()> { reveal_secrets, output, }) => { - return export::run(export::ExportArgs { - endpoints: etcd, - prefix, - reveal_secrets, - output, - }) - .await; + return tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()? + .block_on(export::run(export::ExportArgs { + endpoints: etcd, + prefix, + reveal_secrets, + output, + })); } None => {} } @@ -153,10 +153,24 @@ async fn async_main() -> anyhow::Result<()> { .config .expect("clap enforces --config unless a subcommand is given"); - // Steps 1-2: config. + // Steps 1-2: config. Read before the runtime exists because + // `proxy.thread_per_core` and `proxy.workers` decide how many threads + // this runtime gets — and, in thread-per-core mode, that it is not + // the runtime serving proxy traffic at all. let cfg = Config::load_from_path(Some(&config_path)) .map_err(|e| anyhow::anyhow!("config load failed: {e}"))?; + let mut builder = tokio::runtime::Builder::new_multi_thread(); + builder.enable_all(); + builder.worker_threads(if cfg.proxy.thread_per_core_enabled() { + CONTROL_RUNTIME_THREADS + } else { + cfg.proxy.worker_threads() + }); + builder.build()?.block_on(async_main(cfg)) +} + +async fn async_main(cfg: Config) -> anyhow::Result<()> { // Step 3: tracing + optional OTLP export. init_tracing(&cfg.observability).map_err(|e| anyhow::anyhow!("tracing init failed: {e}"))?; let _otlp = install_otlp_tracer(&cfg.observability) @@ -943,6 +957,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { downstream_idle_timeout, cancel_rx.clone(), "admin", + None, ))) } else { // Drop unused shared components so the compiler can see they @@ -988,6 +1003,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { downstream_idle_timeout, cancel_rx.clone(), "metrics", + None, ))) } else { None @@ -997,6 +1013,10 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { // Step 9: bind + serve the proxy (always). Admin is handled above. let proxy_addr: std::net::SocketAddr = cfg.proxy.addr.parse()?; let proxy_tls = cfg.proxy.tls.clone(); + let proxy_workers = cfg + .proxy + .thread_per_core_enabled() + .then(|| cfg.proxy.worker_threads()); let proxy_serve = serve_http( proxy_addr, proxy_router, @@ -1004,6 +1024,7 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { downstream_idle_timeout, cancel_rx.clone(), "proxy", + proxy_workers, ); // Step 10: shutdown coordinator. Whichever of (signal, proxy, admin) @@ -1402,21 +1423,8 @@ async fn serve_http( idle_timeout: Option, cancel: watch::Receiver, label: &'static str, + workers: Option, ) -> anyhow::Result<()> { - // Bench-only: BENCH_RT_TPC=N short-circuits the plain-HTTP proxy - // listener into N thread-per-core workers (one OS thread + one - // current_thread runtime + one listener each). BENCH_RT_TPC_MODE - // picks the listener layout: "reuseport" (default) binds every worker - // to the same port with SO_REUSEPORT so the kernel spreads accepted - // connections by 4-tuple hash; "stride" binds port, port+1, ... like - // the macOS phase did. - if label == "proxy" && tls.is_none() { - if let Ok(n) = std::env::var("BENCH_RT_TPC") { - let workers: usize = n.parse().expect("BENCH_RT_TPC must be a number"); - return serve_http_tpc(addr, router, idle_timeout, cancel, label, workers).await; - } - } - // Resolved before binding so a bad cert path still fails with the // same error it always did, before a port is taken. let tls_config = match tls { @@ -1434,6 +1442,21 @@ async fn serve_http( None => None, }; + // Thread-per-core serving is the proxy's; the admin and metrics + // surfaces are control traffic and keep their single listener. + if let Some(workers) = workers { + return serve_http_tpc( + addr, + router, + tls_config, + idle_timeout, + cancel, + label, + workers, + ) + .await; + } + let listener = std::net::TcpListener::bind(addr) .map_err(|e| anyhow::anyhow!("{label} listener bind {addr} failed: {e}"))?; listener.set_nonblocking(true)?; @@ -1473,99 +1496,196 @@ async fn serve_http( Ok(()) } -/// Bench-only (BENCH_RT_TPC): serve the proxy from N OS threads, each with -/// its own current_thread runtime and its own listener. The router (and so -/// the shared AppState / reqwest client) is cloned per thread, matching the -/// macOS-phase experiment semantics. +/// Serve from `workers` independent threads, each with its own runtime +/// and its own `SO_REUSEPORT` listener on the same address. +/// +/// A connection is accepted, read, dispatched, and answered on one +/// thread, and the upstream call it makes runs on that thread's own +/// connection pool (see `upstream_tls::mark_worker_thread`). That is the +/// point of the mode: the shared runtime hands a request between threads +/// roughly once per request — first when a worker steals the task, again +/// when the upstream response lands on whichever thread happens to own +/// that connection — and each handoff costs a wakeup and a context +/// switch. Here there are none. +/// +/// The kernel decides which listener gets each connection, hashing the +/// 4-tuple. That spreads evenly across many client connections and +/// unevenly across few, which is why the mode is documented as a +/// throughput setting rather than a latency one. async fn serve_http_tpc( addr: std::net::SocketAddr, router: axum::Router, + tls_config: Option, idle_timeout: Option, cancel: watch::Receiver, label: &'static str, workers: usize, ) -> anyhow::Result<()> { - let mode = std::env::var("BENCH_RT_TPC_MODE").unwrap_or_else(|_| "reuseport".to_string()); - let reuseport = match mode.as_str() { - "reuseport" => true, - "stride" => false, - other => anyhow::bail!("BENCH_RT_TPC_MODE must be reuseport|stride, got {other}"), - }; - let mut threads = Vec::with_capacity(workers); - for i in 0..workers { - let bind_addr = if reuseport { - addr - } else { - std::net::SocketAddr::new(addr.ip(), addr.port() + u16::try_from(i)?) - }; - // Bound on the parent thread so a bind failure aborts startup - // before any worker begins serving. - let listener = bind_tpc_listener(bind_addr, reuseport) - .map_err(|e| anyhow::anyhow!("{label} tpc listener bind {bind_addr} failed: {e}"))?; + // An address someone else already holds has to stay a loud boot + // failure. Every socket on a `SO_REUSEPORT` address has to set the + // option, so a socket that does not set it fails to bind exactly + // when something else is there — which is the check the worker + // sockets below deliberately give up, since they must co-bind with + // each other. Without this a second gateway would start silently + // and split traffic with the first. + drop( + std::net::TcpListener::bind(addr) + .map_err(|e| anyhow::anyhow!("{label} listener bind {addr} failed: {e}"))?, + ); + + // One slot per worker so no worker blocks reporting its exit. + let (exit_tx, exit_rx) = std::sync::mpsc::sync_channel::>(workers); + for worker in 0..workers { + // Bound on this thread so a bind failure aborts startup before + // any worker begins serving, exactly as the single listener does. + let listener = bind_reuseport_listener(addr) + .map_err(|e| anyhow::anyhow!("{label} listener bind {addr} failed: {e}"))?; listener.set_nonblocking(true)?; let router = router.clone(); let cancel = cancel.clone(); - let thread = std::thread::Builder::new().name(format!("tpc-{i}")).spawn( - move || -> anyhow::Result<()> { - let rt = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build()?; - rt.block_on(async move { - let handle = axum_server::Handle::new(); - tokio::spawn({ - let handle = handle.clone(); - async move { - shutdown_signal(cancel, label).await; - handle.graceful_shutdown(None); - } - }); - tracing::info!( - %bind_addr, - label, - worker = i, - reuseport, - "aisix listening (http, tpc)" - ); - let make_service = - router.into_make_service_with_connect_info::(); - let mut server = axum_server::from_tcp(listener).handle(handle); - apply_idle_timeout(server.http_builder(), idle_timeout); - server.serve(make_service).await?; - Ok(()) - }) - }, - )?; - threads.push(thread); + let tls_config = tls_config.clone(); + let exit_tx = exit_tx.clone(); + std::thread::Builder::new() + // Names the mode in `ps -T` / `top -H`: `tpc-N` here, + // tokio's own `tokio-rt-worker` on the shared runtime. + .name(format!("tpc-{worker}")) + .spawn(move || { + let mut exit = WorkerExit { + tx: exit_tx, + worker, + outcome: None, + }; + exit.outcome = Some(run_tpc_worker( + addr, + listener, + router, + tls_config, + idle_timeout, + cancel, + label, + worker, + )); + })?; } + // Only the workers hold senders from here, so a `RecvError` means + // every worker is gone. + drop(exit_tx); + + // The first worker to leave decides the outcome for all of them: a + // graceful shutdown ends them together, and any other exit — an + // accept loop failing, a panic unwinding — has to bring the process + // down rather than leave it serving on fewer listeners than it + // reported binding. tokio::task::spawn_blocking(move || { - for t in threads { - t.join() - .map_err(|_| anyhow::anyhow!("tpc worker thread panicked"))??; - } - Ok(()) + exit_rx + .recv() + .unwrap_or_else(|_| Err(anyhow::anyhow!("no proxy worker reported its exit"))) }) .await? } -/// Listener for the bench-only TPC path: plain bind for "stride" mode, -/// SO_REUSEPORT (set before bind) for same-port kernel load balancing. -fn bind_tpc_listener( +/// Reports a worker's exit exactly once, including when a panic unwinds +/// out of it and there is no return value to send. +struct WorkerExit { + tx: std::sync::mpsc::SyncSender>, + worker: usize, + outcome: Option>, +} + +impl Drop for WorkerExit { + fn drop(&mut self) { + let outcome = self.outcome.take().unwrap_or_else(|| { + Err(anyhow::anyhow!( + "proxy worker {} panicked; see the panic above", + self.worker + )) + }); + let _ = self.tx.send(outcome); + } +} + +/// One thread-per-core worker: its own current-thread runtime, its own +/// listener, its own upstream connection pool. +#[allow(clippy::too_many_arguments)] +fn run_tpc_worker( addr: std::net::SocketAddr, - reuseport: bool, -) -> std::io::Result { - let domain = if addr.is_ipv4() { - socket2::Domain::IPV4 - } else { - socket2::Domain::IPV6 - }; - let socket = socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))?; - socket.set_reuse_address(true)?; - if reuseport { + listener: std::net::TcpListener, + router: axum::Router, + tls_config: Option, + idle_timeout: Option, + cancel: watch::Receiver, + label: &'static str, + worker: usize, +) -> anyhow::Result<()> { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + // Every dispatch from this thread now uses this thread's pool, so an + // upstream response is read by the same runtime that is waiting for + // it. Marked inside the worker because the marker is per thread. + aisix_gateway::upstream_tls::mark_worker_thread(); + rt.block_on(async move { + let handle = axum_server::Handle::new(); + tokio::spawn({ + let handle = handle.clone(); + async move { + shutdown_signal(cancel, label).await; + // `None` = drain without a deadline, as on the shared + // runtime: an in-flight LLM stream can run for minutes. + handle.graceful_shutdown(None); + } + }); + + let make_service = router.into_make_service_with_connect_info::(); + match tls_config { + None => { + tracing::info!(%addr, label, worker, "aisix listening (http, thread-per-core)"); + let mut server = axum_server::from_tcp(listener).handle(handle); + apply_idle_timeout(server.http_builder(), idle_timeout); + server.serve(make_service).await?; + } + Some(tls_config) => { + tracing::info!(%addr, label, worker, "aisix listening (https, thread-per-core)"); + let mut server = axum_server::from_tcp_rustls(listener, tls_config).handle(handle); + apply_idle_timeout(server.http_builder(), idle_timeout); + server.serve(make_service).await?; + } + } + Ok(()) + }) +} + +/// A listener that shares `addr` with the other workers' listeners. +/// +/// `SO_REUSEPORT` has to be set before the bind, and every socket on the +/// address has to set it, which is why this cannot go through +/// `TcpListener::bind`. +fn bind_reuseport_listener(addr: std::net::SocketAddr) -> std::io::Result { + #[cfg(not(unix))] + return Err(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "thread-per-core serving needs SO_REUSEPORT, which this platform \ + does not have; set proxy.thread_per_core: false", + )); + + #[cfg(unix)] + { + let domain = if addr.is_ipv4() { + socket2::Domain::IPV4 + } else { + socket2::Domain::IPV6 + }; + let socket = + socket2::Socket::new(domain, socket2::Type::STREAM, Some(socket2::Protocol::TCP))?; + socket.set_reuse_address(true)?; socket.set_reuse_port(true)?; + socket.bind(&addr.into())?; + // Matches the backlog `std::net::TcpListener::bind` requests, so + // the two modes queue the same number of pending connections per + // socket. + socket.listen(128)?; + Ok(socket.into()) } - socket.bind(&addr.into())?; - socket.listen(1024)?; - Ok(socket.into()) } /// Close an accepted HTTP/1.1 connection that sits idle for `idle_timeout`. diff --git a/tests/e2e/src/cases/listener-tls-e2e.test.ts b/tests/e2e/src/cases/listener-tls-e2e.test.ts index 0b2d5704..8c9e6f5d 100644 --- a/tests/e2e/src/cases/listener-tls-e2e.test.ts +++ b/tests/e2e/src/cases/listener-tls-e2e.test.ts @@ -6,7 +6,7 @@ import { randomUUID } from "node:crypto"; import { stringify as yamlStringify } from "yaml"; import { Agent, request } from "undici"; import { afterAll, beforeAll, describe, expect, test } from "vitest"; -import { EtcdClient, pickFreePorts } from "../harness/index.js"; +import { EtcdClient, pickFreePorts, suiteThreadPerCore } from "../harness/index.js"; // E2E: `proxy.tls` / `admin.tls` actually serve HTTPS. // @@ -77,6 +77,13 @@ describe("listener TLS (#473)", () => { addr: `127.0.0.1:${proxyPort}`, request_body_limit_bytes: 10485760, tls: { cert_file: certFile, key_file: keyFile }, + // This file spawns the binary itself, so it has to honour the + // suite-wide serving mode the shared harness applies — HTTPS is + // served by a different code path in each mode, and this is the + // only test that covers either of them. + ...(suiteThreadPerCore !== undefined + ? { thread_per_core: suiteThreadPerCore } + : {}), }, admin: { addr: `127.0.0.1:${adminPort}`, diff --git a/tests/e2e/src/cases/upstream-pool-idle-e2e.test.ts b/tests/e2e/src/cases/upstream-pool-idle-e2e.test.ts index 574a8f9f..c874662b 100644 --- a/tests/e2e/src/cases/upstream-pool-idle-e2e.test.ts +++ b/tests/e2e/src/cases/upstream-pool-idle-e2e.test.ts @@ -131,9 +131,20 @@ describe("upstream pool idle timeout", () => { expiringUpstream = await startCountingUpstream(); holdingUpstream = await startCountingUpstream(); - expiring = await spawnApp({ extra: { upstream: { pool_idle_timeout_secs: POOL_IDLE_S } } }); + // One pool, so the connection counts below measure the idle deadline + // and nothing else. Under thread-per-core serving each worker keeps + // its own pool and the kernel decides which worker takes each + // connection, which would make "the second call reused the first + // call's connection" a question about that choice instead. + expiring = await spawnApp({ + threadPerCore: false, + extra: { upstream: { pool_idle_timeout_secs: POOL_IDLE_S } }, + }); // 0 switches the knob off, leaving reqwest's own 90s pool lifetime. - holding = await spawnApp({ extra: { upstream: { pool_idle_timeout_secs: 0 } } }); + holding = await spawnApp({ + threadPerCore: false, + extra: { upstream: { pool_idle_timeout_secs: 0 } }, + }); await seedInto(expiring, etcd, expiringUpstream.baseUrl); await seedInto(holding, etcd, holdingUpstream.baseUrl); diff --git a/tests/e2e/src/harness/app.ts b/tests/e2e/src/harness/app.ts index d5bc9ea8..8362df85 100644 --- a/tests/e2e/src/harness/app.ts +++ b/tests/e2e/src/harness/app.ts @@ -58,6 +58,17 @@ export interface AppOverrides { * AccessKey deliberately never travels on the config path. */ extraEnv?: Record; + /** + * `proxy.thread_per_core`. Omitted, the binary picks its platform + * default, which is what the suite should normally exercise. + * + * Pin it to `false` in a test whose subject is per-connection or + * per-pool state: with thread-per-core serving the kernel picks which + * worker accepts each connection, and each worker keeps its own + * upstream pool, so a count taken across two calls depends on that + * choice. Pinning keeps such an assertion measuring its own subject. + */ + threadPerCore?: boolean; /** * Log level for the spawned binary. Defaults to `warn` — quiet enough * that the suite's output stays readable. Tests that assert on a line @@ -119,6 +130,21 @@ const BIN_PATH = const READY_TIMEOUT_MS = 10_000; const SHUTDOWN_GRACE_MS = 3_000; +/** + * Suite-wide `proxy.thread_per_core`, from `E2E_THREAD_PER_CORE`, so CI + * can run the whole suite in each serving mode. Unset leaves the binary + * on its platform default; a per-test `threadPerCore` still wins over + * both. + * + * Every site that spawns the binary has to read this — a spawn site that + * ignores it silently stays on one mode forever, and the leg that was + * supposed to cover the other one goes green without ever running it. + */ +export const suiteThreadPerCore: boolean | undefined = + process.env.E2E_THREAD_PER_CORE === undefined + ? undefined + : process.env.E2E_THREAD_PER_CORE !== "false"; + /** * Per-test handle to a spawned `aisix` binary. Each call writes a fresh * config YAML into a tmp dir, picks three free ports (proxy, admin, @@ -211,6 +237,9 @@ async function spawnAppOnce(overrides: AppOverrides = {}): Promise { addr: `127.0.0.1:${proxyPort}`, request_body_limit_bytes: overrides.requestBodyLimitBytes ?? 10485760, ...(overrides.realIp ? { real_ip: overrides.realIp } : {}), + ...((overrides.threadPerCore ?? suiteThreadPerCore) !== undefined + ? { thread_per_core: overrides.threadPerCore ?? suiteThreadPerCore } + : {}), }, admin: adminEnabled ? { addr: `127.0.0.1:${adminPort}`, admin_keys: [adminKey] } diff --git a/tests/e2e/src/harness/index.ts b/tests/e2e/src/harness/index.ts index f5bd5ce6..7830d1fb 100644 --- a/tests/e2e/src/harness/index.ts +++ b/tests/e2e/src/harness/index.ts @@ -1,4 +1,4 @@ -export { spawnApp, type SpawnedApp, type AppOverrides } from "./app.js"; +export { spawnApp, suiteThreadPerCore, type SpawnedApp, type AppOverrides } from "./app.js"; export { AdminClient, waitConfigPropagation, awaitWindowHeadroom } from "./admin.js"; export { ProxyClient } from "./proxy.js"; export { EtcdClient } from "./etcd.js"; From d0e735adb33caa9cfaea7fd3fab4a157d3b84994 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 6 Aug 2026 19:04:15 +0800 Subject: [PATCH 3/4] fix(server): drain every thread-per-core worker before exiting on shutdown The exit watcher waited for a single worker report, so on SIGTERM the first idle worker to finish draining ended the process while siblings still carried in-flight streams - the graceful-shutdown window collapsed to ~0 in the mode that is the Linux default. Wait for every worker on the graceful path; any error exit and any exit without a shutdown signal stays immediately fatal. Also bind all listeners before spawning any worker, so a mid-loop bind failure (e.g. fd exhaustion) aborts startup before a single connection is accepted, matching the comment that claimed as much. The new shutdown-drain-e2e pins the contract in both serving modes: a request in flight at SIGTERM receives its complete response. --- crates/aisix-server/src/main.rs | 49 +++++-- .../e2e/src/cases/shutdown-drain-e2e.test.ts | 136 ++++++++++++++++++ 2 files changed, 172 insertions(+), 13 deletions(-) create mode 100644 tests/e2e/src/cases/shutdown-drain-e2e.test.ts diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 7d3ab188..f214ee30 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -1531,14 +1531,21 @@ async fn serve_http_tpc( .map_err(|e| anyhow::anyhow!("{label} listener bind {addr} failed: {e}"))?, ); - // One slot per worker so no worker blocks reporting its exit. - let (exit_tx, exit_rx) = std::sync::mpsc::sync_channel::>(workers); - for worker in 0..workers { - // Bound on this thread so a bind failure aborts startup before - // any worker begins serving, exactly as the single listener does. + // Every listener is bound before any worker spawns, so a bind + // failure — fd exhaustion on the last socket included — aborts + // startup before a single connection is accepted, exactly as the + // single listener does. + let mut listeners = Vec::with_capacity(workers); + for _ in 0..workers { let listener = bind_reuseport_listener(addr) .map_err(|e| anyhow::anyhow!("{label} listener bind {addr} failed: {e}"))?; listener.set_nonblocking(true)?; + listeners.push(listener); + } + + // One slot per worker so no worker blocks reporting its exit. + let (exit_tx, exit_rx) = std::sync::mpsc::sync_channel::>(workers); + for (worker, listener) in listeners.into_iter().enumerate() { let router = router.clone(); let cancel = cancel.clone(); let tls_config = tls_config.clone(); @@ -1569,15 +1576,31 @@ async fn serve_http_tpc( // every worker is gone. drop(exit_tx); - // The first worker to leave decides the outcome for all of them: a - // graceful shutdown ends them together, and any other exit — an - // accept loop failing, a panic unwinding — has to bring the process - // down rather than leave it serving on fewer listeners than it - // reported binding. + // A graceful shutdown ends the workers together, but each drains its + // own connections on its own clock — an idle worker returns at once + // while a sibling may hold an in-flight stream for minutes. Wait for + // every worker, so the fastest drain cannot end the process under + // the slowest. Everything else stays immediately fatal: an accept + // loop failing, a panic unwinding, or a worker stopping without a + // shutdown signal has to bring the process down rather than leave it + // serving on fewer listeners than it reported binding. + let shutdown_seen = cancel; tokio::task::spawn_blocking(move || { - exit_rx - .recv() - .unwrap_or_else(|_| Err(anyhow::anyhow!("no proxy worker reported its exit"))) + for _ in 0..workers { + match exit_rx.recv() { + Ok(Ok(())) if *shutdown_seen.borrow() => {} + Ok(Ok(())) => { + return Err(anyhow::anyhow!( + "a proxy worker stopped serving without a shutdown signal" + )); + } + Ok(Err(e)) => return Err(e), + Err(_) => { + return Err(anyhow::anyhow!("a proxy worker exited without reporting")); + } + } + } + Ok(()) }) .await? } diff --git a/tests/e2e/src/cases/shutdown-drain-e2e.test.ts b/tests/e2e/src/cases/shutdown-drain-e2e.test.ts new file mode 100644 index 00000000..472909e1 --- /dev/null +++ b/tests/e2e/src/cases/shutdown-drain-e2e.test.ts @@ -0,0 +1,136 @@ +import { createHash } from "node:crypto"; +import { createServer, type Server } from "node:http"; +import { afterAll, beforeAll, describe, expect, test } from "vitest"; +import { + EtcdClient, + SeedClient, + spawnApp, + waitConfigPropagation, + type SpawnedApp, +} from "../harness/index.js"; + +// E2E: graceful shutdown drains in-flight requests, in whichever serving +// mode the suite leg selects. +// +// Pinned contract: a request the gateway has accepted before SIGTERM +// receives its complete response, and only then does the process exit. +// In thread-per-core mode every worker drains on its own clock, and a +// worker with no connections finishes instantly — that fast drain must +// not end the process while a sibling still holds an in-flight upstream +// call (the regression this file guards). + +const KEY = "sk-shutdown-drain-e2e"; +const sha256 = (s: string) => createHash("sha256").update(s).digest("hex"); +// Long enough that SIGTERM reliably lands mid-flight, short enough to +// finish inside the harness's 3s SIGTERM→SIGKILL escalation window. +const UPSTREAM_DELAY_MS = 2_000; + +describe("graceful shutdown drains in-flight requests", () => { + let app: SpawnedApp | undefined; + let upstream: Server | undefined; + let etcdReachable = false; + let sawRequest = () => {}; + const upstreamGotRequest = new Promise((resolve) => { + sawRequest = resolve; + }); + + beforeAll(async () => { + const etcd = new EtcdClient(); + etcdReachable = await etcd.ping(); + if (!etcdReachable) return; + + // A deliberately slow upstream: answers every POST with a valid + // chat completion after UPSTREAM_DELAY_MS, so one request is + // reliably in flight when SIGTERM lands. + upstream = createServer((req, res) => { + req.resume(); + req.on("end", () => { + sawRequest(); + setTimeout(() => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + id: "chatcmpl-drain", + object: "chat.completion", + created: 1, + model: "gpt-4o-mini", + choices: [ + { + index: 0, + message: { role: "assistant", content: "drained" }, + finish_reason: "stop", + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, + }), + ); + }, UPSTREAM_DELAY_MS); + }); + }); + await new Promise((resolve) => upstream!.listen(0, "127.0.0.1", resolve)); + const upstreamPort = (upstream!.address() as { port: number }).port; + + app = await spawnApp({}); + const seed = new SeedClient(new EtcdClient(), app.etcdPrefix); + const pk = await seed.createProviderKey({ + display_name: "drain-pk", + secret: "sk-mock", + api_base: `http://127.0.0.1:${upstreamPort}/v1`, + }); + await seed.createModel({ + display_name: "drain-model", + provider: "openai", + model_name: "gpt-4o-mini", + provider_key_id: pk.id, + }); + await seed.createApiKey({ + key_hash: sha256(KEY), + allowed_models: ["drain-model"], + }); + + await waitConfigPropagation(async () => { + const res = await fetch(`${app!.proxyUrl}/v1/models`, { + headers: { authorization: `Bearer ${KEY}` }, + }); + if (res.status !== 200) return false; + const body = (await res.json()) as { data?: Array<{ id?: string }> }; + return body.data?.some((m) => m.id === "drain-model") ?? false; + }); + }, 60_000); + + afterAll(async () => { + await app?.exit(); + if (upstream) await new Promise((resolve) => upstream!.close(() => resolve())); + }); + + test("a request in flight at SIGTERM still completes", async (ctx) => { + if (!etcdReachable || !app) return ctx.skip(); + + const inFlight = fetch(`${app.proxyUrl}/v1/chat/completions`, { + method: "POST", + headers: { + authorization: `Bearer ${KEY}`, + "content-type": "application/json", + }, + body: JSON.stringify({ + model: "drain-model", + messages: [{ role: "user", content: "hold the line" }], + }), + }); + + // Only signal once the request has actually reached the upstream — + // from here the gateway holds it in flight for UPSTREAM_DELAY_MS. + await upstreamGotRequest; + app.signal("SIGTERM"); + + // The drain contract: the response completes despite the shutdown. + // A process that exits from under the request turns this await into + // a socket error instead of a response. + const res = await inFlight; + expect(res.status).toBe(200); + const body = (await res.json()) as { + choices?: Array<{ message?: { content?: string } }>; + }; + expect(body.choices?.[0]?.message?.content).toBe("drained"); + }, 20_000); +}); From b31de2036e53af0b8b7fc8d5d48798eaa4598645 Mon Sep 17 00:00:00 2001 From: Yuansheng Date: Thu, 6 Aug 2026 19:04:15 +0800 Subject: [PATCH 4/4] fix(e2e,ci): close two serving-mode coverage gaps - status-config-e2e spawns the binary itself but ignored the suite-wide serving mode, so the work-stealing CI leg silently ran it thread-per-core (the Linux platform default). Spread suiteThreadPerCore into its proxy config like listener-tls-e2e does. - Both e2e legs uploaded coverage as the same lcov.info filename, which merge-multiple flattens to one file - the gate would measure a single leg the moment the suite starts emitting coverage. Tag the file per leg before upload. --- .github/workflows/ci.yml | 11 ++++++++++- tests/e2e/src/cases/status-config-e2e.test.ts | 9 ++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 64ada348..9e45dc13 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -293,13 +293,22 @@ jobs: env: E2E_THREAD_PER_CORE: ${{ matrix.thread_per_core }} run: pnpm test + - name: tag coverage by leg + # The gate downloads every leg with `merge-multiple: true`, which + # flattens same-named files to one — the legs must upload distinct + # filenames or only one leg's coverage survives the merge. + if: always() + run: | + if [ -f tests/e2e/coverage/lcov.info ]; then + mv tests/e2e/coverage/lcov.info "tests/e2e/coverage/lcov-${{ matrix.serving }}.info" + fi - uses: actions/upload-artifact@v7 if: always() with: # Per leg: two uploads under one name is a hard conflict the # moment the suite starts emitting coverage. name: coverage-e2e-${{ matrix.serving }} - path: tests/e2e/coverage/lcov.info + path: tests/e2e/coverage/lcov-${{ matrix.serving }}.info if-no-files-found: ignore retention-days: 7 diff --git a/tests/e2e/src/cases/status-config-e2e.test.ts b/tests/e2e/src/cases/status-config-e2e.test.ts index 1996d7c2..80172a36 100644 --- a/tests/e2e/src/cases/status-config-e2e.test.ts +++ b/tests/e2e/src/cases/status-config-e2e.test.ts @@ -12,6 +12,7 @@ import { pickFreePorts, spawnApp, startOpenAiUpstream, + suiteThreadPerCore, waitConfigPropagation, type OpenAiUpstream, type SpawnedApp, @@ -345,7 +346,13 @@ async function spawnPointedAtDeadEtcd(): Promise { dial_timeout_ms: 2000, request_timeout_ms: 2000, }, - proxy: { addr: `127.0.0.1:${proxyPort}`, request_body_limit_bytes: 10485760 }, + proxy: { + addr: `127.0.0.1:${proxyPort}`, + request_body_limit_bytes: 10485760, + // This file spawns the binary itself, so it has to honour the + // suite-wide serving mode the shared harness applies. + ...(suiteThreadPerCore !== undefined ? { thread_per_core: suiteThreadPerCore } : {}), + }, admin: { addr: `127.0.0.1:${adminPort}`, admin_keys: [`admin-${randomUUID()}`] }, observability: { service_name: "aisix-status-nl",