diff --git a/.github/workflows/gateway-proxy-tests.yml b/.github/workflows/gateway-proxy-tests.yml new file mode 100644 index 000000000..9a91d8031 --- /dev/null +++ b/.github/workflows/gateway-proxy-tests.yml @@ -0,0 +1,81 @@ +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 + +name: Gateway proxy tests + +# The gateway's proxy data path has two opt-in optimisations (`tcp_splice`, +# `ktls`) whose behaviour depends on kernel capabilities and on a per-connection +# gate. Unit tests cover the relay functions; this runs a real gateway process +# and asserts on what actually reaches the wire. +on: + push: + branches: [ master, next, dev-* ] + paths: + - 'dstack/gateway/**' + - 'dstack/vendor/ktls/**' + - '.github/workflows/gateway-proxy-tests.yml' + pull_request: + branches: [ master, next, dev-* ] + paths: + - 'dstack/gateway/**' + - 'dstack/vendor/ktls/**' + - '.github/workflows/gateway-proxy-tests.yml' + +permissions: + contents: read + +env: + CARGO_TERM_COLOR: always + +jobs: + proxy-integration: + runs-on: ${{ vars.CI_RUNNER || 'ubuntu-latest' }} + # Each of the ~25 arms restarts the gateway, and the idle-timeout arms wait + # out a real timeout, so this is minutes rather than seconds. + timeout-minutes: 30 + steps: + - uses: actions/checkout@v5 + + - name: Install Rust + uses: dtolnay/rust-toolchain@1.92.0 + + - name: Cache cargo + uses: actions/cache@v4 + with: + path: | + ~/.cargo/registry + ~/.cargo/git + dstack/target + key: gateway-proxy-${{ runner.os }}-${{ hashFiles('dstack/Cargo.lock') }} + restore-keys: gateway-proxy-${{ runner.os }}- + + - name: Build the gateway + working-directory: dstack + run: cargo build --release -p dstack-gateway + + - name: Record kernel capabilities + # The suite adapts to what the kernel offers, so the log needs to say + # what it had: a run that skipped kTLS looks the same as one that + # covered it otherwise. + run: | + echo "kernel: $(uname -r)" + sudo modprobe tls 2>&1 || echo "no TLS ULP available" + echo "tls module loaded: $(lsmod | grep -c '^tls ' || true)" + grep -B2 -A3 'gcm(aes)' /proc/crypto | grep -E '^(driver|priority)' \ + | paste - - | sort -u || true + + - name: Proxy integration tests + working-directory: dstack/gateway/test-run + env: + GATEWAY_BIN: ${{ github.workspace }}/dstack/target/release/dstack-gateway + run: ./test_proxy.sh + + - name: Upload logs on failure + if: failure() + uses: actions/upload-artifact@v4 + with: + name: gateway-proxy-test-logs + path: /tmp/dstack-gw-proxy-test.*/logs/ + if-no-files-found: ignore + retention-days: 7 diff --git a/REUSE.toml b/REUSE.toml index 16f7a6031..236f5b008 100644 --- a/REUSE.toml +++ b/REUSE.toml @@ -181,6 +181,13 @@ SPDX-FileCopyrightText = "NONE" SPDX-License-Identifier = "Apache-2.0" precedence = "override" +# ktls 6.0.2 plus rustls/ktls#70; see dstack/vendor/README.md. +[[annotations]] +path = "dstack/vendor/ktls/**" +SPDX-FileCopyrightText = "Copyright (c) 2022 Amos Wenger " +SPDX-License-Identifier = "MIT OR Apache-2.0" +precedence = "override" + # Generated files [[annotations]] diff --git a/dstack/Cargo.lock b/dstack/Cargo.lock index 0ece65480..1bbfd4d31 100644 --- a/dstack/Cargo.lock +++ b/dstack/Cargo.lock @@ -1887,6 +1887,8 @@ dependencies = [ "insta", "ipnet", "jemallocator", + "ktls", + "libc", "load_config", "nix 0.29.0", "or-panic", @@ -1908,6 +1910,7 @@ dependencies = [ "sha2 0.10.9", "shared_child", "smallvec", + "socket2 0.5.10", "tdx-attest", "tempfile", "tokio", @@ -3468,7 +3471,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2", + "socket2 0.6.4", "tokio", "tower-service", "tracing", @@ -3774,7 +3777,7 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4d40460c0ce33d6ce4b0630ad68ff63d6661961c48b6dba35e5a4d81cfb48222" dependencies = [ - "socket2", + "socket2 0.6.4", "widestring", "windows-registry", "windows-result 0.4.1", @@ -4035,6 +4038,31 @@ dependencies = [ "libc", ] +[[package]] +name = "ktls" +version = "6.0.2" +dependencies = [ + "futures-util", + "ktls-sys", + "libc", + "memoffset 0.9.1", + "nix 0.29.0", + "num_enum", + "pin-project-lite", + "rustls", + "smallvec", + "thiserror 1.0.69", + "tokio", + "tokio-rustls", + "tracing", +] + +[[package]] +name = "ktls-sys" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ed84c81d133bc00e291085cff51c15cb07d3cede0aade1f2d82dd33df82d7e7" + [[package]] name = "lazy_static" version = "1.5.0" @@ -4522,6 +4550,7 @@ dependencies = [ "cfg-if", "cfg_aliases", "libc", + "memoffset 0.9.1", ] [[package]] @@ -4746,6 +4775,28 @@ dependencies = [ "libc", ] +[[package]] +name = "num_enum" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26" +dependencies = [ + "num_enum_derive", + "rustversion", +] + +[[package]] +name = "num_enum_derive" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8" +dependencies = [ + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "nvml-wrapper" version = "0.12.1" @@ -5542,7 +5593,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls", - "socket2", + "socket2 0.6.4", "thiserror 2.0.18", "tokio", "tracing", @@ -5580,7 +5631,7 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2", + "socket2 0.6.4", "tracing", "windows-sys 0.60.2", ] @@ -6476,7 +6527,7 @@ dependencies = [ "lazy_static", "libc", "s2n-quic-core", - "socket2", + "socket2 0.6.4", "tokio", ] @@ -6722,6 +6773,7 @@ name = "serde-duration" version = "0.6.0" dependencies = [ "serde", + "serde_json", ] [[package]] @@ -7158,6 +7210,16 @@ dependencies = [ "syn 1.0.109", ] +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "socket2" version = "0.6.4" @@ -7654,7 +7716,7 @@ dependencies = [ "parking_lot 0.12.5", "pin-project-lite", "signal-hook-registry", - "socket2", + "socket2 0.6.4", "tokio-macros", "windows-sys 0.61.2", ] diff --git a/dstack/Cargo.toml b/dstack/Cargo.toml index ca5375628..848809b03 100644 --- a/dstack/Cargo.toml +++ b/dstack/Cargo.toml @@ -71,6 +71,11 @@ members = [ "crates/build-info", "crates/mock-attestation", ] +# Vendored third-party crates are path dependencies but deliberately not members: +# `--all-features` applies to members, and ktls declares its `ring` and +# `aws_lc_rs` features mutually exclusive, so membership makes +# `cargo test --all-features` fail to compile. +exclude = ["vendor/ktls"] resolver = "2" [workspace.dependencies] @@ -141,6 +146,9 @@ tracing-subscriber = { version = "0.3.20", features = ["env-filter"] } safe-write = "0.1.2" rustix = { version = "0.38", features = ["fs"] } nix = "0.29.0" +# Vendored: 6.0.2 does not build for musl. See vendor/README.md. +ktls = { path = "vendor/ktls" } +socket2 = { version = "0.5", features = ["all"] } sd-notify = "0.4.5" listenfd = "1.0" jemallocator = "0.5.4" diff --git a/dstack/gateway/Cargo.toml b/dstack/gateway/Cargo.toml index f81713bfe..b4c05113c 100644 --- a/dstack/gateway/Cargo.toml +++ b/dstack/gateway/Cargo.toml @@ -67,7 +67,10 @@ dstack-api-auth.workspace = true cached-cell.workspace = true [target.'cfg(unix)'.dependencies] -nix = { workspace = true, features = ["resource"] } +nix = { workspace = true, features = ["resource", "fs", "socket", "zerocopy"] } +ktls.workspace = true +libc.workspace = true +socket2.workspace = true [[bin]] name = "gen_debug_key" @@ -76,3 +79,6 @@ path = "src/gen_debug_key.rs" [dev-dependencies] insta.workspace = true tempfile.workspace = true +# `test-util` gives the idle-watchdog tests a paused clock, so they assert on +# the window without waiting for it in wall-clock time. +tokio = { workspace = true, features = ["test-util"] } diff --git a/dstack/gateway/gateway.toml b/dstack/gateway/gateway.toml index 1772a67e4..a93bff51b 100644 --- a/dstack/gateway/gateway.toml +++ b/dstack/gateway/gateway.toml @@ -77,18 +77,44 @@ tls_versions = ["1.2"] listen_addr = "0.0.0.0" listen_port = 8443 agent_port = 8090 -buffer_size = 8192 +buffer_size = 65536 # number of hosts to try to connect to connect_top_n = 3 localhost_enabled = false app_address_ns_prefix = "_dstack-app-address" app_address_ns_compat = true workers = 32 +# One runtime + SO_REUSEPORT listener per worker (thread-per-core). Linux only. +thread_per_core = true +# Hand new connections to a less loaded core when SO_REUSEPORT skews them. +connection_rebalance = true external_port = 443 # Maximum concurrent connections per app. 0 means unlimited. max_connections_per_app = 2000 # Whether to read PROXY protocol from inbound connections (e.g. from Cloudflare). inbound_pp_enabled = false +# splice(2) zero-copy relaying for the TLS-passthrough path (Linux only), and +# kernel TLS offload for the terminate path, are both configured by an optional +# section. Omitting the section disables the optimisation; an empty section +# engages it from the first byte; `after_bytes` / `after_duration` are two +# independent gates and whichever fires first engages it. +# +# [core.proxy.tcp_splice] +# # Bulk transfers trip this within milliseconds. +# after_bytes = 65536 +# # Long-lived low-rate streams (e.g. LLM token streaming) never accumulate +# # bytes fast enough for the gate above, so they need this one instead. +# after_duration = "5s" +# # Park a relay's pipe in the pool while it waits for the next chunk. A spliced +# # connection otherwise pins 4 descriptors for its whole lifetime, so bursty +# # traffic (a token every 25 ms) holds them >99% idle. +# release_idle_pipes = true +# +# kTLS hands session keys to the kernel; see config.rs for the security note. +# Gated offload also requires [core.proxy.tcp_splice]. +# +# [core.proxy.ktls] +# after_bytes = 65536 [core.proxy.port_policy_fetch] # Background lazy-fetch of port_policy from legacy CVM agents. diff --git a/dstack/gateway/rpc/proto/gateway_rpc.proto b/dstack/gateway/rpc/proto/gateway_rpc.proto index 046b731ed..54778d8b7 100644 --- a/dstack/gateway/rpc/proto/gateway_rpc.proto +++ b/dstack/gateway/rpc/proto/gateway_rpc.proto @@ -104,6 +104,28 @@ message StatusResponse { repeated GatewayNodeInfo nodes = 6; // Peer uuid bytes uuid = 7; + // What the data path is actually doing with kTLS and splice. + ProxyAccelStatus accel = 8; +} + +// Effective state of the data-path acceleration options. +// +// The configured value does not say what is running: the startup probe can turn +// kTLS off on a kernel without the TLS ULP, and both options engage per +// connection only once their gate fires. These are reported so an operator can +// tell the difference without reading the startup log. +message ProxyAccelStatus { + // Effective kTLS mode: "off", "immediate", "after 64 KiB or 5s", or + // "disabled (kernel has no TLS ULP)" when the startup probe cleared it. + string ktls_mode = 1; + // Effective splice mode, same encoding minus the probe case. + string splice_mode = 2; + // Connections handed to the kernel's TLS ULP since start. + uint64 ktls_offloaded = 3; + // Connections whose handover to the kernel failed. + uint64 ktls_offload_failed = 4; + // Connections that entered a zero-copy splice relay, on either path. + uint64 splice_engaged = 5; } // HostInfo is the information of a host. diff --git a/dstack/gateway/src/admin_service.rs b/dstack/gateway/src/admin_service.rs index 8a02bdb19..ca04693ad 100644 --- a/dstack/gateway/src/admin_service.rs +++ b/dstack/gateway/src/admin_service.rs @@ -31,7 +31,7 @@ use crate::{ kv::{DnsCredential, DnsProvider, NodeStatus, PortFlags, PortPolicy, ZtDomainConfig}, main_service::Proxy, models::PortPolicyView, - proxy::NUM_CONNECTIONS, + proxy::{stats::accel_status, NUM_CONNECTIONS}, }; pub struct AdminRpcHandler { @@ -74,6 +74,9 @@ impl AdminRpcHandler { nodes: state.get_all_nodes(), hosts, num_connections: NUM_CONNECTIONS.load(Ordering::Relaxed), + // Reads the post-probe config, so this is what the data path is + // running rather than what the file asked for. + accel: Some(accel_status(&state.config.proxy)), }) } } diff --git a/dstack/gateway/src/config.rs b/dstack/gateway/src/config.rs index f3a8dcf58..94f37c8f2 100644 --- a/dstack/gateway/src/config.rs +++ b/dstack/gateway/src/config.rs @@ -11,7 +11,7 @@ use rocket::figment::Figment; use serde::{Deserialize, Serialize}; use std::net::Ipv4Addr; use std::path::PathBuf; -use std::time::Duration; +use std::time::{Duration, Instant}; use tracing::info; #[derive(Debug, Clone, Deserialize)] @@ -103,6 +103,10 @@ where } } +fn default_true() -> bool { + true +} + #[derive(Debug, Clone, Deserialize)] pub struct ProxyConfig { pub tls_crypto_provider: CryptoProvider, @@ -111,10 +115,62 @@ pub struct ProxyConfig { #[serde(deserialize_with = "deserialize_port_range")] pub listen_port: Vec, pub timeouts: Timeouts, + /// Relay buffer size, per direction, for connections that copy through + /// userspace -- TLS terminate, and passthrough before the splice gate. + /// + /// Costs `2 * buffer_size` of address space per such connection, of which + /// only the pages actually touched become resident: measured at 2 000 + /// concurrent streaming connections, the userspace relay path sat at ~52 KB + /// RSS per connection. Budget for it before raising this on a gateway that + /// fronts many idle-ish connections; the same measurement with kTLS, where + /// the payload is spliced and never enters the process, was ~12 KB. + /// + /// 64 KiB is the bulk-throughput sweet spot: it is large enough to keep a + /// 1 MiB pipe fed without the syscall rate 8 KiB imposed. pub buffer_size: usize, pub connect_top_n: usize, pub localhost_enabled: bool, pub workers: usize, + /// Run one single-threaded runtime per worker, each with its own + /// `SO_REUSEPORT` listener, instead of one accept thread feeding a shared + /// work-stealing runtime. + /// + /// A connection is then accepted and served entirely on one thread. The + /// default model costs ~0.6 context switches per request (accept-thread + /// handoff plus work-stealing migrations); HAProxy's thread-per-core design + /// measures ~0. Linux-only (needs SO_REUSEPORT). + #[serde(default)] + pub thread_per_core: bool, + /// Hand a freshly accepted connection to a less loaded core when the + /// accepting one is running ahead. + /// + /// `SO_REUSEPORT` picks a listener by hashing the connection's 4-tuple and + /// thread-per-core cannot move a connection afterwards, so a core can sit + /// starved: measured at 16 connections over 4 cores, utilisation came out + /// `[99, 42, 101, 101]`. Rebalancing costs one channel send per *rebalanced + /// connection*, never per request. Requires `thread_per_core`. + /// + /// On by default. It used to be opt-in, because handing a connection over + /// cost 2-3% wherever the hash was already even -- that turned out to be a + /// bug (the migrated socket kept its registration on the accepting core's + /// reactor), and with it fixed the trade is one-sided. Measured on a 4-core + /// gateway, 3-5 runs per arm after warmup: + /// + /// | workload | off | on | | + /// |---|---|---|---| + /// | passthrough small-request, 8 conns | 113 350 | 143 341 | **+26.5%** | + /// | passthrough small-request, 16 conns | 231 043 | 258 787 | **+12.0%** | + /// | passthrough small-request, 50 conns | 294 937 | 295 311 | +0.1% | + /// | TLS terminate small-request, 50 conns | 244 304 | 246 228 | +0.8% | + /// | TLS terminate, connections/s | 41 024 | 43 019 | +4.9% | + /// | passthrough, connections/s | 17 431 | 17 319 | -0.6% | + /// | passthrough, bulk throughput | 12.10 GB/s | 12.16 GB/s | +0.5% | + /// + /// The gain is largest where the hash has fewest connections to spread and + /// the worst-loaded core would otherwise starve: at 16 connections the + /// quietest core goes from 63% to 97% busy. + #[serde(default = "default_true")] + pub connection_rebalance: bool, #[serde(default)] pub base_domain: Option, #[serde(default)] @@ -132,10 +188,229 @@ pub struct ProxyConfig { /// (e.g. when behind a PP-aware load balancer like Cloudflare). #[serde(default)] pub inbound_pp_enabled: bool, + /// Use `splice(2)` zero-copy relaying for the TLS-passthrough path. Both + /// sides are raw TCP there, so payload never needs to enter userspace. + /// Linux-only; ignored for the TLS-terminate path. + /// + /// Absent disables splice entirely; see [`EngageAfter`] for what a present + /// section means. + /// + /// Tradeoff (measured on a 4-core gateway): bulk passthrough throughput + /// +~12% with a lower tail latency under load, but small-request latency + /// regresses (each tiny message pays an extra pipe hop). splice costs ~17 + /// syscalls per connection to move a small response (fill pipe, drain pipe, + /// readiness retries) where a read/write pair needs two: its benefit is per + /// byte, its cost is per connection, which is what the gates amortise. + #[serde(default)] + pub tcp_splice: Option, + /// Offload TLS record encryption to the kernel (kTLS) on the + /// TLS-terminate path. The handshake still runs in rustls; only the + /// symmetric crypto moves into the kernel afterwards. Linux-only. + /// + /// Absent disables kTLS entirely; see [`EngageAfter`] for what a present + /// section means. Gated offload additionally requires `tcp_splice`, since + /// the point of handing the socket to the kernel is to then splice it. + /// + /// A kernel built without `CONFIG_TLS` cannot honour this, so startup + /// probes for the TLS ULP and clears this section with a warning if it is + /// missing, rather than letting every connection discover it at the gate. + /// + /// kTLS costs ~30% of connection setup rate but wins ~25% on bulk + /// throughput, so paying the setup cost up front is wrong for short + /// request/response connections. + /// + /// On token-streaming traffic the throughput win does not materialise, but + /// a memory win does. Measured on the terminate path, 10k connections with + /// 2k streaming a 64 B record every 25 ms, 2 runs per arm: + /// + /// | | userspace rustls | kTLS | + /// |---|---|---| + /// | latency p50 | 0.174 / 0.181 ms | 0.180 / 0.174 ms | + /// | latency p99 | 0.602 / 0.526 ms | 0.859 / 0.537 ms | + /// | **RSS** | **349 MB** | **206 MB** | + /// + /// Latency is unchanged, as expected: at 64 B per record the per-record + /// overhead dominates and there is almost no symmetric crypto to move into + /// the kernel. The 41% RSS drop is the real effect and was not predicted -- + /// with kTLS the payload is spliced without ever entering this process, so + /// the per-connection userspace relay buffers disappear (~14 KB/connection + /// here). Weigh that against handing session keys to the kernel. + /// + /// Security note: enabling this hands the negotiated session keys to the + /// kernel via `dangerous_extract_secrets`, so the keys live outside + /// rustls' control. Inside a CVM the kernel is part of the measured TCB, + /// but on a non-TEE host this widens key exposure. Off by default. + #[serde(default)] + pub ktls: Option, /// Background lazy-fetch behaviour for `port_policy` (legacy CVMs). pub port_policy_fetch: PortPolicyFetchConfig, } +impl ProxyConfig { + /// The idle window every relay enforces, or `None` when data timeouts are + /// off. Computed in one place so the buffered bridge and the gated fast + /// paths cannot end up enforcing different things. + pub fn idle_timeout(&self) -> Option { + self.timeouts + .data_timeout_enabled + .then_some(self.timeouts.idle) + } +} + +/// Configuration for `splice(2)` relaying on the TLS-passthrough path. +#[derive(Debug, Clone, Default, Deserialize)] +pub struct SpliceConfig { + /// When splice should take over from the buffered relay. + #[serde(flatten)] + pub engage: EngageAfter, + /// Return a splice pipe to the thread-local pool while waiting for the + /// next chunk, instead of holding it for the connection's lifetime. + /// + /// A relay holds one pipe per direction, so a spliced connection pins four + /// descriptors. Held for the connection's lifetime that is proportional to + /// *connections*: 50k streaming connections need 300k descriptors. Released + /// while idle it is proportional to *chunks actually in flight*, which for + /// bursty traffic is far smaller -- LLM token streaming moves a ~64 B record + /// every 25 ms and spends over 99% of the connection idle. + /// + /// The pipe is provably empty at the point it is released: the relay only + /// waits for readability after the previous chunk has been fully drained + /// into the destination, so nothing is left to corrupt the next borrower. + /// + /// The cost is a pool pop and push per idle-to-active transition, both + /// `Vec` operations on a thread-local. Pipes are only created when the pool + /// is empty, so the live pipe count converges on the peak number of + /// concurrent in-flight chunks rather than churning `pipe2`/`close`. + /// + /// Measured on a 4-core gateway streaming a 64 B record every 25 ms per + /// connection, 2 runs per arm at each scale: + /// + /// | | off | on | + /// |---|---|---| + /// | pipe fds, 2k connections | 8 000 | **8** | + /// | pipe fds, 50k conns / 10k streams | 53 864 / 55 212 | **8** | + /// | RSS, 50k connections | 1 203 MB | 1 202 / 1 204 MB | + /// | latency p50, 50k | 20.9 / 20.3 ms | 21.6 / 21.7 ms | + /// | latency p999, 50k | 111 / 113 ms | 68 / 79 ms | + /// + /// The descriptor result is the point, and it is flat in the connection + /// count: 8 descriptors is four pipes, one per worker thread, for the whole + /// gateway, at 2k connections and at 50k alike. That is the bound this knob + /// exists to impose. + /// + /// Latency is close to a wash and should not be used to justify the knob. At + /// 2k connections there was no consistent difference at all (the single + /// fastest run of seven was an `off` run). At 50k, where the box is + /// saturated, `on` costs ~1 ms on p50 and saves ~40 ms on p999; the p999 + /// gain is consistent across repeats but comes from a regime that is already + /// over budget. RSS is unchanged either way. + /// + /// Mixing bulk transfers with token streams does not break the pooling, and + /// this was the failure worth checking: a connection that is never idle + /// never reaches the release point, so bulk traffic could in principle hold + /// every pipe and force each token to allocate a fresh one. Measured with + /// 10k streaming connections alongside 256 bulk connections saturating the + /// passthrough path at 12.8 GB/s, 2 runs per arm: + /// + /// | | off | on | + /// |---|---|---| + /// | pipe fds | 9 024 (= 2000*4 + 256*4) | **8** | + /// | bulk throughput | 12.8 GB/s | 12.8 GB/s | + /// | stream latency p50 | 0.099 / 0.102 ms | 0.093 / 0.089 ms | + /// | bulk latency p999 | 14.3 / 11.4 ms | 6.9 / 7.7 ms | + /// + /// The pool never degenerates because even a saturated bulk connection is + /// idle for tens of microseconds between chunks (measured inter-record gap + /// 62 us) and releases in that window. Throughput is identical and both + /// traffic classes are slightly better off with `on`, so there is no bulk + /// regression to trade against the descriptor saving. + #[serde(default)] + pub release_idle_pipes: bool, +} + +/// When an adaptive optimisation should engage on a connection. +/// +/// Both gates are optional and independent, and the optimisation engages as +/// soon as *either* fires. They catch different traffic and neither subsumes +/// the other: +/// +/// - `after_bytes` catches high-rate connections almost immediately -- a bulk +/// transfer trips a 64 KiB gate within milliseconds -- but is blind to +/// long-lived low-rate streams. LLM token streaming at 40 tok/s of ~64 B +/// records needs ~25 s of wall time to move 64 KiB, so a byte gate leaves the +/// whole early phase of every stream on the copy path, and never promotes +/// short conversations at all. +/// Measured: with a 64 KiB byte gate alone, connections streaming a 64 B +/// record every 25 ms promoted 25 s after they started streaming, matching +/// `65536 / (64 B * 40/s)`. Adding `after_duration = "5s"` moved that to 5 s, +/// a 5x earlier handover, with no change in added latency or RSS. +/// - `after_duration` catches exactly those long-lived low-rate streams, but is +/// blind to short high-rate ones, which finish before it fires. +/// +/// With neither gate set there is nothing to wait for, so the optimisation +/// engages from the first byte. "Never engage" is expressed by omitting the +/// whole section rather than by a sentinel value here, so every state has +/// exactly one representation. +#[derive(Debug, Clone, Default, Deserialize)] +pub struct EngageAfter { + /// Bytes the connection must transfer first. Absent = this gate never + /// fires. + #[serde(default)] + pub after_bytes: Option, + /// Wall time the connection must stay alive first, measured from the point + /// the relay starts (upstream already connected). Absent = this gate never + /// fires. + #[serde(default, with = "serde_duration::option")] + pub after_duration: Option, +} + +impl EngageAfter { + /// No gate configured, so there is nothing to wait for. + pub fn is_immediate(&self) -> bool { + self.after_bytes.is_none() && self.after_duration.is_none() + } + + /// Whether either gate has been reached. + /// + /// `start` is only read when the duration gate is configured, so a + /// bytes-only config pays no clock read per message. + pub fn reached(&self, moved: u64, start: Instant) -> bool { + self.after_bytes.is_some_and(|bytes| moved >= bytes) + || self + .after_duration + .is_some_and(|limit| start.elapsed() >= limit) + } +} + +/// Rendered for the dashboard and the `Status` RPC, so it reads as the answer to +/// "when does this engage?" rather than as a struct dump. +impl std::fmt::Display for EngageAfter { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match (self.after_bytes, self.after_duration) { + (None, None) => write!(f, "immediate"), + (Some(bytes), None) => write!(f, "after {}", DisplayBytes(bytes)), + (None, Some(after)) => write!(f, "after {after:?}"), + (Some(bytes), Some(after)) => write!(f, "after {} or {after:?}", DisplayBytes(bytes)), + } + } +} + +/// A byte threshold in the units the config file writes it in: binary units +/// when they divide evenly, raw bytes otherwise. +struct DisplayBytes(u64); + +impl std::fmt::Display for DisplayBytes { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + const KIB: u64 = 1 << 10; + const MIB: u64 = 1 << 20; + match self.0 { + bytes if bytes >= MIB && bytes % MIB == 0 => write!(f, "{} MiB", bytes / MIB), + bytes if bytes >= KIB && bytes % KIB == 0 => write!(f, "{} KiB", bytes / KIB), + bytes => write!(f, "{bytes} B"), + } + } +} + #[derive(Debug, Clone, Deserialize)] pub struct PortPolicyFetchConfig { /// Timeout for a single `Info()` RPC attempt. @@ -172,6 +447,13 @@ pub struct Timeouts { pub data_timeout_enabled: bool, #[serde(with = "serde_duration")] pub idle: Duration, + /// No longer read. The per-operation write timer was replaced by the + /// connection-level progress watchdog in `io_bridge`, which catches a + /// stalled write through `idle` instead: a write that makes no progress + /// stops bumping the direction's progress counter, and the watchdog fires. + /// The key is still accepted so existing configs -- and the CVM app + /// entrypoint's `TIMEOUT_WRITE` -- keep parsing. + #[allow(dead_code)] #[serde(with = "serde_duration")] pub write: Duration, #[serde(with = "serde_duration")] @@ -431,4 +713,71 @@ mod tests { "Client IP range is not in the network" ); } + + fn engage_after(toml: &str) -> EngageAfter { + Figment::from(Toml::string(toml)) + .extract() + .expect("valid EngageAfter") + } + + #[test] + fn no_gate_engages_immediately() { + let gate = engage_after(""); + assert!(gate.is_immediate()); + assert!(gate.after_bytes.is_none()); + assert!(gate.after_duration.is_none()); + } + + #[test] + fn a_configured_gate_is_not_immediate() { + assert!(!engage_after("after_bytes = 65536").is_immediate()); + assert!(!engage_after("after_duration = \"5s\"").is_immediate()); + } + + #[test] + fn byte_gate_ignores_elapsed_time() { + let gate = engage_after("after_bytes = 1024"); + let long_ago = Instant::now() - Duration::from_secs(3600); + assert!(!gate.reached(1023, long_ago)); + assert!(gate.reached(1024, long_ago)); + } + + #[test] + fn duration_gate_ignores_bytes() { + let gate = engage_after("after_duration = \"5s\""); + assert!(!gate.reached(u64::MAX, Instant::now())); + assert!(gate.reached(0, Instant::now() - Duration::from_secs(5))); + } + + #[test] + fn gates_are_independent_and_either_fires() { + // The case the byte gate alone cannot express: a low-rate stream that + // stays well under `after_bytes` but outlives `after_duration`. + let gate = engage_after("after_bytes = 65536\nafter_duration = \"5s\""); + let just_started = Instant::now(); + assert!(!gate.reached(64, just_started)); + assert!(gate.reached(65536, just_started)); + assert!(gate.reached(64, Instant::now() - Duration::from_secs(5))); + } + + #[test] + fn splice_section_is_optional() { + #[derive(Deserialize)] + struct Holder { + #[serde(default)] + tcp_splice: Option, + } + let absent: Holder = Figment::from(Toml::string("")).extract().unwrap(); + assert!( + absent.tcp_splice.is_none(), + "absent section disables splice" + ); + + let present: Holder = Figment::from(Toml::string("[tcp_splice]\nafter_duration = \"5s\"")) + .extract() + .unwrap(); + let gate = present.tcp_splice.expect("section present"); + assert_eq!(gate.after_duration, Some(Duration::from_secs(5))); + assert!(gate.after_bytes.is_none()); + } } diff --git a/dstack/gateway/src/main.rs b/dstack/gateway/src/main.rs index f689d2d7f..eb0eecec9 100644 --- a/dstack/gateway/src/main.rs +++ b/dstack/gateway/src/main.rs @@ -240,11 +240,14 @@ async fn main() -> Result<()> { let args = Args::parse(); let figment = config::load_config_figment(args.config.as_deref()); - let config = figment.focus("core").extract::()?; + let mut config = figment.focus("core").extract::()?; // Validate node_id if config.sync.enabled && config.sync.node_id == 0 { anyhow::bail!("node_id must be greater than 0"); } + // Before anything reads `proxy.ktls`: the acceptor built later decides + // whether to extract session secrets from it. + proxy::disable_ktls_if_unsupported(&mut config.proxy); config::setup_wireguard(&config.wg)?; diff --git a/dstack/gateway/src/models.rs b/dstack/gateway/src/models.rs index 77d2cfe2e..318869334 100644 --- a/dstack/gateway/src/models.rs +++ b/dstack/gateway/src/models.rs @@ -2,7 +2,7 @@ // // SPDX-License-Identifier: Apache-2.0 -use dstack_gateway_rpc::{AcmeInfoResponse, StatusResponse}; +use dstack_gateway_rpc::{AcmeInfoResponse, ProxyAccelStatus, StatusResponse}; use rinja::Template; use serde::{Deserialize, Serialize}; use std::{ @@ -168,4 +168,7 @@ pub struct WgConf<'a> { pub struct Dashboard { pub status: StatusResponse, pub acme_info: AcmeInfoResponse, + /// Lifted out of `status` so the template does not have to unwrap the + /// proto's optional message on every field. + pub accel: ProxyAccelStatus, } diff --git a/dstack/gateway/src/proxy.rs b/dstack/gateway/src/proxy.rs index 4451ee359..58d73d7d9 100644 --- a/dstack/gateway/src/proxy.rs +++ b/dstack/gateway/src/proxy.rs @@ -22,7 +22,7 @@ use tokio::{ runtime::Runtime, time::timeout, }; -use tracing::{debug, error, info, info_span, Instrument}; +use tracing::{debug, debug_span, error, info, warn, Instrument}; use crate::{ config::ProxyConfig, @@ -42,9 +42,15 @@ pub(crate) struct AddressInfo { pub(crate) type AddressGroup = smallvec::SmallVec<[AddressInfo; 4]>; +mod adaptive_ktls; +mod balance; +mod idle; mod io_bridge; pub(crate) mod port_policy; +mod reuseport; mod sni; +mod splice; +pub(crate) mod stats; mod tls_passthough; mod tls_terminate; @@ -141,7 +147,7 @@ async fn handle_connection(inbound: TcpStream, state: Proxy) -> Result<()> { .await .context("proxy protocol header timeout")? .context("failed to read proxy protocol header")?; - info!("client address: {}", DisplayAddr(&pp_header)); + debug!("client address: {}", DisplayAddr(&pp_header)); let (sni, buffer) = timeout(timeouts.handshake, take_sni(&mut inbound)) .await @@ -168,64 +174,172 @@ async fn handle_connection(inbound: TcpStream, state: Proxy) -> Result<()> { } } -#[inline(never)] -pub async fn proxy_main(rt: &Runtime, config: &ProxyConfig, proxy: Proxy) -> Result<()> { +/// Bind one listener per configured port. +/// +/// With `reuse_port` every worker binds its own listener on the same port and +/// the kernel spreads incoming connections across them, so each worker can +/// accept and serve its connections without any cross-thread handoff. +/// Accept queue depth. tokio's default is 1024, which is where SYN drops start +/// under connection bursts; both listen paths use this so they behave alike. +const LISTEN_BACKLOG: i32 = 4096; + +async fn bind_listeners(config: &ProxyConfig, reuse_port: bool) -> Result> { let mut tcp_listeners = Vec::new(); for &port in &config.listen_port { - let listener = TcpListener::bind((config.listen_addr, port)) - .await - .with_context(|| format!("failed to bind {}:{}", config.listen_addr, port))?; + let listener = { + let addr = std::net::SocketAddr::from((config.listen_addr, port)); + let socket = socket2::Socket::new( + socket2::Domain::IPV4, + socket2::Type::STREAM, + Some(socket2::Protocol::TCP), + ) + .context("failed to create listening socket")?; + if reuse_port { + socket + .set_reuse_port(true) + .context("failed to set SO_REUSEPORT")?; + } + socket.set_reuse_address(true).ok(); + socket.set_nonblocking(true).ok(); + socket + .bind(&addr.into()) + .with_context(|| format!("failed to bind {addr}"))?; + socket.listen(LISTEN_BACKLOG).context("failed to listen")?; + TcpListener::from_std(std::net::TcpListener::from(socket)) + .context("failed to register listener with tokio")? + }; info!("tcp bridge listening on {}:{}", config.listen_addr, port); tcp_listeners.push(listener); } + Ok(tcp_listeners) +} + +#[inline(never)] +pub async fn proxy_main(rt: &Runtime, config: &ProxyConfig, proxy: Proxy) -> Result<()> { + let tcp_listeners = bind_listeners(config, false).await?; + accept_loop(tcp_listeners, proxy, Some(rt), None).await +} + +/// The per-connection task: everything a single proxied connection does. +/// Was this failure just the peer hanging up? +/// +/// Clients disconnecting mid-connection is routine -- a browser navigating away, +/// a mobile network dropping, a load generator ending its run. Logging those at +/// error level buries the failures that are actually the gateway's fault: a 40 +/// minute soak produced 379 such lines and no real errors. +fn is_peer_disconnect(err: &anyhow::Error) -> bool { + err.chain().any(|cause| { + cause.downcast_ref::().is_some_and(|io| { + matches!( + io.kind(), + std::io::ErrorKind::ConnectionReset + | std::io::ErrorKind::BrokenPipe + | std::io::ErrorKind::UnexpectedEof + | std::io::ErrorKind::ConnectionAborted + ) + }) + }) +} + +fn conn_task( + inbound: TcpStream, + from: std::net::SocketAddr, + proxy: Proxy, + slot: Option, +) -> impl std::future::Future + Send + 'static { + let span = debug_span!("conn", id = next_connection_id()); + let conn_entered = EnteredCounter::new(&NUM_CONNECTIONS); + async move { + let _conn_entered = conn_entered; + let _slot = slot; + debug!(%from, "new connection"); + let timeouts = &proxy.config.proxy.timeouts; + match timeout(timeouts.total, handle_connection(inbound, proxy)).await { + Ok(Ok(_)) => debug!("connection closed"), + Ok(Err(e)) if is_peer_disconnect(&e) => debug!("peer disconnected: {e:#}"), + Ok(Err(e)) => error!("connection error: {e:#}"), + Err(_) => error!("connection kept too long, force closing"), + } + } + .instrument(span) +} + +/// Accept connections forever. +/// +/// `rt` selects where connections run: `Some(worker_rt)` hands them to a shared +/// multi-threaded runtime, `None` keeps them on the calling thread's own +/// runtime (thread-per-core), which avoids the cross-thread handoff and the +/// work-stealing migrations that come with it. +async fn accept_loop( + tcp_listeners: Vec, + proxy: Proxy, + rt: Option<&Runtime>, + mut balance: Option<( + balance::Balancer, + tokio::sync::mpsc::Receiver, + )>, +) -> Result<()> { if tcp_listeners.is_empty() { bail!("no tcp listen ports configured"); } - let poll_counter = AtomicUsize::new(0); loop { // Accept from any TCP listener via round-robin poll. let poll_start = poll_counter.fetch_add(1, Ordering::Relaxed); let n = tcp_listeners.len(); - let accepted: std::io::Result<(TcpStream, std::net::SocketAddr)> = - std::future::poll_fn(|cx| { - for j in 0..n { - let i = (poll_start + j) % n; - if let Poll::Ready(result) = tcp_listeners[i].poll_accept(cx) { - return Poll::Ready(result); + let accept_next = std::future::poll_fn(|cx| { + for j in 0..n { + let i = (poll_start + j) % n; + if let Poll::Ready(result) = tcp_listeners[i].poll_accept(cx) { + return Poll::Ready(result); + } + } + Poll::Pending + }); + // Also take connections other cores decided to give us. + let accepted: std::io::Result<(TcpStream, std::net::SocketAddr)> = match balance.as_mut() { + Some((_, rx)) => { + tokio::select! { + r = accept_next => r, + Some((raw, from, slot)) = rx.recv() => { + // Register the handed-over socket with *this* core's reactor. + match TcpStream::from_std(raw) { + Ok(stream) => { + let task = conn_task(stream, from, proxy.clone(), Some(slot)); + tokio::spawn(task); + } + Err(e) => error!("failed to adopt handed-over connection: {e}"), + } + continue; } } - Poll::Pending - }) - .await; + } + None => accept_next.await, + }; match accepted { Ok((inbound, from)) => { - let span = info_span!("conn", id = next_connection_id()); - let _enter = span.enter(); - let conn_entered = EnteredCounter::new(&NUM_CONNECTIONS); - - info!(%from, "new connection"); - let proxy = proxy.clone(); - rt.spawn( - async move { - let _conn_entered = conn_entered; - let timeouts = &proxy.config.proxy.timeouts; - let result = - timeout(timeouts.total, handle_connection(inbound, proxy)).await; - match result { - Ok(Ok(_)) => { - info!("connection closed"); - } - Ok(Err(e)) => { - error!("connection error: {e:#}"); - } - Err(_) => { - error!("connection kept too long, force closing"); - } + // Disable Nagle: this is a latency-sensitive proxy and small + // request/response traffic otherwise stalls on delayed ACKs. + let _ = inbound.set_nodelay(true); + // In thread-per-core mode, hand the connection to a less loaded + // core when this one is running ahead; SO_REUSEPORT's hash can + // leave a core starved and it cannot be fixed later. + let placed: Option<(TcpStream, Option)> = match balance.as_ref() + { + Some((b, _)) => b.place(inbound, from).map(|(s, slot)| (s, Some(slot))), + None => Some((inbound, None)), + }; + if let Some((inbound, slot)) = placed { + let task = conn_task(inbound, from, proxy.clone(), slot); + match rt { + Some(rt) => { + rt.spawn(task); + } + None => { + tokio::spawn(task); } } - .in_current_span(), - ); + } } Err(e) => { error!("failed to accept connection: {e:?}"); @@ -240,6 +354,26 @@ fn next_connection_id() -> usize { } pub fn start(config: ProxyConfig, app_state: Proxy) -> Result<()> { + if config.thread_per_core { + // Probe SO_REUSEPORT before committing: it is the one prerequisite the + // thread-per-core model cannot work without, and failing to serve at all + // is far worse than losing the optimisation. + match probe_reuse_port(&config) { + Ok(()) => return start_thread_per_core(config, app_state), + Err(err) => warn!( + "thread_per_core requested but SO_REUSEPORT is unavailable ({err:#}); \ + falling back to the shared-runtime proxy" + ), + } + } + if config.connection_rebalance { + // It is only wired up by the thread-per-core path: there is nothing to + // rebalance between when every connection lands on a shared runtime. + warn!( + "connection_rebalance is set but has no effect without thread_per_core; \ + the shared-runtime proxy balances connections through its scheduler" + ); + } std::thread::Builder::new() .name("proxy-main".to_string()) .spawn(move || { @@ -268,10 +402,216 @@ pub fn start(config: ProxyConfig, app_state: Proxy) -> Result<()> { Ok(()) } +/// Check that a `SO_REUSEPORT` listener can actually be created and bound. +fn probe_reuse_port(config: &ProxyConfig) -> Result<()> { + let port = *config + .listen_port + .first() + .context("no tcp listen ports configured")?; + let addr = std::net::SocketAddr::from((config.listen_addr, port)); + let socket = socket2::Socket::new( + socket2::Domain::IPV4, + socket2::Type::STREAM, + Some(socket2::Protocol::TCP), + ) + .context("failed to create probe socket")?; + socket + .set_reuse_port(true) + .context("SO_REUSEPORT not supported")?; + socket.set_reuse_address(true).ok(); + socket + .bind(&addr.into()) + .with_context(|| format!("failed to bind {addr} with SO_REUSEPORT"))?; + Ok(()) +} + +/// Turn kTLS off when the kernel cannot provide it. +/// +/// Without the TLS ULP an immediate offload fails every handshake, which is at +/// least loud. A gated offload is worse: the connection is served from +/// userspace up to the threshold and only then fails, so the client gets a +/// successful response truncated at exactly the gate. Both are worse than never +/// offloading, and whether the ULP is there is not something the config can +/// know -- so ask the kernel once, at startup, before the acceptor is built +/// (that is also what decides whether rustls extracts session secrets at all). +/// +/// This only establishes that the ULP exists. A cipher suite the kernel does +/// not implement still fails per connection, at offload time. +pub fn disable_ktls_if_unsupported(config: &mut ProxyConfig) { + if config.ktls.is_none() { + return; + } + if let Err(err) = probe_ktls() { + warn!( + "kTLS is configured but unavailable ({err:#}); \ + falling back to userspace TLS record encryption" + ); + config.ktls = None; + stats::mark_ktls_unsupported(); + } +} + +/// Check that the kernel exposes the TLS upper-layer protocol. +/// +/// `TCP_ULP` is only accepted on an established socket, so this sets up a +/// throwaway loopback connection instead of probing a fresh one, which would +/// fail with `ENOTCONN` whether or not the ULP exists. +#[cfg(target_os = "linux")] +fn probe_ktls() -> Result<()> { + use std::net::{Ipv4Addr, TcpListener, TcpStream}; + use std::os::fd::AsRawFd; + + /// `include/uapi/linux/tcp.h`; not exposed by the `libc` crate. + const TCP_ULP: libc::c_int = 31; + + let listener = + TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).context("failed to bind the probe listener")?; + let addr = listener + .local_addr() + .context("failed to read the probe listener address")?; + let client = TcpStream::connect(addr).context("failed to connect the probe socket")?; + let _server = listener + .accept() + .context("failed to accept the probe socket")?; + + let name = c"tls"; + // SAFETY: `name` outlives the call and `len` matches it, NUL included. + let rc = unsafe { + libc::setsockopt( + client.as_raw_fd(), + libc::IPPROTO_TCP, + TCP_ULP, + name.as_ptr().cast(), + name.to_bytes_with_nul().len() as libc::socklen_t, + ) + }; + if rc != 0 { + return Err(std::io::Error::last_os_error()) + .context("kernel rejected the TLS ULP (is CONFIG_TLS enabled?)"); + } + Ok(()) +} + +#[cfg(not(target_os = "linux"))] +fn probe_ktls() -> Result<()> { + bail!("kernel TLS is only available on Linux") +} + +/// Thread-per-core proxy: `workers` threads, each with its own single-threaded +/// runtime and its own `SO_REUSEPORT` listener. +/// +/// The kernel load-balances new connections across the listeners, and a +/// connection is accepted and served entirely on one thread. That removes the +/// accept-thread -> worker handoff and the work-stealing scheduler's task +/// migrations, which together accounted for ~0.6 context switches per request +/// in the default model. +fn start_thread_per_core(config: ProxyConfig, app_state: Proxy) -> Result<()> { + let workers = config.workers.max(1); + // Bind every listener here, in order, rather than letting each thread bind + // its own. The CBPF steering program that originally needed a known group + // order is gone (see `super::reuseport`), but binding in one place is kept: + // it keeps listener-to-core assignment deterministic across restarts, and it + // fails startup as a whole if any bind fails, instead of leaving some cores + // serving and others dead. + let mut per_worker: Vec> = + (0..workers).map(|_| Vec::new()).collect(); + for &port in &config.listen_port { + let addr = std::net::SocketAddr::from((config.listen_addr, port)); + let group = reuseport::bind_group(addr, workers, LISTEN_BACKLOG) + .with_context(|| format!("failed to bind reuseport group on {addr}"))?; + info!( + "tcp bridge listening on {}:{} across {} listeners", + config.listen_addr, port, workers + ); + for (i, l) in group.into_iter().enumerate() { + per_worker[i].push(l); + } + } + + // Per-core balancers share the connection counts; each core also gets the + // receiving end of its handoff channel. + let (balancers, receivers) = if config.connection_rebalance { + let (b, r) = balance::Balancer::build(workers); + ( + b.into_iter().map(Some).collect(), + r.into_iter().map(Some).collect(), + ) + } else { + ( + (0..workers).map(|_| None).collect::>(), + (0..workers).map(|_| None).collect::>(), + ) + }; + + let config = Arc::new(config); + for (i, ((std_listeners, balancer), receiver)) in per_worker + .into_iter() + .zip(balancers) + .zip(receivers) + .enumerate() + { + let config = config.clone(); + let app_state = app_state.clone(); + std::thread::Builder::new() + .name(format!("proxy-core-{i}")) + .spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .or_panic("Failed to build Tokio runtime"); + let result = rt.block_on(async { + let mut listeners = Vec::with_capacity(std_listeners.len()); + for l in std_listeners { + listeners.push( + TcpListener::from_std(l) + .context("failed to register listener with tokio")?, + ); + } + let bal = balancer.zip(receiver); + accept_loop(listeners, app_state, None, bal).await + }); + if let Err(err) = result { + error!( + "proxy core {i} error on {}:{:?}: {err:?}", + config.listen_addr, config.listen_port + ); + } + }) + .with_context(|| format!("Failed to spawn proxy-core-{i} thread"))?; + } + Ok(()) +} + #[cfg(test)] mod tests { use super::*; + fn default_proxy_config() -> ProxyConfig { + crate::config::load_config_figment(None) + .focus("core.proxy") + .extract() + .expect("the shipped default config should parse") + } + + #[test] + fn ktls_probe_leaves_an_unconfigured_gateway_alone() { + let mut config = default_proxy_config(); + config.ktls = None; + disable_ktls_if_unsupported(&mut config); + assert!(config.ktls.is_none()); + } + + #[test] + fn ktls_survives_the_probe_only_when_the_kernel_supports_it() { + let mut config = default_proxy_config(); + config.ktls = Some(crate::config::EngageAfter::default()); + disable_ktls_if_unsupported(&mut config); + // Whichever way this kernel answers, the config must agree with it: + // keeping kTLS on a kernel without the ULP is what truncates responses + // at the gate. + assert_eq!(config.ktls.is_some(), probe_ktls().is_ok()); + } + #[test] fn test_parse_destination() { // Test basic app_id only diff --git a/dstack/gateway/src/proxy/adaptive_ktls.rs b/dstack/gateway/src/proxy/adaptive_ktls.rs new file mode 100644 index 000000000..c7b951585 --- /dev/null +++ b/dstack/gateway/src/proxy/adaptive_ktls.rs @@ -0,0 +1,294 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Traffic-triggered kernel TLS offload. +//! +//! Measurements show the two halves of kTLS pull in opposite directions: +//! enabling it costs ~30% of connection setup rate (secret extraction plus +//! kernel ULP setup per connection) but wins ~25% on bulk throughput once +//! combined with splice. Short request/response connections therefore pay the +//! setup cost and never earn it back. +//! +//! This module keeps the connection in userspace rustls after the handshake and +//! only hands it to the kernel once it has proven itself: once the configured +//! [`EngageAfter`] gate fires, the stream is drained at a TLS record boundary +//! and switched to kTLS + splice for the remainder. +//! +//! Handing over mid-stream is sound because the secrets rustls exports carry +//! the current record sequence numbers, and `CorkStream` exists precisely to +//! stop reads at a record boundary so nothing is left half-parsed. + +use std::time::{Duration, Instant}; + +use anyhow::{bail, Context, Result}; +use ktls::CorkStream; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::TcpStream; +use tokio_rustls::server::TlsStream; +use tracing::debug; + +use super::idle::IdleWatchdog; +use super::splice::{splice_bidirectional, CloseKind}; +use super::tls_terminate::SocketParts; +use crate::config::{EngageAfter, SpliceConfig}; + +/// Why the userspace relay phase stopped. +enum Phase { + /// The connection proved itself worth the offload. + Gated, + /// One side closed before the gate fired. + Eof, +} + +/// Relay both directions in userspace until either side closes or `gate` fires. +async fn relay_until( + tls: &mut S, + upstream: &mut TcpStream, + gate: &EngageAfter, + idle: Option, +) -> Result +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let (mut tr, mut tw) = tokio::io::split(tls); + let (mut ur, mut uw) = upstream.split(); + let mut down = vec![0u8; 32 * 1024]; + let mut up = vec![0u8; 32 * 1024]; + let mut moved: u64 = 0; + let start = Instant::now(); + let mut watchdog: Option = match idle { + Some(idle) => { + let mut w = IdleWatchdog::new(idle); + w.tick().await; // the first tick completes immediately + Some(w) + } + None => None, + }; + + let phase = loop { + // One side closing is not the end of the connection: a client that ends + // its request with close_notify still expects the response. Propagate + // the EOF to that direction's peer, then drain the other direction + // before giving up on the connection. + macro_rules! finish_one { + ($closing:expr, $r:expr, $w:expr, $buf:expr) => {{ + $closing.shutdown().await.ok(); + loop { + let n = tokio::select! { + () = async { match watchdog.as_mut() { + Some(w) => w.tick().await, + None => std::future::pending().await, + } } => { + // The drain is watched too: a backend that accepts the + // request and then never answers would otherwise hold + // the connection until `timeouts.total`. + if watchdog.as_mut().is_some_and(|w| w.stalled(moved)) { + bail!("idle timeout"); + } + continue; + } + r = $r.read(&mut $buf) => r.context("read error")?, + }; + if n == 0 { + break; + } + // Not `write_all`: it is not cancel-safe, so it cannot sit + // in a `select!`, and leaving it outside meant a client + // that stopped reading blocked the drain in its write with + // the watchdog unpolled -- the mirror of the silent-backend + // stall, and just as good for holding a connection to + // `timeouts.total`. Single `write` calls are cancel-safe + // (nothing is written when the other branch wins), so the + // partial-write loop is ours to drive. + let mut written = 0usize; + while written < n { + let count = tokio::select! { + () = async { match watchdog.as_mut() { + Some(w) => w.tick().await, + None => std::future::pending().await, + } } => { + if watchdog.as_mut().is_some_and(|w| w.stalled(moved)) { + bail!("idle timeout"); + } + continue; + } + r = $w.write(&$buf[written..n]) => r.context("write error")?, + }; + if count == 0 { + bail!("write accepted no bytes"); + } + written += count; + // Per partial write, so a peer draining slowly still + // counts as progress and is not reaped for being slow. + moved += count as u64; + } + } + $w.shutdown().await.ok(); + break Phase::Eof; + }}; + } + tokio::select! { + () = async { match watchdog.as_mut() { + Some(w) => w.tick().await, + // No idle timeout configured: this arm must never win. + None => std::future::pending().await, + } } => { + if watchdog.as_mut().is_some_and(|w| w.stalled(moved)) { + bail!("idle timeout"); + } + continue; + } + r = tr.read(&mut down) => { + let n = r.context("read from client failed")?; + if n == 0 { finish_one!(uw, ur, tw, up); } + uw.write_all(&down[..n]).await.context("write to app failed")?; + moved += n as u64; + } + r = ur.read(&mut up) => { + let n = r.context("read from app failed")?; + if n == 0 { finish_one!(tw, tr, uw, down); } + tw.write_all(&up[..n]).await.context("write to client failed")?; + moved += n as u64; + } + } + if gate.reached(moved, start) { + // Flush before handing the socket to the kernel so no plaintext is + // still sitting in a rustls write buffer. + tw.flush().await.context("flush before offload failed")?; + break Phase::Gated; + } + }; + Ok(phase) +} + +/// Relay a freshly accepted TLS connection, upgrading it to kTLS + splice once +/// the kTLS gate fires. `splice` supplies the relay settings used afterwards. +pub(crate) async fn relay_with_adaptive_offload( + mut tls: TlsStream>, + mut upstream: TcpStream, + ktls: &EngageAfter, + splice: &SpliceConfig, + idle: Option, +) -> Result<()> +where + IO: AsyncRead + AsyncWrite + Unpin + std::os::fd::AsRawFd + ktls::AsyncReadReady, + IO: SocketParts, +{ + match relay_until(&mut tls, &mut upstream, ktls, idle).await? { + Phase::Eof => return Ok(()), + Phase::Gated => {} + } + debug!("offloading connection to kTLS after {ktls:?}"); + + // config_ktls_server corks the stream, drains rustls to a record boundary + // and installs the current traffic secrets into the kernel. + let ktls_stream = super::stats::record_ktls_offload(ktls::config_ktls_server(tls).await) + .context("failed to switch connection to kernel TLS")?; + let (drained, io) = ktls_stream.into_raw(); + if let Some(drained) = drained { + if !drained.is_empty() { + upstream + .write_all(&drained) + .await + .context("failed to flush drained data to app")?; + } + } + // Same rule as the immediate offload path: the sniff remainder is raw + // ciphertext, so it can neither be forwarded to the app nor pushed back + // into a socket the kernel now owns. A completed handshake always consumes + // it, so this is unreachable -- refuse rather than corrupt a stream if it + // ever is not. + let (buffered, tcp) = io.into_socket_parts(); + if !buffered.is_empty() { + bail!( + "{} bytes of unconsumed ciphertext at kTLS handover", + buffered.len() + ); + } + // The client side is now a kTLS socket, so its close needs a close_notify. + splice_bidirectional( + tcp, + upstream, + splice.release_idle_pipes, + idle, + CloseKind::KernelTls, + ) + .await + .context("splice after kTLS offload failed") +} + +#[cfg(test)] +mod tests { + use super::*; + use tokio::net::TcpListener; + + async fn connected_pair() -> (TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let client = TcpStream::connect(addr).await.unwrap(); + let (server, _) = listener.accept().await.unwrap(); + (client, server) + } + + /// A gate no test connection reaches, so the relay stays in the userspace + /// phase for the whole exchange. `relay_until` is generic over the client + /// stream, so a plain socket stands in for the TLS one and the half-close + /// handling is exercised without a handshake. + fn ungated() -> EngageAfter { + EngageAfter { + after_bytes: Some(1 << 30), + after_duration: None, + } + } + + #[tokio::test] + async fn response_survives_a_client_half_close_before_the_gate() { + let (mut client, mut tls_side) = connected_pair().await; + let (mut upstream, mut backend) = connected_pair().await; + let relay = tokio::spawn(async move { + relay_until(&mut tls_side, &mut upstream, &ungated(), None).await + }); + + client.write_all(b"ping").await.unwrap(); + client.shutdown().await.unwrap(); + + let mut req = vec![0u8; 4]; + backend.read_exact(&mut req).await.unwrap(); + assert_eq!(&req, b"ping"); + let mut trailing = Vec::new(); + backend.read_to_end(&mut trailing).await.unwrap(); + assert!(trailing.is_empty()); + backend.write_all(b"pong").await.unwrap(); + drop(backend); + + let mut resp = Vec::new(); + client.read_to_end(&mut resp).await.unwrap(); + assert_eq!(resp, b"pong", "client lost the response after half-closing"); + assert!(matches!(relay.await.unwrap().unwrap(), Phase::Eof)); + } + + #[tokio::test] + async fn request_survives_an_app_half_close_before_the_gate() { + let (mut client, mut tls_side) = connected_pair().await; + let (mut upstream, mut backend) = connected_pair().await; + let relay = tokio::spawn(async move { + relay_until(&mut tls_side, &mut upstream, &ungated(), None).await + }); + + backend.write_all(b"early").await.unwrap(); + backend.shutdown().await.unwrap(); + + let mut resp = vec![0u8; 5]; + client.read_exact(&mut resp).await.unwrap(); + assert_eq!(&resp, b"early"); + + client.write_all(b"late").await.unwrap(); + drop(client); + + let mut got = Vec::new(); + backend.read_to_end(&mut got).await.unwrap(); + assert_eq!(got, b"late", "app lost the request after half-closing"); + assert!(matches!(relay.await.unwrap().unwrap(), Phase::Eof)); + } +} diff --git a/dstack/gateway/src/proxy/balance.rs b/dstack/gateway/src/proxy/balance.rs new file mode 100644 index 000000000..a065c4ae9 --- /dev/null +++ b/dstack/gateway/src/proxy/balance.rs @@ -0,0 +1,205 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Evening out connections across the thread-per-core workers. +//! +//! `SO_REUSEPORT` assigns a connection to a listener by hashing its 4-tuple, and +//! thread-per-core cannot move a connection afterwards, so a core that draws few +//! connections idles while its neighbours saturate. Measured at 16 connections +//! over 4 cores, utilisation came out like `[99, 42, 101, 101]` and throughput +//! varied 46% between restarts on nothing but how the hash fell. +//! +//! Steering the kernel's choice was tried first and failed (see +//! [`super::reuseport`]). This instead lets the accepting core hand the +//! connection to a less loaded one. The cost is a channel send per *rebalanced +//! connection* -- not per connection, and never per request -- so the +//! thread-per-core property that made this model fast in the first place is kept +//! for everything already in flight. +//! +//! This is the same shape as HAProxy's multi-queue accept, which picks the least +//! loaded of three candidate threads and pushes the connection onto its ring. +//! Two of its choices were tried here and did not transfer: capping accepts per +//! wakeup (its `maxaccept`) and routing every connection through the channel even +//! when it stays local both measured within noise. Its shared listening socket +//! measured clearly worse for us -- 227k against 245k at 16 connections -- since +//! every worker's reactor then wakes on every connection. Its own numbers say the +//! same thing from the other side: put HAProxy on per-thread reuseport listeners +//! (`shards by-thread`) and it drops 15%, below our thread-per-core. + +use std::net::SocketAddr; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use tokio::net::TcpStream; +use tokio::sync::mpsc; +use tracing::{debug, warn}; + +/// A connection accounted to one core, released when the connection ends. +pub(crate) struct CoreSlot { + counts: Arc>, + core: usize, +} + +impl CoreSlot { + fn claim(counts: Arc>, core: usize) -> Self { + counts[core].fetch_add(1, Ordering::Relaxed); + Self { counts, core } + } +} + +impl Drop for CoreSlot { + fn drop(&mut self) { + self.counts[self.core].fetch_sub(1, Ordering::Relaxed); + } +} + +/// A connection handed over from another core. +/// +/// Carried as a `std::net::TcpStream`, deliberately: a tokio `TcpStream` stays +/// registered with the reactor that accepted it, so moving the object to another +/// core would leave its readiness handling behind and cost a cross-core hop on +/// every read and write for the life of the connection. Handing over the plain +/// socket lets the target core register it with its own reactor. +pub(crate) type Handoff = (std::net::TcpStream, SocketAddr, CoreSlot); + +/// Per-core view of the shared connection counts and handoff channels. +pub(crate) struct Balancer { + counts: Arc>, + senders: Arc>>, + me: usize, +} + +/// Handoffs a core may have waiting before others stop offering it work. +/// +/// The queue is bounded on purpose. A core only receives connections while it +/// is the *least* loaded, so in steady state this never fills. What it protects +/// against is a core that stops draining -- wedged or gone: with an unbounded +/// queue the others keep succeeding at `send`, and every connection they hand +/// over sits unserved until `timeouts.total` while its `CoreSlot` keeps the +/// core looking busier, so the queue and the memory behind it only grow. Full +/// means "this core is not actually taking work", and the sender keeps the +/// connection instead. +const HANDOFF_QUEUE: usize = 1024; + +/// How far above the least loaded core this one has to be before handing a +/// connection over. +/// +/// Proportional, not fixed: a slack of 2 stops a core being starved at 4 +/// connections per core but keeps churning at 12, and every migration costs a +/// little locality. Measured against a fixed slack of 2 on passthrough +/// small-request: +6.5% at 16 connections, and the worst-loaded core goes from +/// 81% to 97% busy. +fn migration_threshold(least: usize) -> usize { + least + 1 + least / 4 +} + +impl Balancer { + /// Build one balancer per core, plus the receiver each core listens on. + pub(crate) fn build(workers: usize) -> (Vec, Vec>) { + let counts = Arc::new( + (0..workers) + .map(|_| AtomicUsize::new(0)) + .collect::>(), + ); + let mut senders = Vec::with_capacity(workers); + let mut receivers = Vec::with_capacity(workers); + for _ in 0..workers { + let (tx, rx) = mpsc::channel(HANDOFF_QUEUE); + senders.push(tx); + receivers.push(rx); + } + let senders = Arc::new(senders); + let balancers = (0..workers) + .map(|me| Self { + counts: counts.clone(), + senders: senders.clone(), + me, + }) + .collect(); + (balancers, receivers) + } + + /// Which core should own a freshly accepted connection. + fn target(&self) -> usize { + if self.counts.len() < 2 { + return self.me; + } + let mine = self.counts[self.me].load(Ordering::Relaxed); + let mut best = self.me; + let mut best_val = mine; + for (i, c) in self.counts.iter().enumerate() { + let v = c.load(Ordering::Relaxed); + if v < best_val { + best = i; + best_val = v; + } + } + if mine >= migration_threshold(best_val) { + best + } else { + self.me + } + } + + /// Account a newly accepted connection, handing it to a less loaded core if + /// this one is running ahead. + /// + /// Returns the slot to keep alongside the connection when it stays here, or + /// `None` once the connection has been handed away. + /// + /// Handing over is best-effort: every failure path falls back to serving the + /// connection on this core, because a connection served by a busier core is + /// strictly better than one that is dropped. The one exception is + /// [`TcpStream::into_std`], which consumes the stream and closes the socket + /// when it fails -- there is nothing left to fall back to, so that case is + /// logged rather than silently counted as a handover. + pub(crate) fn place( + &self, + stream: TcpStream, + from: SocketAddr, + ) -> Option<(TcpStream, CoreSlot)> { + let target = self.target(); + if target == self.me { + let slot = CoreSlot::claim(self.counts.clone(), self.me); + return Some((stream, slot)); + } + // Drop this core's reactor registration before handing the socket over. + let raw = match stream.into_std() { + Ok(raw) => raw, + Err(err) => { + // `into_std` took ownership, so the socket is already closed. + warn!("dropping connection from {from}: failed to deregister for handover: {err}"); + return None; + } + }; + let slot = CoreSlot::claim(self.counts.clone(), target); + // `try_send`, not `send`: this runs on the accept path and must not wait + // on a core that is not draining. See `HANDOFF_QUEUE`. + let raw = match self.senders[target].try_send((raw, from, slot)) { + Ok(()) => return None, + Err(mpsc::error::TrySendError::Full((raw, _, _))) => { + debug!( + "core {target} handoff queue is full; keeping connection on core {}", + self.me + ); + raw + } + Err(mpsc::error::TrySendError::Closed((raw, _, _))) => { + debug!( + "core {target} is gone; keeping connection on core {}", + self.me + ); + raw + } + }; + match TcpStream::from_std(raw) { + Ok(stream) => Some((stream, CoreSlot::claim(self.counts.clone(), self.me))), + Err(err) => { + // Same as above: `from_std` consumed the socket. + warn!("dropping connection from {from}: failed to re-adopt after handover: {err}"); + None + } + } + } +} diff --git a/dstack/gateway/src/proxy/idle.rs b/dstack/gateway/src/proxy/idle.rs new file mode 100644 index 000000000..34ce04f64 --- /dev/null +++ b/dstack/gateway/src/proxy/idle.rs @@ -0,0 +1,133 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! One idle watchdog for every relay path. +//! +//! `timeouts.idle` used to be enforced per read, which meant a connection died +//! when *one* direction went quiet even while the other was busy. It is now a +//! connection-level watchdog that samples a monotonic progress counter: if +//! neither direction has moved for the idle window, the connection is stalled. +//! That costs one timer per window rather than one per operation, which is why +//! the fast paths can afford it. +//! +//! It lives here rather than inside a relay because every relay needs it and +//! they are shaped differently: the buffered bridge and the pre-gate relays are +//! `select!` loops that can poll it as one more branch, while a spliced +//! connection has no loop to hang it on and races it against the transfer +//! instead. Both call the same sampling logic, so the three paths cannot drift +//! into enforcing different things -- which is exactly what happened when +//! splice and kTLS silently had no idle timeout at all. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use tokio::time::{interval, Interval, MissedTickBehavior}; + +/// Samples a progress counter a few times per idle window. +pub(crate) struct IdleWatchdog { + ticker: Interval, + /// Consecutive ticks that saw no progress. + idle_ticks: u32, + /// Ticks without progress that add up to the configured idle window. + max_idle_ticks: u32, + last_seen: u64, +} + +impl IdleWatchdog { + pub(crate) fn new(idle: Duration) -> Self { + // Four samples per window bounds the overshoot at 25% while keeping the + // timer rate proportional to the window rather than to the traffic. The + // floor stops a tiny `idle` turning into a busy loop. + let tick = (idle / 4).max(Duration::from_millis(500)); + let mut ticker = interval(tick); + ticker.set_missed_tick_behavior(MissedTickBehavior::Delay); + Self { + ticker, + idle_ticks: 0, + max_idle_ticks: (idle.as_millis() / tick.as_millis()).max(1) as u32, + last_seen: 0, + } + } + + /// Wait for the next sample point. + /// + /// Cancel-safe: `Interval::tick` is, and nothing else is held across it, so + /// this can sit in a `select!` arm that loses the race. + pub(crate) async fn tick(&mut self) { + self.ticker.tick().await; + } + + /// Record where the relay has got to. `true` means it has been stalled for + /// the whole idle window. + pub(crate) fn stalled(&mut self, progress: u64) -> bool { + if progress == self.last_seen { + self.idle_ticks += 1; + self.idle_ticks >= self.max_idle_ticks + } else { + self.idle_ticks = 0; + self.last_seen = progress; + false + } + } + + /// Resolve once the relay has been idle for the whole window. + /// + /// For relays with no loop of their own to poll from -- a spliced + /// connection is two joined transfers, not a `select!` -- so this is raced + /// against the transfer instead. + pub(crate) async fn wait_until_stalled(mut self, progress: &AtomicU64) { + // The first tick completes immediately; consume it so the first real + // sample is a full tick away. + self.tick().await; + loop { + self.tick().await; + if self.stalled(progress.load(Ordering::Relaxed)) { + return; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + #[tokio::test(start_paused = true)] + async fn a_stalled_relay_is_reported_after_the_window() { + let progress = Arc::new(AtomicU64::new(0)); + let watchdog = IdleWatchdog::new(Duration::from_secs(4)); + let start = tokio::time::Instant::now(); + watchdog.wait_until_stalled(&progress).await; + // Sampling four times per window means it fires within one tick of the + // window, never before it. + let waited = start.elapsed(); + assert!( + waited >= Duration::from_secs(4) && waited <= Duration::from_secs(6), + "fired after {waited:?}" + ); + } + + #[tokio::test(start_paused = true)] + async fn a_relay_that_keeps_moving_is_never_reported() { + let progress = Arc::new(AtomicU64::new(0)); + let bump = progress.clone(); + tokio::spawn(async move { + for _ in 0..20 { + tokio::time::sleep(Duration::from_secs(1)).await; + bump.fetch_add(1, Ordering::Relaxed); + } + }); + let watchdog = IdleWatchdog::new(Duration::from_secs(4)); + let fired = tokio::time::timeout( + Duration::from_secs(15), + watchdog.wait_until_stalled(&progress), + ) + .await; + assert!( + fired.is_err(), + "watchdog fired on a connection that was moving" + ); + } +} diff --git a/dstack/gateway/src/proxy/io_bridge.rs b/dstack/gateway/src/proxy/io_bridge.rs index 8c5dd7662..70cda0a77 100644 --- a/dstack/gateway/src/proxy/io_bridge.rs +++ b/dstack/gateway/src/proxy/io_bridge.rs @@ -2,10 +2,12 @@ // // SPDX-License-Identifier: Apache-2.0 +use super::idle::IdleWatchdog; use crate::config::ProxyConfig; -use anyhow::{Context, Result}; +use anyhow::{bail, Context, Result}; use bytes::BytesMut; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::TcpStream; use tokio::time::timeout; use tracing::{debug, trace}; @@ -25,6 +27,10 @@ struct OneDirection<'a, R, W> { reader: &'a mut R, writer: &'a mut W, next_step: NextStep, + /// Bumped whenever this direction makes progress. The watchdog samples it + /// instead of the clock, so the hot path costs an integer increment rather + /// than arming a timer per read/write/flush. + progress: u64, } impl OneDirection<'_, R, W> @@ -35,11 +41,12 @@ where async fn step(&mut self) -> Result { match self.next_step { NextStep::Read => { - let n = timeout(self.cfg.timeouts.idle, self.reader.read_buf(&mut self.buf)) + let n = self + .reader + .read_buf(&mut self.buf) .await - .ok() - .context("idle timeout")? .context("read error")?; + self.progress += 1; trace!(direction = %self.dir, "read: {n} bytes"); if n == 0 { self.next_step = NextStep::Shutdown; @@ -49,25 +56,19 @@ where Ok(false) } NextStep::Write => { - timeout( - self.cfg.timeouts.write, - self.writer.write_buf(&mut self.buf), - ) - .await - .ok() - .context("write timeout")? - .context("write error")?; + self.writer + .write_buf(&mut self.buf) + .await + .context("write error")?; + self.progress += 1; if self.buf.is_empty() { self.next_step = NextStep::Flush; } Ok(false) } NextStep::Flush => { - timeout(self.cfg.timeouts.write, self.writer.flush()) - .await - .ok() - .context("flush timeout")? - .context("flush error")?; + self.writer.flush().await.context("flush error")?; + self.progress += 1; self.next_step = NextStep::Read; Ok(false) } @@ -90,6 +91,30 @@ enum Rest { B2a(B), } +/// Relay between two TCP sockets. +/// +/// Same logic as [`bridge`], but splitting a `TcpStream` with its own `split()` +/// hands out borrowed halves, where the generic `tokio::io::split` has to wrap +/// the stream in a `BiLock`. That lock showed up at 1.1% of total time on the +/// small-request passthrough profile, and passthrough always has a plain socket +/// on both sides, so it never needs the generic path. +pub(crate) async fn bridge_tcp( + mut a: TcpStream, + mut b: TcpStream, + config: &ProxyConfig, +) -> Result<()> { + let buf_size = config.buffer_size; + if !config.timeouts.data_timeout_enabled { + tokio::io::copy_bidirectional_with_sizes(&mut a, &mut b, buf_size, buf_size) + .await + .context("failed to copy")?; + return Ok(()); + } + let (mut ra, mut wa) = a.split(); + let (mut rb, mut wb) = b.split(); + relay(&mut ra, &mut wa, &mut rb, &mut wb, config).await +} + pub(crate) async fn bridge(mut a: A, mut b: B, config: &ProxyConfig) -> Result<()> where A: AsyncRead + AsyncWrite + Unpin, @@ -106,31 +131,67 @@ where let (mut ra, mut wa) = tokio::io::split(a); let (mut rb, mut wb) = tokio::io::split(b); + relay(&mut ra, &mut wa, &mut rb, &mut wb, config).await +} +/// Drive both directions until each has seen EOF and been shut down. +async fn relay( + ra: &mut RA, + wa: &mut WA, + rb: &mut RB, + wb: &mut WB, + config: &ProxyConfig, +) -> Result<()> +where + RA: AsyncRead + Unpin, + WA: AsyncWrite + Unpin, + RB: AsyncRead + Unpin, + WB: AsyncWrite + Unpin, +{ + let buf_size = config.buffer_size; let mut a2b = OneDirection { dir: "a2b", cfg: config, buf: BytesMut::with_capacity(buf_size), - reader: &mut ra, - writer: &mut wb, + reader: ra, + writer: wb, next_step: NextStep::Read, + progress: 0, }; let mut b2a = OneDirection { dir: "b2a", cfg: config, buf: BytesMut::with_capacity(buf_size), - reader: &mut rb, - writer: &mut wa, + reader: rb, + writer: wa, next_step: NextStep::Read, + progress: 0, }; + // One watchdog for the whole connection replaces the per-operation timeouts: + // it samples both directions' progress counters, so a connection only dies + // when neither has moved. See `super::idle`. + let mut watchdog = IdleWatchdog::new(config.timeouts.idle); + watchdog.tick().await; // the first tick completes immediately + let mut rest; + // Progress of the direction that finishes first, frozen at that point. The + // watchdog samples a monotonic counter, so the drain phase has to keep + // adding it rather than restart from the surviving direction alone. + // Assigned on every path that leaves the loop, like `rest`. + let finished: u64; // Transfer data between a and b bidirectionally. loop { tokio::select! { + _ = watchdog.tick() => { + if watchdog.stalled(a2b.progress + b2a.progress) { + bail!("idle timeout"); + } + } done = a2b.step() => { if done? { // a to b is EOF, switch to b to a only + finished = a2b.progress; rest = Rest::B2a(b2a); drop(a2b); break; @@ -139,6 +200,7 @@ where done = b2a.step() => { if done? { // b to a is EOF, switch to a to b only + finished = b2a.progress; rest = Rest::A2b(a2b); drop(b2a); break; @@ -147,18 +209,39 @@ where } } - // One of the direction is closed, copy the other direction. + // One direction is closed; drain the other -- still watched. Half-close is + // not a licence to hang: before the watchdog existed each read carried its + // own `idle` timeout, so this phase was covered, and leaving it bare let a + // client hold a connection open until `timeouts.total` (5h) by + // half-closing against a backend that never replies. match &mut rest { - Rest::A2b(a2b) => loop { - if a2b.step().await? { - break; + Rest::A2b(a2b) => drain(a2b, &mut watchdog, finished).await, + Rest::B2a(b2a) => drain(b2a, &mut watchdog, finished).await, + } +} + +/// Pump the surviving direction to EOF, giving up if it stalls for `idle`. +async fn drain( + dir: &mut OneDirection<'_, R, W>, + watchdog: &mut IdleWatchdog, + finished: u64, +) -> Result<()> +where + R: AsyncRead + Unpin, + W: AsyncWrite + Unpin, +{ + loop { + tokio::select! { + _ = watchdog.tick() => { + if watchdog.stalled(finished + dir.progress) { + bail!("idle timeout"); + } } - }, - Rest::B2a(b2a) => loop { - if b2a.step().await? { - break; + done = dir.step() => { + if done? { + return Ok(()); + } } - }, + } } - Ok(()) } diff --git a/dstack/gateway/src/proxy/reuseport.rs b/dstack/gateway/src/proxy/reuseport.rs new file mode 100644 index 000000000..b2be7a93d --- /dev/null +++ b/dstack/gateway/src/proxy/reuseport.rs @@ -0,0 +1,58 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! `SO_REUSEPORT` listener groups for the thread-per-core proxy. +//! +//! By default the kernel picks a listener from a reuseport group by hashing the +//! connection's 4-tuple. With a moderate number of long-lived connections that +//! distribution is visibly uneven, and because thread-per-core cannot migrate a +//! connection between runtimes, the cores that drew fewer connections sit idle +//! while the others saturate. Measured at 16 connections over 4 cores, per-core +//! utilisation came out like `[99, 42, 101, 101]` and throughput swung 46% +//! between restarts purely on how the hash fell. +//! +//! `SO_ATTACH_REUSEPORT_CBPF` steering was tried as a fix and measured worse: +//! classic BPF has no maps, so round-robin is not expressible, and the one +//! signal it does expose -- the CPU handling the packet -- collapses the +//! distribution rather than evening it out. Connections are established in a +//! burst, so only the one or two client CPUs active at that moment are sampled, +//! and `cpu % n` sent every connection to the same one or two listeners: 2.3 of +//! 4 cores active and 167k rps against 245k for the plain hash. A real fix needs +//! per-connection state, i.e. an eBPF program with a counter map. +//! +//! What is kept here is binding the group in one place, in order, which keeps +//! listener order deterministic and the setup out of the per-thread path. + +use std::net::{SocketAddr, TcpListener}; + +use anyhow::{bail, Context, Result}; + +/// Bind `count` `SO_REUSEPORT` listeners on `addr`, in order. +/// +/// Returns one listener per worker, in the same order the kernel indexes them. +pub(crate) fn bind_group(addr: SocketAddr, count: usize, backlog: i32) -> Result> { + if count == 0 { + bail!("reuseport group needs at least one listener"); + } + let mut listeners = Vec::with_capacity(count); + for _ in 0..count { + let socket = socket2::Socket::new( + socket2::Domain::IPV4, + socket2::Type::STREAM, + Some(socket2::Protocol::TCP), + ) + .context("failed to create listening socket")?; + socket + .set_reuse_port(true) + .context("failed to set SO_REUSEPORT")?; + socket.set_reuse_address(true).ok(); + socket.set_nonblocking(true).ok(); + socket + .bind(&addr.into()) + .with_context(|| format!("failed to bind {addr}"))?; + socket.listen(backlog).context("failed to listen")?; + listeners.push(TcpListener::from(socket)); + } + Ok(listeners) +} diff --git a/dstack/gateway/src/proxy/splice.rs b/dstack/gateway/src/proxy/splice.rs new file mode 100644 index 000000000..6cf7b538f --- /dev/null +++ b/dstack/gateway/src/proxy/splice.rs @@ -0,0 +1,769 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! Zero-copy TCP relay using `splice(2)`. +//! +//! For the TLS-passthrough path both sides of the proxy are raw `TcpStream`s +//! and the gateway never inspects the (encrypted) payload. Instead of copying +//! bytes through a userspace buffer we move them kernel-side through a pipe with +//! `splice(2)`, which avoids two copies per direction and the associated CPU. +//! +//! Only used on Linux and only for the passthrough path; TLS-terminate still +//! uses the buffered bridge because one side is a decrypted rustls stream. + +use std::os::fd::{AsRawFd, OwnedFd}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use anyhow::{bail, Context, Result}; +use nix::fcntl::{fcntl, splice, FcntlArg, SpliceFFlags}; +use nix::sys::socket::{shutdown, Shutdown}; +use nix::unistd::pipe; +use or_panic::OptionOrPanic; +use tokio::io::{AsyncReadExt, AsyncWriteExt, Interest}; +use tokio::net::TcpStream; + +use super::idle::IdleWatchdog; +use crate::config::{EngageAfter, SpliceConfig}; + +/// How a relay endpoint's write side has to be closed. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum CloseKind { + /// Plain socket: a FIN says everything there is to say. + Tcp, + /// Kernel TLS: the peer also needs a `close_notify` alert, or it cannot + /// tell an orderly close from a truncated stream. + KernelTls, +} + +/// Bytes moved per `splice` syscall. Also the target pipe capacity so a full +/// read can be buffered kernel-side before draining to the destination. +const PIPE_CAPACITY: usize = 1 << 20; // 1 MiB + +fn set_pipe_capacity(fd: &OwnedFd, size: usize) { + // Best-effort: larger pipes mean fewer syscalls for bulk transfers. If the + // kernel rejects the size (e.g. over /proc/sys/fs/pipe-max-size) we simply + // keep the default capacity. + let _ = fcntl( + fd.as_raw_fd(), + FcntlArg::F_SETPIPE_SZ(size as std::os::raw::c_int), + ); +} + +fn errno_to_io(e: nix::errno::Errno) -> std::io::Error { + std::io::Error::from_raw_os_error(e as i32) +} + +/// Per-thread cache of splice pipes. +/// +/// Creating a pipe per direction per connection costs two `pipe2` calls plus +/// four descriptor closes, which is pure overhead for short connections: it +/// measurably raised passthrough connection-setup CPU (134 -> 142 us per +/// connection) while HAProxy, which pools its pipes, went the other way. +/// Reusing them keeps the bulk-transfer win without paying setup per +/// connection. Thread-local, so no locking -- and with thread-per-core a +/// connection stays on the thread that took the pipe. +/// Per thread. Each pooled pipe holds two descriptors, so the cap costs +/// `PIPE_POOL_MAX * 2 * workers` file descriptors at steady state -- 512 for the +/// 4-worker bench, ~4096 for a 32-worker gateway. That is fine given +/// `set_ulimit` raises RLIMIT_NOFILE to the hard limit, but it is why the number +/// is not larger. A 40-minute soak confirmed the pool fills to the cap and then +/// stops (fds 487 -> 543 -> flat, RSS flat at ~44 MB). +/// +/// With `SpliceConfig::release_idle_pipes` this cap stops being just a cache +/// size and becomes the actual descriptor bound: idle connections park their +/// pipes here instead of holding them, so steady-state use is +/// `2 * PIPE_POOL_MAX * workers` plus the pipes carrying data, rather than +/// `4 * connections`. Note that parking does not close anything -- the saving +/// comes from connections sharing a bounded set of pipes, not from idle +/// connections costing zero. +const PIPE_POOL_MAX: usize = 64; + +thread_local! { + static PIPE_POOL: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; + /// Pipes currently checked out of the pool. Test-only bookkeeping: together + /// with `PIPE_POOL.len()` it gives this thread's live pipe count exactly, + /// where counting `/proc/self/fd` would race with tests on other threads. + #[cfg(test)] + static PIPES_BORROWED: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// A splice pipe borrowed from the thread-local pool. +/// +/// Returned to the pool on drop, but only if the transfer drained it: a pipe +/// still holding bytes would corrupt the next connection that used it. +struct PooledPipe { + rd: Option, + wr: Option, + drained: bool, +} + +impl PooledPipe { + fn get() -> Result { + #[cfg(test)] + PIPES_BORROWED.with(|n| n.set(n.get() + 1)); + if let Some((rd, wr)) = PIPE_POOL.with(|p| p.borrow_mut().pop()) { + return Ok(Self { + rd: Some(rd), + wr: Some(wr), + drained: true, + }); + } + let (rd, wr) = pipe().context("failed to create splice pipe")?; + set_pipe_capacity(&wr, PIPE_CAPACITY); + Ok(Self { + rd: Some(rd), + wr: Some(wr), + drained: true, + }) + } + + /// The ends are `Option` only so `Drop` can move them into the pool, and + /// `Drop` is the last thing that runs, so both are always present here. + fn rd(&self) -> &OwnedFd { + self.rd.as_ref().or_panic("pipe read end present") + } + + fn wr(&self) -> &OwnedFd { + self.wr.as_ref().or_panic("pipe write end present") + } +} + +impl Drop for PooledPipe { + fn drop(&mut self) { + let (Some(rd), Some(wr)) = (self.rd.take(), self.wr.take()) else { + return; + }; + #[cfg(test)] + PIPES_BORROWED.with(|n| n.set(n.get() - 1)); + if !self.drained { + // Unknown residue: close instead of poisoning the pool. + return; + } + PIPE_POOL.with(|p| { + let mut pool = p.borrow_mut(); + if pool.len() < PIPE_POOL_MAX { + pool.push((rd, wr)); + } + }); + } +} + +/// Copy one direction (`src` -> `dst`) with splice until EOF, then half-close +/// the destination's write side. +/// +/// With `release_idle_pipes` the pipe is handed back to the pool whenever the +/// source runs dry, so a connection only pins descriptors while it actually has +/// bytes in flight. See `SpliceConfig::release_idle_pipes` for why that is safe +/// and what it costs. +async fn splice_one( + src: Arc, + dst: Arc, + release_idle_pipes: bool, + progress: &AtomicU64, + src_kind: CloseKind, + dst_close: CloseKind, +) -> Result<()> { + let mut pipe = PooledPipe::get()?; + + loop { + // Move a chunk from the source socket into the pipe. + // Try the syscall first and only wait for readiness when it actually + // blocks. `try_io` decides with one atomic read of the cached readiness + // flag, where `readable()` builds, polls and drops a `Readiness` future + // every time -- which showed up as ~2.4% of total time on the + // small-request passthrough profile, paid once per splice. + let n = loop { + match src.try_io(Interest::READABLE, || { + splice( + src.as_ref(), + None, + pipe.wr(), + None, + PIPE_CAPACITY, + SpliceFFlags::SPLICE_F_MOVE | SpliceFFlags::SPLICE_F_NONBLOCK, + ) + .map_err(errno_to_io) + }) { + Ok(n) => break n, + Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => { + // The pipe is empty on every path that reaches here: it was + // either just taken from the pool, or the drain below ran to + // completion, and a `WouldBlock` from the fill above added + // nothing. So it can be parked while the source is idle. + if release_idle_pipes { + drop(pipe); + src.readable().await.context("readable error")?; + pipe = PooledPipe::get()?; + } else { + src.readable().await.context("readable error")?; + } + } + // A kTLS socket refuses to splice a record that is not + // application data, and reports it as EINVAL. For a byte relay + // that is the end of the data stream, not a failure: it is how + // the client's close_notify arrives once the socket belongs to + // the kernel. Treating it as fatal took down the *other* + // direction too, so a client closing its request discarded the + // response the app had already written. + Err(ref e) + if matches!(src_kind, CloseKind::KernelTls) + && e.raw_os_error() == Some(libc::EINVAL) => + { + break 0; + } + Err(e) => return Err(e).context("splice src->pipe failed"), + } + }; + if n == 0 { + break; // EOF on source; the pipe was left empty by the last drain + } + // One relaxed increment per chunk -- not per byte, and dwarfed by the + // splice syscall it accompanies -- is what lets the idle watchdog see + // this connection without a timer per operation. + progress.fetch_add(1, Ordering::Relaxed); + // A chunk is in the pipe now: not safe to recycle until fully drained. + pipe.drained = false; + + // Drain the pipe fully into the destination socket. + let mut left = n; + while left > 0 { + match dst.try_io(Interest::WRITABLE, || { + splice( + pipe.rd(), + None, + dst.as_ref(), + None, + left, + SpliceFFlags::SPLICE_F_MOVE | SpliceFFlags::SPLICE_F_NONBLOCK, + ) + .map_err(errno_to_io) + }) { + Ok(m) => { + left -= m; + if left == 0 { + // Fully drained: the pipe is empty again, so it can go + // back to the pool even if the connection dies next. + pipe.drained = true; + } + } + Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => { + dst.writable().await.context("writable error")?; + } + Err(e) => return Err(e).context("splice pipe->dst failed"), + } + } + } + + // Propagate EOF: half-close the write side so the peer sees the close. + // + // A kTLS socket needs the TLS-level close first. `KtlsStream::poll_shutdown` + // would have sent it, but the offload path hands the bare descriptor to + // splice and never goes through that type again, so a plain shutdown here + // reaches the client as a FIN with no close_notify -- indistinguishable from + // a truncation attack to a client that checks. + if matches!(dst_close, CloseKind::KernelTls) { + let _ = ktls::send_close_notify(dst.as_raw_fd()); + } + let _ = shutdown(dst.as_raw_fd(), Shutdown::Write); + Ok(()) +} + +/// Bidirectional zero-copy relay between two TCP streams. +/// +/// `a` is the client side, `b` the app side; `a_close` describes what kind of +/// socket `a` is, which decides both how its write side is closed and how a +/// non-data record from it is interpreted. +/// `idle` is `None` when data timeouts are disabled. +pub(crate) async fn splice_bidirectional( + a: TcpStream, + b: TcpStream, + release_idle_pipes: bool, + idle: Option, + a_close: CloseKind, +) -> Result<()> { + // The single funnel for zero-copy relaying, so counting here covers both the + // passthrough gate and the post-kTLS handover. + super::stats::record_splice_engaged(); + let a = Arc::new(a); + let b = Arc::new(b); + let progress = AtomicU64::new(0); + let relay = async { + // `a` is the client side: it is the one that may be a kTLS socket, so + // it is the source kind for a2b and the destination kind for b2a. + let a2b = splice_one( + a.clone(), + b.clone(), + release_idle_pipes, + &progress, + a_close, + CloseKind::Tcp, + ); + let b2a = splice_one(b, a, release_idle_pipes, &progress, CloseKind::Tcp, a_close); + tokio::try_join!(a2b, b2a)?; + Ok(()) + }; + let Some(idle) = idle else { + return relay.await; + }; + // Spliced bytes never enter this process, so there is no read to hang a + // timeout on; the watchdog races the transfer instead. + tokio::select! { + result = relay => result, + () = IdleWatchdog::new(idle).wait_until_stalled(&progress) => bail!("idle timeout"), + } +} + +/// Per-thread cache of relay buffers, for the same reason as the pipe pool: +/// a pair of `buffer_size` allocations per connection is a real cost when the +/// connection only carries a few hundred bytes. +const BUF_POOL_MAX: usize = 32; +/// Phase 1 only runs until the splice threshold, so it does not need the full +/// `buffer_size` used for bulk copying. +const RELAY_BUF_SIZE: usize = 16 * 1024; + +thread_local! { + static BUF_POOL: std::cell::RefCell, Vec)>> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +struct PooledBufs { + a: Vec, + b: Vec, +} + +impl PooledBufs { + fn get(buf_size: usize) -> Self { + let size = buf_size.clamp(4096, RELAY_BUF_SIZE); + if let Some((a, b)) = BUF_POOL.with(|p| p.borrow_mut().pop()) { + return Self { a, b }; + } + Self { + a: vec![0u8; size], + b: vec![0u8; size], + } + } +} + +impl Drop for PooledBufs { + fn drop(&mut self) { + let a = std::mem::take(&mut self.a); + let b = std::mem::take(&mut self.b); + BUF_POOL.with(|p| { + let mut pool = p.borrow_mut(); + if pool.len() < BUF_POOL_MAX { + pool.push((a, b)); + } + }); + } +} + +/// Relay both directions with plain reads/writes until `gate` is reached, then +/// report whether splice should take over. +/// +/// Returns `true` if the gate was reached and both sockets are still open, +/// `false` if the connection finished first (in which case it is fully done). +/// +/// The gate is only tested at the tail of the loop, which is the one point +/// where both directions are quiescent: the `select!` arms each finish their +/// `write_all` before falling through, so nothing is buffered in userspace and +/// the sockets can be handed to the kernel safely. A timer arm inside the +/// `select!` would also be safe, but it would promote connections that are +/// merely idle -- allocating a pipe for a stream with nothing to move -- so +/// checking the clock on activity is both cheaper and better behaved. +async fn relay_until( + a: &mut TcpStream, + b: &mut TcpStream, + gate: &EngageAfter, + buf_size: usize, + idle: Option, +) -> Result { + let (mut ar, mut aw) = a.split(); + let (mut br, mut bw) = b.split(); + // Buffers come from a thread-local pool: allocating two of them per + // connection cost more than the pipe setup this phase exists to avoid. + let mut bufs = PooledBufs::get(buf_size); + let mut moved: u64 = 0; + let start = Instant::now(); + // `moved` doubles as the progress counter: it only advances when a transfer + // happened, which is exactly what the watchdog samples for. + let mut watchdog: Option = match idle { + Some(idle) => { + let mut w = IdleWatchdog::new(idle); + w.tick().await; // the first tick completes immediately + Some(w) + } + None => None, + }; + + loop { + // `finish_one` drains a half-closed connection without splice: once one + // side is done there is no long-lived stream left to optimise. + // + // Half-close is not end-of-connection. A client that finishes its + // request with `shutdown(SHUT_WR)` still expects the response, so the + // EOF is propagated to the *peer of the direction that ended* + // (`$closing`) and the opposite direction is pumped to completion. + // Shutting down the writer we are about to pump into instead would + // deliver the peer's EOF and drop everything still in flight. + macro_rules! finish_one { + ($closing:expr, $r:expr, $w:expr, $buf:expr) => {{ + $closing.shutdown().await.ok(); + loop { + let n = tokio::select! { + () = async { match watchdog.as_mut() { + Some(w) => w.tick().await, + None => std::future::pending().await, + } } => { + // The drain is watched too: a backend that accepts the + // request and then never answers would otherwise hold + // the connection until `timeouts.total`. + if watchdog.as_mut().is_some_and(|w| w.stalled(moved)) { + bail!("idle timeout"); + } + continue; + } + r = $r.read(&mut $buf) => r.context("read error")?, + }; + if n == 0 { + break; + } + // Not `write_all`: it is not cancel-safe, so it cannot sit + // in a `select!`, and leaving it outside meant a client + // that stopped reading blocked the drain in its write with + // the watchdog unpolled -- the mirror of the silent-backend + // stall, and just as good for holding a connection to + // `timeouts.total`. Single `write` calls are cancel-safe + // (nothing is written when the other branch wins), so the + // partial-write loop is ours to drive. + let mut written = 0usize; + while written < n { + let count = tokio::select! { + () = async { match watchdog.as_mut() { + Some(w) => w.tick().await, + None => std::future::pending().await, + } } => { + if watchdog.as_mut().is_some_and(|w| w.stalled(moved)) { + bail!("idle timeout"); + } + continue; + } + r = $w.write(&$buf[written..n]) => r.context("write error")?, + }; + if count == 0 { + bail!("write accepted no bytes"); + } + written += count; + // Per partial write, so a peer draining slowly still + // counts as progress and is not reaped for being slow. + moved += count as u64; + } + } + // Both directions are drained now; let the other peer see EOF. + $w.shutdown().await.ok(); + return Ok(false); + }}; + } + tokio::select! { + () = async { match watchdog.as_mut() { + Some(w) => w.tick().await, + // No idle timeout configured: this arm must never win. + None => std::future::pending().await, + } } => { + if watchdog.as_mut().is_some_and(|w| w.stalled(moved)) { + bail!("idle timeout"); + } + continue; + } + r = ar.read(&mut bufs.a) => { + let n = r.context("read from client failed")?; + // Client is done sending: tell the app, keep relaying its reply. + if n == 0 { finish_one!(bw, br, aw, bufs.b); } + bw.write_all(&bufs.a[..n]).await.context("write to app failed")?; + moved += n as u64; + } + r = br.read(&mut bufs.b) => { + let n = r.context("read from app failed")?; + // App is done replying: tell the client, keep relaying its input. + if n == 0 { finish_one!(aw, ar, bw, bufs.a); } + aw.write_all(&bufs.b[..n]).await.context("write to client failed")?; + moved += n as u64; + } + } + if gate.reached(moved, start) { + return Ok(true); + } + } +} + +/// Bidirectional relay that only switches to splice once the connection has +/// proven itself worth the syscalls. +/// +/// splice moves ~17 syscalls per connection to shift a small response (fill the +/// pipe, drain the pipe, plus readiness retries), where a read/write pair needs +/// two. Its benefit is per byte, its cost is per connection -- the same shape as +/// kTLS. Short request/response connections therefore never touch a pipe, while +/// connections that trip either gate still get zero-copy. +pub(crate) async fn splice_bidirectional_after( + mut a: TcpStream, + mut b: TcpStream, + config: &SpliceConfig, + buf_size: usize, + idle: Option, +) -> Result<()> { + if relay_until(&mut a, &mut b, &config.engage, buf_size, idle).await? { + splice_bidirectional(a, b, config.release_idle_pipes, idle, CloseKind::Tcp).await + } else { + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + use tokio::net::TcpListener; + + fn pool_len() -> usize { + PIPE_POOL.with(|pool| pool.borrow().len()) + } + + /// Descriptors this thread currently holds on a pipe: two per live pipe, + /// whether the pipe is parked in the pool or checked out by a relay. + /// + /// This is the number capacity planning cares about, so the tests assert on + /// it directly rather than trusting the pool alone as a proxy. + fn pipe_fds() -> usize { + 2 * (pool_len() + PIPES_BORROWED.with(|n| n.get())) + } + + async fn connected_pair() -> (TcpStream, TcpStream) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let client = TcpStream::connect(addr).await.unwrap(); + let (server, _) = listener.accept().await.unwrap(); + (client, server) + } + + /// Wire up `client <-> relay <-> backend` and start the relay. + /// + /// The relay runs on the same thread as the test (a `#[tokio::test]` + /// runtime is single-threaded), so it shares the thread-local pipe pool and + /// the test can observe what the relay parks there. + async fn start_relay(release_idle_pipes: bool) -> (TcpStream, TcpStream) { + let (client, inbound) = connected_pair().await; + let (outbound, backend) = connected_pair().await; + tokio::spawn(splice_bidirectional( + inbound, + outbound, + release_idle_pipes, + None, + CloseKind::Tcp, + )); + (client, backend) + } + + fn clear_pool() { + PIPE_POOL.with(|pool| pool.borrow_mut().clear()); + } + + /// Let the relay reach its next idle wait. + async fn settle() { + tokio::time::sleep(Duration::from_millis(50)).await; + } + + async fn expect(stream: &mut TcpStream, want: &[u8]) { + let mut got = vec![0u8; want.len()]; + stream.read_exact(&mut got).await.unwrap(); + assert_eq!(got, want); + } + + #[tokio::test] + async fn an_idle_relay_owns_no_pipe_when_release_is_on() { + clear_pool(); + let (mut client, mut backend) = start_relay(true).await; + client.write_all(b"ping").await.unwrap(); + expect(&mut backend, b"ping").await; + settle().await; + + // Only one pipe is ever created: each direction hands its pipe back + // before the other one asks for it, so an idle bidirectional relay + // converges on a single pipe shared through the pool -- and owns none + // of its own. + assert_eq!(pool_len(), 1, "the relay parked everything it borrowed"); + } + + #[tokio::test] + async fn an_idle_relay_pins_two_pipes_when_release_is_off() { + clear_pool(); + let (mut client, mut backend) = start_relay(false).await; + client.write_all(b"ping").await.unwrap(); + expect(&mut backend, b"ping").await; + settle().await; + + assert_eq!(pool_len(), 0, "both pipes stay pinned to the connection"); + } + + /// The headline claim: with release on, descriptor use tracks pipes that + /// are actually carrying data, not open connections. + /// + /// Note what this does *not* say. Released pipes stay open in the pool, so + /// the saving is not "idle connections cost nothing" but "idle connections + /// share": the steady-state bound moves from `4 * connections` to + /// `2 * PIPE_POOL_MAX * workers` plus whatever is in flight. + #[tokio::test] + async fn idle_connections_share_pipes_instead_of_each_pinning_four() { + clear_pool(); + assert_eq!(pipe_fds(), 0, "test starts with no pipes on this thread"); + let mut ends = Vec::new(); + for _ in 0..8 { + let (mut client, mut backend) = start_relay(true).await; + client.write_all(b"ping").await.unwrap(); + expect(&mut backend, b"ping").await; + // Without this the previous relay has not yet been polled back to + // its idle wait, so it still owns its pipe when the next one asks + // for one -- which is the in-flight case, not the idle case. + settle().await; + ends.push((client, backend)); + } + + // Held for the connection's lifetime this would be 8 * 4 = 32. + assert_eq!(pipe_fds(), 2, "8 idle connections share one pooled pipe"); + drop(ends); + } + + #[tokio::test] + async fn idle_connections_each_cost_four_descriptors_when_release_is_off() { + clear_pool(); + assert_eq!(pipe_fds(), 0, "test starts with no pipes on this thread"); + let mut ends = Vec::new(); + for _ in 0..8 { + let (mut client, mut backend) = start_relay(false).await; + client.write_all(b"ping").await.unwrap(); + expect(&mut backend, b"ping").await; + ends.push((client, backend)); + } + settle().await; + + assert_eq!(pipe_fds(), 8 * 4, "two pipes pinned per relay"); + drop(ends); + } + + #[tokio::test] + async fn data_survives_repeated_park_and_reacquire() { + clear_pool(); + let (mut client, mut backend) = start_relay(true).await; + + // Each gap forces the relay to park its pipe and take a fresh one from + // the pool, which is where a stale or half-drained pipe would corrupt + // the stream. + for i in 0..5u8 { + let up = [b'a' + i; 16]; + client.write_all(&up).await.unwrap(); + expect(&mut backend, &up).await; + settle().await; + + let down = [b'A' + i; 16]; + backend.write_all(&down).await.unwrap(); + expect(&mut client, &down).await; + settle().await; + } + + assert_eq!(pool_len(), 1); + } + + #[tokio::test] + async fn a_payload_larger_than_one_splice_still_arrives_intact() { + let (mut client, mut backend) = start_relay(true).await; + let payload: Vec = (0..512 * 1024).map(|i| (i % 251) as u8).collect(); + let sender = tokio::spawn({ + let payload = payload.clone(); + async move { + client.write_all(&payload).await.unwrap(); + client + } + }); + let mut got = vec![0u8; payload.len()]; + backend.read_exact(&mut got).await.unwrap(); + assert_eq!(got, payload); + drop(sender.await.unwrap()); + } + + #[tokio::test] + async fn eof_propagates_with_release_enabled() { + let (client, mut backend) = start_relay(true).await; + drop(client); + let mut got = Vec::new(); + backend.read_to_end(&mut got).await.unwrap(); + assert!(got.is_empty()); + } + + /// A gate that no test connection will ever reach, so the relay stays in + /// the pre-splice phase for the whole exchange. + fn ungated() -> SpliceConfig { + SpliceConfig { + engage: EngageAfter { + after_bytes: Some(1 << 30), + after_duration: None, + }, + release_idle_pipes: false, + } + } + + async fn start_gated_relay() -> (TcpStream, TcpStream) { + let (client, inbound) = connected_pair().await; + let (outbound, backend) = connected_pair().await; + tokio::spawn(async move { + splice_bidirectional_after(inbound, outbound, &ungated(), 16 * 1024, None).await + }); + (client, backend) + } + + /// Half-closing a request must not cost the response: the client shuts down + /// its write side, and the backend replies afterwards. + #[tokio::test] + async fn response_survives_a_client_half_close_before_the_gate() { + let (mut client, mut backend) = start_gated_relay().await; + + client.write_all(b"ping").await.unwrap(); + client.shutdown().await.unwrap(); + + let mut req = vec![0u8; 4]; + backend.read_exact(&mut req).await.unwrap(); + assert_eq!(&req, b"ping"); + // The app sees the client's EOF but is still free to answer. + let mut trailing = Vec::new(); + backend.read_to_end(&mut trailing).await.unwrap(); + assert!(trailing.is_empty()); + backend.write_all(b"pong").await.unwrap(); + drop(backend); + + let mut resp = Vec::new(); + client.read_to_end(&mut resp).await.unwrap(); + assert_eq!(resp, b"pong", "client lost the response after half-closing"); + } + + /// The mirror image: the backend finishes first and the client is still + /// sending. Its remaining bytes have to reach the app. + #[tokio::test] + async fn request_survives_a_backend_half_close_before_the_gate() { + let (mut client, mut backend) = start_gated_relay().await; + + backend.write_all(b"early").await.unwrap(); + backend.shutdown().await.unwrap(); + + let mut resp = vec![0u8; 5]; + client.read_exact(&mut resp).await.unwrap(); + assert_eq!(&resp, b"early"); + + client.write_all(b"late").await.unwrap(); + drop(client); + + let mut got = Vec::new(); + backend.read_to_end(&mut got).await.unwrap(); + assert_eq!(got, b"late", "app lost the request after half-closing"); + } +} diff --git a/dstack/gateway/src/proxy/stats.rs b/dstack/gateway/src/proxy/stats.rs new file mode 100644 index 000000000..f468e6412 --- /dev/null +++ b/dstack/gateway/src/proxy/stats.rs @@ -0,0 +1,125 @@ +// SPDX-FileCopyrightText: © 2024-2025 Phala Network +// +// SPDX-License-Identifier: Apache-2.0 + +//! What the data path is actually doing with kTLS and splice. +//! +//! Neither option is simply on or off. kTLS can be cleared at startup by the +//! capability probe, and both engage per connection only once their gate fires, +//! so the configured value does not tell an operator what is running -- and the +//! difference is worth 2-4x in per-connection memory. Two things are reported: +//! the effective mode, which answers "is it on at all", and per-connection +//! counters, which answer "is it engaging". +//! +//! All counters are monotonic since process start and are read without +//! synchronisation, so a snapshot can be marginally inconsistent between +//! fields. That is the right trade for numbers whose purpose is to be watched +//! over time; the alternative costs an ordering on the connection path. + +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; + +use dstack_gateway_rpc::ProxyAccelStatus; + +use crate::config::{EngageAfter, ProxyConfig}; + +/// Connections handed to the kernel's TLS ULP. +static KTLS_OFFLOADED: AtomicU64 = AtomicU64::new(0); +/// Connections where that handover failed. Non-zero here means connections are +/// being dropped or truncated at the gate, not merely running unaccelerated. +static KTLS_OFFLOAD_FAILED: AtomicU64 = AtomicU64::new(0); +/// Connections that entered a zero-copy splice relay. Counted on both paths -- +/// TLS passthrough, and terminate once kTLS has handed the socket over -- since +/// from the relay's point of view they are the same thing. +static SPLICE_ENGAGED: AtomicU64 = AtomicU64::new(0); +/// Set when the startup probe finds no TLS ULP. Without it a probe-disabled +/// gateway is indistinguishable from one that was never configured for kTLS, +/// because both end up with `ProxyConfig::ktls` unset. +static KTLS_UNSUPPORTED: AtomicBool = AtomicBool::new(false); + +/// Record that the kernel cannot do kTLS, so the reported mode can say why it +/// is off rather than just that it is. +pub(crate) fn mark_ktls_unsupported() { + KTLS_UNSUPPORTED.store(true, Ordering::Relaxed); +} + +/// Record the outcome of handing one connection to the kernel. +/// +/// Takes the result so both offload sites -- immediate and gated -- count the +/// same way without repeating the match. +pub(crate) fn record_ktls_offload(result: Result) -> Result { + let counter = if result.is_ok() { + &KTLS_OFFLOADED + } else { + &KTLS_OFFLOAD_FAILED + }; + counter.fetch_add(1, Ordering::Relaxed); + result +} + +/// Record that one connection started splice relaying. +pub(crate) fn record_splice_engaged() { + SPLICE_ENGAGED.fetch_add(1, Ordering::Relaxed); +} + +/// Snapshot the effective acceleration state for the `Status` RPC. +pub fn accel_status(config: &ProxyConfig) -> ProxyAccelStatus { + ProxyAccelStatus { + ktls_mode: ktls_mode(config), + splice_mode: match &config.tcp_splice { + Some(splice) => splice.engage.to_string(), + None => "off".to_string(), + }, + ktls_offloaded: KTLS_OFFLOADED.load(Ordering::Relaxed), + ktls_offload_failed: KTLS_OFFLOAD_FAILED.load(Ordering::Relaxed), + splice_engaged: SPLICE_ENGAGED.load(Ordering::Relaxed), + } +} + +fn ktls_mode(config: &ProxyConfig) -> String { + describe_ktls( + config.ktls.as_ref(), + KTLS_UNSUPPORTED.load(Ordering::Relaxed), + ) +} + +/// Split out from the statics so it can be tested without mutating global state +/// that other tests in this process would then see. +fn describe_ktls(ktls: Option<&EngageAfter>, unsupported: bool) -> String { + match ktls { + Some(engage) => engage.to_string(), + None if unsupported => "disabled (kernel has no TLS ULP)".to_string(), + None => "off".to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn gates_read_as_when_they_engage() { + let gate = |bytes, secs: Option| EngageAfter { + after_bytes: bytes, + after_duration: secs.map(std::time::Duration::from_secs), + }; + assert_eq!(gate(None, None).to_string(), "immediate"); + assert_eq!(gate(Some(65536), None).to_string(), "after 64 KiB"); + assert_eq!(gate(Some(1 << 20), None).to_string(), "after 1 MiB"); + assert_eq!(gate(Some(1000), None).to_string(), "after 1000 B"); + assert_eq!(gate(None, Some(5)).to_string(), "after 5s"); + assert_eq!(gate(Some(65536), Some(5)).to_string(), "after 64 KiB or 5s"); + } + + #[test] + fn a_probe_disabled_gateway_does_not_look_unconfigured() { + assert_eq!(describe_ktls(None, false), "off"); + assert_eq!( + describe_ktls(None, true), + "disabled (kernel has no TLS ULP)" + ); + // A kernel that cannot do kTLS is only interesting while kTLS is off; + // once it is on, the gate is the useful thing to report. + let gate = EngageAfter::default(); + assert_eq!(describe_ktls(Some(&gate), true), "immediate"); + } +} diff --git a/dstack/gateway/src/proxy/tls_passthough.rs b/dstack/gateway/src/proxy/tls_passthough.rs index 3848883d1..0d4f9abcd 100644 --- a/dstack/gateway/src/proxy/tls_passthough.rs +++ b/dstack/gateway/src/proxy/tls_passthough.rs @@ -20,7 +20,7 @@ use crate::{ }; use super::{ - io_bridge::bridge, + io_bridge::bridge_tcp, port_policy::{filter_allowed_addresses, should_send_pp}, AddressGroup, }; @@ -199,8 +199,29 @@ pub(crate) async fn connect_multiple_hosts( ) -> Result<(TcpStream, EnteredCounter, String)> { check_connection_limit(&addresses, max_connections, app_id)?; + let mut candidates = addresses.into_iter(); + let Some(first) = candidates.next() else { + bail!("no addresses to connect to app <{app_id}>"); + }; + + // Fast path: with a single candidate there is nothing to race, so skip the + // JoinSet and the task spawn it needs. That allocation and scheduling + // happened on every connection, and single-address apps are the common + // case. + if candidates.as_slice().is_empty() { + let addr = first; + let counter = addr.counter.enter(); + let ip = addr.ip; + debug!("connecting to {ip}:{port}"); + let connection = TcpStream::connect((ip, port)) + .await + .map_err(|e| anyhow::anyhow!("failed to connect to app@{ip}:{port}: {e}"))?; + let _ = connection.set_nodelay(true); + return Ok((connection, counter, addr.instance_id)); + } + let mut join_set = JoinSet::new(); - for addr in addresses { + for addr in std::iter::once(first).chain(candidates) { let counter = addr.counter.enter(); let ip = addr.ip; let instance_id = addr.instance_id; @@ -228,6 +249,9 @@ pub(crate) async fn connect_multiple_hosts( } } }; + // Disable Nagle on the upstream socket for the same reason as the inbound + // side: avoid delayed-ACK stalls on small proxied messages. + let _ = connection.set_nodelay(true); debug!("connected to {:?}", connection.peer_addr()); Ok((connection, counter, instance_id)) } @@ -259,9 +283,31 @@ pub(crate) async fn proxy_to_app( .write_all(&buffer) .await .context("failed to write to app")?; - bridge(inbound, outbound, &state.config.proxy) - .await - .context("failed to copy between inbound and outbound")?; + if let Some(gate) = &state.config.proxy.tcp_splice { + // Passthrough is a pure TCP relay: move bytes kernel-side with splice. + // Both ends are plain sockets here, so a FIN is the whole close. + let idle = state.config.proxy.idle_timeout(); + if gate.engage.is_immediate() { + super::splice::splice_bidirectional( + inbound, + outbound, + gate.release_idle_pipes, + idle, + super::splice::CloseKind::Tcp, + ) + .await + .context("failed to splice between inbound and outbound")?; + } else { + let buf_size = state.config.proxy.buffer_size; + super::splice::splice_bidirectional_after(inbound, outbound, gate, buf_size, idle) + .await + .context("failed to relay between inbound and outbound")?; + } + } else { + bridge_tcp(inbound, outbound, &state.config.proxy) + .await + .context("failed to copy between inbound and outbound")?; + } Ok(()) } diff --git a/dstack/gateway/src/proxy/tls_terminate.rs b/dstack/gateway/src/proxy/tls_terminate.rs index 48136b1db..558ee67d2 100644 --- a/dstack/gateway/src/proxy/tls_terminate.rs +++ b/dstack/gateway/src/proxy/tls_terminate.rs @@ -110,6 +110,15 @@ pub(crate) fn create_acceptor_with_cert_resolver( CryptoProvider::AwsLcRs => rustls::crypto::aws_lc_rs::default_provider(), CryptoProvider::Ring => rustls::crypto::ring::default_provider(), }; + // Stateless session tickets. TLS 1.2 resumption already works via the + // in-memory session-ID cache, but TLS 1.3 resumption requires a ticketer; + // without one, every reconnect pays a full handshake (a large RSA signing + // cost on the server). Installing a ticketer restores resumption for 1.3. + let ticketer = match proxy_config.tls_crypto_provider { + CryptoProvider::AwsLcRs => rustls::crypto::aws_lc_rs::Ticketer::new(), + CryptoProvider::Ring => rustls::crypto::ring::Ticketer::new(), + } + .context("failed to create TLS session ticketer")?; let supported_versions = proxy_config .tls_versions .iter() @@ -125,6 +134,15 @@ pub(crate) fn create_acceptor_with_cert_resolver( .with_no_client_auth() .with_cert_resolver(cert_resolver); + config.ticketer = ticketer; + + // kTLS needs the negotiated traffic secrets so it can install them into + // the kernel's TLS ULP. This is opt-in because it moves session keys + // outside rustls' control (see the `ktls` config docs). + if proxy_config.ktls.is_some() { + config.enable_secret_extraction = true; + } + if h2 { config.alpn_protocols = vec![b"h2".to_vec()]; } @@ -270,6 +288,66 @@ impl Proxy { Ok(tls_stream) } + /// Accept a TLS connection keeping the `CorkStream` wrapper, so the + /// connection can be handed to the kernel later without re-wrapping. + async fn tls_accept_corked( + &self, + inbound: TcpStream, + buffer: Vec, + h2: bool, + ) -> Result>> { + let stream = ktls::CorkStream::new(MergedStream { + buffer, + buffer_cursor: 0, + inbound, + }); + let acceptor = if h2 { + &self.h2_acceptor + } else { + &self.acceptor + }; + timeout( + self.config.proxy.timeouts.handshake, + acceptor.accept(stream), + ) + .await + .context("handshake timeout")? + .context("failed to accept tls connection") + } + + /// Accept a TLS connection and hand the socket over to kernel TLS. + /// + /// The handshake still runs in rustls; afterwards the negotiated keys are + /// installed into the kernel TLS ULP so record encryption happens there. + /// The inner IO must be wrapped in `CorkStream` because that is the only + /// way to drain a rustls stream cleanly at a record boundary. + async fn tls_accept_ktls( + &self, + inbound: TcpStream, + buffer: Vec, + h2: bool, + ) -> Result> { + let stream = ktls::CorkStream::new(MergedStream { + buffer, + buffer_cursor: 0, + inbound, + }); + let acceptor = if h2 { + &self.h2_acceptor + } else { + &self.acceptor + }; + let tls_stream = timeout( + self.config.proxy.timeouts.handshake, + acceptor.accept(stream), + ) + .await + .context("handshake timeout")? + .context("failed to accept tls connection")?; + super::stats::record_ktls_offload(ktls::config_ktls_server(tls_stream).await) + .context("failed to enable kernel TLS") + } + pub(super) async fn proxy( &self, inbound: TcpStream, @@ -291,20 +369,136 @@ impl Proxy { .with_context(|| format!("app <{app_id}> not found"))?; let addresses = filter_allowed_addresses(self, addresses, app_id, port)?; debug!("selected top n hosts: {addresses:?}"); - let tls_stream = self.tls_accept(inbound, buffer, h2).await?; + if let Some(ktls) = &self.config.proxy.ktls { + let splice = self.config.proxy.tcp_splice.as_ref(); + // A gated offload only pays off if the socket is spliced afterwards, + // so it needs both sections configured. + if let Some(splice) = splice.filter(|_| !ktls.is_immediate()) { + // Adaptive: stay in userspace rustls until the connection proves + // itself worth the offload, then hand it to the kernel. + let tls_stream = self.tls_accept_corked(inbound, buffer, h2).await?; + let (mut outbound, _counter, instance_id) = + self.connect_upstream(addresses, port, app_id).await?; + self.send_pp_header(&mut outbound, &instance_id, port, pp_header) + .await?; + return super::adaptive_ktls::relay_with_adaptive_offload( + tls_stream, + outbound, + ktls, + splice, + self.config.proxy.idle_timeout(), + ) + .await; + } + let tls_stream = self.tls_accept_ktls(inbound, buffer, h2).await?; + if let Some(splice) = splice { + // With kTLS the socket carries plaintext from userspace's point + // of view, so the payload can be relayed with splice and never + // enters this process at all. + let (mut outbound, _counter, instance_id) = + self.connect_upstream(addresses, port, app_id).await?; + self.send_pp_header(&mut outbound, &instance_id, port, pp_header) + .await?; + let (drained, stream) = tls_stream.into_raw(); + let (buffered, tcp) = stream.into_parts(); + // These two are not the same kind of data and must not be + // treated as interchangeable: `drained` is plaintext rustls + // decrypted past the handshake and owes to the app, while + // `buffered` is whatever raw *ciphertext* was left in the SNI + // sniff buffer. Forwarding the latter would hand the app TLS + // records to interpret as application data. + // + // A completed handshake always consumes the sniff buffer, so + // this is unreachable rather than merely unlikely -- but it is + // cheap to refuse instead of finding out by corrupting a stream. + if !buffered.is_empty() { + bail!( + "{} bytes of unconsumed ciphertext at kTLS handover", + buffered.len() + ); + } + // Plaintext rustls already decrypted has to reach the app before + // the kernel starts moving bytes directly. + let drained = drained.unwrap_or_default(); + if !drained.is_empty() { + outbound + .write_all(&drained) + .await + .context("failed to flush drained data to app")?; + } + // `tcp` is now a kernel-TLS socket, so closing it needs a + // close_notify and not just a FIN. + return super::splice::splice_bidirectional( + tcp, + outbound, + splice.release_idle_pipes, + self.config.proxy.idle_timeout(), + super::splice::CloseKind::KernelTls, + ) + .await + .context("ktls splice error"); + } + self.relay_to_app(tls_stream, addresses, port, app_id, pp_header) + .await + } else { + let tls_stream = self.tls_accept(inbound, buffer, h2).await?; + self.relay_to_app(tls_stream, addresses, port, app_id, pp_header) + .await + } + } + + /// Connect to the app and relay an already-terminated TLS stream to it. + /// + /// Generic over the accepted stream so the userspace-rustls and kTLS + /// paths share the same connect / PROXY-protocol / bridging logic. + /// Race a connection to the app's top-N addresses. + async fn connect_upstream( + &self, + addresses: super::AddressGroup, + port: u16, + app_id: &str, + ) -> Result<(TcpStream, crate::models::EnteredCounter, String)> { let max_connections = self.config.proxy.max_connections_per_app; - let (mut outbound, _counter, instance_id) = timeout( + timeout( self.config.proxy.timeouts.connect, connect_multiple_hosts(addresses, port, max_connections, app_id), ) .await .map_err(|_| anyhow!("connecting timeout"))? - .context("failed to connect to app")?; - if should_send_pp(self, &instance_id, port) { + .context("failed to connect to app") + } + + /// Forward the client's address to the app when its port policy asks for it. + async fn send_pp_header( + &self, + outbound: &mut TcpStream, + instance_id: &str, + port: u16, + pp_header: ProxyHeader, + ) -> Result<()> { + if should_send_pp(self, instance_id, port) { let pp_header_bin = proxy_protocol::encode(pp_header).context("failed to encode pp header")?; outbound.write_all(&pp_header_bin).await?; } + Ok(()) + } + + async fn relay_to_app( + &self, + tls_stream: S, + addresses: super::AddressGroup, + port: u16, + app_id: &str, + pp_header: ProxyHeader, + ) -> Result<()> + where + S: AsyncRead + AsyncWrite + Unpin, + { + let (mut outbound, _counter, instance_id) = + self.connect_upstream(addresses, port, app_id).await?; + self.send_pp_header(&mut outbound, &instance_id, port, pp_header) + .await?; bridge( IgnoreUnexpectedEofStream::new(tls_stream), outbound, @@ -316,6 +510,17 @@ impl Proxy { } } +/// Give up the raw socket *and* anything still buffered in front of it. +/// +/// Deliberately not `Into`: that conversion existed, and it dropped +/// the remainder silently. Whatever is left here is raw ciphertext from the SNI +/// sniff, which cannot be forwarded to an app expecting plaintext and cannot be +/// pushed back once the socket belongs to the kernel -- so the only safe thing +/// is to make every caller look at it. +pub(crate) trait SocketParts { + fn into_socket_parts(self) -> (Vec, TcpStream); +} + #[pin_project::pin_project] struct MergedStream { buffer: Vec, @@ -346,6 +551,37 @@ impl AsyncRead for MergedStream { this.inbound.poll_read(cx, buf) } } +impl MergedStream { + /// Unwrap to the raw socket, returning any bytes still buffered from the + /// pre-handshake sniff. After a completed handshake the buffer is drained, + /// so the returned `Vec` is normally empty; callers that bypass the + /// `AsyncRead` impl (e.g. splice) must still handle a non-empty remainder. + fn into_parts(self) -> (Vec, TcpStream) { + let remaining = self.buffer[self.buffer_cursor.min(self.buffer.len())..].to_vec(); + (remaining, self.inbound) + } +} + +impl SocketParts for MergedStream { + fn into_socket_parts(self) -> (Vec, TcpStream) { + self.into_parts() + } +} + +impl std::os::fd::AsRawFd for MergedStream { + fn as_raw_fd(&self) -> std::os::fd::RawFd { + self.inbound.as_raw_fd() + } +} + +impl ktls::AsyncReadReady for MergedStream { + fn poll_read_ready(&self, cx: &mut Context<'_>) -> Poll> { + // Safe to defer to the socket: by the time kTLS takes over, the + // buffered ClientHello prefix has already been consumed by rustls. + self.inbound.poll_read_ready(cx) + } +} + impl AsyncWrite for MergedStream { fn poll_write( self: std::pin::Pin<&mut Self>, @@ -381,3 +617,46 @@ impl AsyncWrite for MergedStream { self.inbound.is_write_vectored() } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::AsyncReadExt as _; + use tokio::net::TcpListener; + + async fn merged_with(buffer: Vec) -> MergedStream { + let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); + let addr = listener.local_addr().expect("addr"); + let client = TcpStream::connect(addr).await.expect("connect"); + drop(client); + let (inbound, _) = listener.accept().await.expect("accept"); + MergedStream { + buffer, + buffer_cursor: 0, + inbound, + } + } + + /// The kTLS handover reads the remainder through this, and forwarding raw + /// ciphertext to an app expecting plaintext is the failure it guards + /// against -- so an unconsumed sniff buffer has to be visible, not silently + /// swallowed by the unwrap. + #[tokio::test] + async fn an_unconsumed_sniff_buffer_is_surfaced_not_dropped() { + let stream = merged_with(b"leftover ciphertext".to_vec()).await; + let (remainder, _socket) = stream.into_socket_parts(); + assert_eq!(remainder, b"leftover ciphertext"); + } + + /// The normal case: rustls drains the sniff buffer during the handshake, so + /// the handover sees nothing left and proceeds. + #[tokio::test] + async fn a_consumed_sniff_buffer_leaves_no_remainder() { + let mut stream = merged_with(b"clienthello".to_vec()).await; + let mut sink = vec![0u8; 11]; + stream.read_exact(&mut sink).await.expect("read"); + assert_eq!(&sink, b"clienthello"); + let (remainder, _socket) = stream.into_socket_parts(); + assert!(remainder.is_empty(), "got {remainder:?}"); + } +} diff --git a/dstack/gateway/src/web_routes/route_index.rs b/dstack/gateway/src/web_routes/route_index.rs index db414391f..c4e9907e0 100644 --- a/dstack/gateway/src/web_routes/route_index.rs +++ b/dstack/gateway/src/web_routes/route_index.rs @@ -23,7 +23,12 @@ pub async fn index(state: &State) -> anyhow::Result> { .acme_info() .await .context("Failed to get ACME info")?; - let model = Dashboard { status, acme_info }; + let accel = status.accel.clone().unwrap_or_default(); + let model = Dashboard { + status, + acme_info, + accel, + }; let html = model.render().context("Failed to render template")?; Ok(Html(html)) } diff --git a/dstack/gateway/templates/dashboard.html b/dstack/gateway/templates/dashboard.html index 55a296ff8..cf2659ec9 100644 --- a/dstack/gateway/templates/dashboard.html +++ b/dstack/gateway/templates/dashboard.html @@ -452,6 +452,39 @@

This Node

+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +
Data Path
kTLS{{ accel.ktls_mode }}Record encryption in the kernel. Cleared at startup if the + kernel has no TLS ULP.
Splice{{ accel.splice_mode }}Zero-copy relaying, once a connection trips the gate.
kTLS Offloads{{ accel.ktls_offloaded }}Connections handed to the kernel since start.
kTLS Offload Failures{{ accel.ktls_offload_failed }}Non-zero means connections are failing at the handover, not + merely running unaccelerated.
Splice Engagements{{ accel.splice_engaged }}Connections that entered a zero-copy relay, on either path.
+

Global Connections

Loading global connection statistics...

diff --git a/dstack/gateway/test-run/.gitignore b/dstack/gateway/test-run/.gitignore index 25d3f97c9..b1c90f813 100644 --- a/dstack/gateway/test-run/.gitignore +++ b/dstack/gateway/test-run/.gitignore @@ -1,3 +1,4 @@ /run/ .env /e2e/dstack-gateway +__pycache__/ diff --git a/dstack/gateway/test-run/TESTING.md b/dstack/gateway/test-run/TESTING.md index 0ff3269fc..62eccbd19 100644 --- a/dstack/gateway/test-run/TESTING.md +++ b/dstack/gateway/test-run/TESTING.md @@ -189,3 +189,42 @@ gh pr checks --repo Dstack-TEE/dstack --watch=false Expected result: all required checks pass, including `gateway`, `rust-checks`, `prek`, `reuse-lint`, and CodeQL. + +## Proxy data-path integration tests + +`test_proxy.sh` runs a real gateway process and asserts on what reaches the +wire, across every combination of the two gated optimisations and both proxy +paths. It runs in CI (`.github/workflows/gateway-proxy-tests.yml`); unlike +`test_suite.sh` it needs no root for the gateway itself, only one `sudo ip link +add` for the link the gateway expects at startup. + +```bash +cd gateway/test-run +./test_proxy.sh # builds the gateway if needed +GATEWAY_BIN=../../target/release/dstack-gateway ./test_proxy.sh +BASE_PORT=39000 ./test_proxy.sh # if the default range is busy +KEEP_LOGS=1 ./test_proxy.sh # keep the work dir on success +``` + +What it covers: + +| group | asserts | +|---|---| +| data path | payloads survive byte-for-byte on both paths, under and over each gate, and under concurrency | +| close | the app closing reaches the client as an orderly TLS shutdown, including after kTLS offload | +| `timeouts.idle` | the idle watchdog reaps on every relay path -- buffered, spliced, kTLS -- and `data_timeout_enabled = false` still opts out | +| kTLS engagement | offload happens, and only once the gate fires; no TLS decrypt errors | +| capability fallbacks | a kernel without the TLS ULP warns and keeps serving untruncated; an inert `connection_rebalance` warns | +| runtime modes | every `thread_per_core` / `connection_rebalance` combination serves traffic | +| half-close | a client that finishes its request still gets the reply, on both paths and every arm | +| Status RPC | the accel counters reflect the traffic that ran | + +The suite adapts to the host: groups that need a capability the kernel or the +privileges do not provide are reported as `SKIP` with the reason, never silently +passed. + +Half-close needs a TLS client that can send `close_notify` without waiting for +the peer's, which `ssl.SSLSocket` cannot express: its only shutdown is +bidirectional, and dropping to `shutdown(SHUT_WR)` sends a bare FIN, which +mid-TLS is a truncation the peer is right to reject. `proxy/tlsclient.py` +drives the TLS state machine over memory BIOs to do it properly. diff --git a/dstack/gateway/test-run/proxy/gwconfig.py b/dstack/gateway/test-run/proxy/gwconfig.py new file mode 100755 index 000000000..fe4f52b56 --- /dev/null +++ b/dstack/gateway/test-run/proxy/gwconfig.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 +"""Emit a gateway config for one arm of the proxy integration tests. + +Every knob the suite varies is written explicitly rather than left to the +shipped defaults, so a test says what it is testing and a default change shows +up as a test change rather than as a silent shift in what was covered. + +Usage: gwconfig.py DIR [key=value ...] + + splice=off|immediate|after:[:] + ktls=off|immediate|after:[:] + tpc=true|false thread_per_core + rebalance=true|false connection_rebalance + idle= timeouts.idle + data_timeout=true|false timeouts.data_timeout_enabled + workers= +""" + +import sys + + +def gate_section(name: str, spec: str, extra: str = "") -> str: + """Render `[core.proxy.]` for one of the two gated optimisations. + + Absent section = off, empty section = engage immediately, keys = gated; + the same three states the config documents. + """ + if spec == "off": + return "" + body = f"\n[core.proxy.{name}]\n" + if spec != "immediate": + _, _, rest = spec.partition(":") + parts = rest.split(":") + body += f"after_bytes = {parts[0]}\n" + if len(parts) > 1: + body += f'after_duration = "{parts[1]}"\n' + return body + extra + + +def main(): + """Write one arm's config to stdout.""" + d = sys.argv[1].rstrip("/") + o = dict(a.split("=", 1) for a in sys.argv[2:] if "=" in a) + + cert, key = f"{d}/certs/cert.pem", f"{d}/certs/key.pem" + cfg = f"""workers = 2 +address = "127.0.0.1:{o["rpc_port"]}" +[tls] +key = "{key}" +certs = "{cert}" +[tls.mutual] +ca_certs = "{cert}" +[core] +kms_url = "" +rpc_domain = "" +set_ulimit = false +[core.debug] +insecure_skip_attestation = true +insecure_enable_debug_rpc = false +[core.admin] +enabled = true +address = "127.0.0.1:{o["admin_port"]}" +auth_token = "{o["admin_token"]}" +[core.sync] +enabled = false +node_id = 1 +data_dir = "{d}/data" +[core.wg] +public_key = "" +private_key = "" +listen_port = {o["wg_port"]} +ip = "10.90.0.1/24" +reserved_net = ["10.90.0.1/32"] +client_ip_range = "10.90.0.0/25" +config_path = "{d}/wg.conf" +interface = "{o["wg_iface"]}" +endpoint = "10.90.0.1:{o["wg_port"]}" +[core.proxy] +listen_addr = "127.0.0.1" +listen_port = {o["proxy_port"]} +localhost_enabled = true +base_domain = "{o["base_domain"]}" +cert_chain = "{cert}" +cert_key = "{key}" +workers = {o.get("workers", "2")} +max_connections_per_app = 0 +buffer_size = 65536 +tls_versions = ["1.2"] +thread_per_core = {o.get("tpc", "true")} +connection_rebalance = {o.get("rebalance", "true")} + +[core.proxy.timeouts] +idle = "{o.get("idle", "10m")}" +data_timeout_enabled = {o.get("data_timeout", "true")} +""" + cfg += gate_section( + "tcp_splice", o.get("splice", "off"), extra="release_idle_pipes = true\n" + ) + cfg += gate_section("ktls", o.get("ktls", "off")) + print(cfg) + + +if __name__ == "__main__": + main() diff --git a/dstack/gateway/test-run/proxy/origin.py b/dstack/gateway/test-run/proxy/origin.py new file mode 100755 index 000000000..a950d2942 --- /dev/null +++ b/dstack/gateway/test-run/proxy/origin.py @@ -0,0 +1,202 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 +"""Origin server for the proxy integration tests. + +Serves the same content over plain HTTP (what the gateway's TLS-terminate path +talks to) and over TLS (what the passthrough path relays to), so a single test +can compare the two paths against one implementation. + +Endpoints: + /bytes/ `n` bytes of a deterministic pattern + /close/ the same, then close the connection (no keep-alive) + /halfclose reads the request, waits for the client's EOF, *then* replies -- + the shape that used to lose its response before the splice gate + /blackhole waits for the client's EOF and then never replies, holding the + connection open -- what a half-closed request looks like to the + relay's drain phase + /flood/ waits for the client's EOF and then sends `n` bytes as fast as + it can -- the drain phase with the *write* side blocked, if the + client has stopped reading + /trickle/ `n` small records spaced out in time, i.e. token streaming + /health "ok" + +Deterministic payloads mean a test can assert on a digest without a second +fetch, and the pattern is not all-zeroes so a truncated or misaligned relay +cannot accidentally look correct. +""" + +import hashlib +import os +import socket +import ssl +import sys +import threading +import time + +PATTERN = b"dstack-gateway-proxy-test-0123456789abcdef" + + +def payload(n: int) -> bytes: + """Return `n` bytes of the deterministic pattern.""" + reps = n // len(PATTERN) + 1 + return (PATTERN * reps)[:n] + + +def digest(n: int) -> str: + """Return the sha256 of what `payload(n)` returns.""" + return hashlib.sha256(payload(n)).hexdigest() + + +def _respond(conn, body: bytes, close: bool): + head = [ + b"HTTP/1.1 200 OK", + b"Content-Type: application/octet-stream", + b"Content-Length: %d" % len(body), + ] + head.append(b"Connection: close" if close else b"Connection: keep-alive") + conn.sendall(b"\r\n".join(head) + b"\r\n\r\n" + body) + + +def _read_request(conn) -> bytes | None: + buf = b"" + while b"\r\n\r\n" not in buf: + try: + chunk = conn.recv(65536) + except (OSError, ssl.SSLError): + return None + if not chunk: + return None + buf += chunk + return buf + + +def handle(conn): + """Serve one connection until it closes.""" + try: + while True: + req = _read_request(conn) + if req is None: + return + path = req.split(b" ")[1].decode() + + if path == "/health": + _respond(conn, b"ok", close=False) + elif path.startswith("/bytes/"): + _respond(conn, payload(int(path.rsplit("/", 1)[1])), close=False) + elif path.startswith("/close/"): + _respond(conn, payload(int(path.rsplit("/", 1)[1])), close=True) + return + elif path == "/halfclose": + # Wait for the client's half-close before answering. A relay that + # treats one direction's EOF as end-of-connection drops this. + conn.settimeout(20) + try: + while conn.recv(65536): + pass + except (OSError, ssl.SSLError): + pass + _respond(conn, payload(4096), close=True) + return + elif path == "/blackhole": + # Drain the client's half of the conversation, then hold the + # connection open without answering. The relay is now in its + # post-EOF drain, waiting on us; whether it ever gives up is + # what `timeouts.idle` is supposed to decide. + conn.settimeout(600) + try: + while conn.recv(65536): + pass + except (OSError, ssl.SSLError): + pass + time.sleep(600) + return + elif path.startswith("/flood/"): + # Drain the client's half, then push far more than the sockets + # in between can buffer. If the client is not reading, the + # relay's drain blocks in its *write*, which is a different + # stall from /blackhole and has to be watched too. + size = int(path.rsplit("/", 1)[1]) + conn.settimeout(600) + try: + while conn.recv(65536): + pass + except (OSError, ssl.SSLError): + pass + try: + conn.sendall( + b"HTTP/1.1 200 OK\r\nContent-Length: %d\r\n\r\n" % size + ) + chunk = payload(262144) + sent = 0 + while sent < size: + conn.sendall(chunk[: min(len(chunk), size - sent)]) + sent += len(chunk) + except (OSError, ssl.SSLError): + pass + return + elif path.startswith("/trickle/"): + count = int(path.rsplit("/", 1)[1]) + conn.sendall(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\n\r\n") + for _ in range(count): + conn.sendall(b"40\r\n" + payload(64) + b"\r\n") + time.sleep(0.05) + conn.sendall(b"0\r\n\r\n") + return + else: + conn.sendall(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n") + return + except (OSError, ssl.SSLError): + pass + finally: + try: + conn.close() + except OSError: + pass + + +def serve(port: int, ctx: ssl.SSLContext | None): + """Accept forever on `port`, wrapping in TLS when `ctx` is given.""" + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("127.0.0.1", port)) + srv.listen(512) + while True: + raw, _ = srv.accept() + raw.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + if ctx is not None: + try: + raw = ctx.wrap_socket(raw, server_side=True) + except (OSError, ssl.SSLError): + raw.close() + continue + threading.Thread(target=handle, args=(raw,), daemon=True).start() + + +def main(): + """Start the configured listeners and idle.""" + plain = int(os.environ.get("PLAIN_PORT", "0")) + tls = int(os.environ.get("TLS_PORT", "0")) + cert, key = os.environ["CERT"], os.environ["KEY"] + + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.load_cert_chain(cert, key) + + if plain: + threading.Thread(target=serve, args=(plain, None), daemon=True).start() + if tls: + threading.Thread(target=serve, args=(tls, ctx), daemon=True).start() + print("origin ready", flush=True) + try: + while True: + time.sleep(3600) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "digest": + print(digest(int(sys.argv[2]))) + else: + main() diff --git a/dstack/gateway/test-run/proxy/probe.py b/dstack/gateway/test-run/proxy/probe.py new file mode 100755 index 000000000..03695e615 --- /dev/null +++ b/dstack/gateway/test-run/proxy/probe.py @@ -0,0 +1,298 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 +"""Client-side probes for the gateway proxy integration tests. + +Each subcommand exercises one behaviour and prints a single machine-readable +verdict line, so the shell driver stays a list of expectations rather than a +pile of parsing. + +Deliberately uses raw sockets and `ssl` rather than an HTTP client: three of the +behaviours under test (half-close, TLS close_notify, idle reaping) are invisible +to a client that hides connection lifecycle from you. +""" + +import argparse +import hashlib +import socket +import ssl +import sys +import time + +import origin +from tlsclient import HalfCloseTlsClient + + +def tls_context() -> ssl.SSLContext: + """Build a client context that hides nothing about how a connection ended.""" + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + # The gateway ships TLS 1.2 by default; pin it so a version change shows up + # as a config diff rather than as a mystery here. + ctx.maximum_version = ssl.TLSVersion.TLSv1_2 + # Ragged EOF is one of the things under test, so never paper over it. + return ctx + + +def connect(args): + """Open one TLS connection to the gateway, routed by SNI.""" + raw = socket.create_connection((args.host, args.port), timeout=args.timeout) + return tls_context().wrap_socket( + raw, server_hostname=args.sni, suppress_ragged_eofs=False + ) + + +def request(sock, path: str, sni: str, close: bool = False): + """Send one HTTP request.""" + extra = "Connection: close\r\n" if close else "" + sock.sendall(f"GET {path} HTTP/1.1\r\nHost: {sni}\r\n{extra}\r\n".encode()) + + +def read_response(sock) -> tuple[bytes, str]: + """Return (body, how_it_ended).""" + buf = b"" + try: + while b"\r\n\r\n" not in buf: + chunk = sock.recv(65536) + if not chunk: + return b"", "eof_before_headers" + buf += chunk + except ssl.SSLEOFError: + return b"", "truncated_before_headers" + + head, body = buf.split(b"\r\n\r\n", 1) + length = None + for line in head.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + length = int(line.split(b":")[1]) + if length is None: + return body, "no_content_length" + + try: + while len(body) < length: + chunk = sock.recv(65536) + if not chunk: + return body, "clean_eof" + body += chunk + except ssl.SSLEOFError: + return body, "truncated" + return body, "complete" + + +# --- probes ----------------------------------------------------------------- + + +def probe_fetch(args): + """Check the payload survives the proxy byte for byte.""" + sock = connect(args) + request(sock, f"/bytes/{args.size}", args.sni) + body, how = read_response(sock) + sock.close() + got = hashlib.sha256(body).hexdigest() + ok = got == origin.digest(args.size) and how == "complete" + print( + f"verdict={'pass' if ok else 'FAIL'} bytes={len(body)} want={args.size} end={how}" + ) + return ok + + +def probe_halfclose(args): + """Check a client that half-closes its request still gets the response. + + Regression test for the gated relays shutting down the wrong half and + returning an empty, successful response. Needs a client that can send + `close_notify` without waiting for the peer's -- see `tlsclient`. + """ + client = HalfCloseTlsClient(args.host, args.port, args.sni, args.timeout) + try: + client.send(f"GET /halfclose HTTP/1.1\r\nHost: {args.sni}\r\n\r\n".encode()) + client.close_write() + body, how = client.read_http_response() + finally: + client.close() + ok = len(body) == 4096 and how == "complete" + print(f"verdict={'pass' if ok else 'FAIL'} bytes={len(body)} want=4096 end={how}") + return ok + + +def probe_close_notify(args): + """Check the app closing first reaches the client as an orderly TLS shutdown. + + Regression test: after kTLS offload the gateway used to close with a bare + FIN, which a strict client cannot tell from a truncation attack. + """ + sock = connect(args) + request(sock, f"/close/{args.size}", args.sni, close=True) + body, how = read_response(sock) + sock.close() + ok = len(body) == args.size and how in ("complete", "clean_eof") + print(f"verdict={'pass' if ok else 'FAIL'} bytes={len(body)} end={how}") + return ok + + +def probe_idle(args): + """Check whether a connection that goes quiet is reaped, as configured. + + Regression test: configuring splice or kTLS used to bypass the watchdog + entirely, leaving `timeouts.total` (5h) as the only bound. + """ + sock = connect(args) + request(sock, f"/bytes/{args.size}", args.sni) + body, _ = read_response(sock) + if len(body) != args.size: + print(f"verdict=FAIL setup: read {len(body)} of {args.size}") + return False + + time.sleep(args.wait) + try: + sock.settimeout(10) + request(sock, "/bytes/64", args.sni) + alive = bool(sock.recv(200)) + except (OSError, ssl.SSLError): + alive = False + finally: + sock.close() + + want_alive = args.expect == "alive" + ok = alive == want_alive + print( + f"verdict={'pass' if ok else 'FAIL'} " + f"observed={'alive' if alive else 'reaped'} expected={args.expect}" + ) + return ok + + +def probe_stalled_after_halfclose(args): + """Check a half-closed request whose backend never replies is still reaped. + + Regression test: the relays dropped out of their watchdog once one + direction hit EOF and drained the other with bare reads, so a client could + hold a connection open until `timeouts.total` -- five hours by default -- + by half-closing against a backend that accepts and stays silent. + """ + client = HalfCloseTlsClient(args.host, args.port, args.sni, args.timeout) + started = time.monotonic() + try: + client.send(f"GET /blackhole HTTP/1.1\r\nHost: {args.sni}\r\n\r\n".encode()) + client.close_write() + client.sock.settimeout(args.wait) + try: + client.recv() + reaped = True + except Exception: + reaped = False + finally: + client.close() + waited = time.monotonic() - started + print( + f"verdict={'pass' if reaped else 'FAIL'} " + f"{'reaped' if reaped else 'still open'} after {waited:.1f}s" + ) + return reaped + + +def probe_stalled_write_after_halfclose(args): + """Check a half-closed request is reaped when the *client* stops reading. + + The mirror of `stalled-halfclose`: there the backend goes silent and the + relay blocks in `read`, here the backend floods and the relay blocks in + `write` because this client never drains it. Watching only the read half + leaves this one running until `timeouts.total`. + """ + client = HalfCloseTlsClient(args.host, args.port, args.sni, args.timeout) + started = time.monotonic() + try: + client.send( + f"GET /flood/{args.size} HTTP/1.1\r\nHost: {args.sni}\r\n\r\n".encode() + ) + client.close_write() + # Deliberately read nothing: let every buffer between here and the + # backend fill, so the relay is stuck in its write. + time.sleep(args.wait) + # Now drain. A reaped connection ends after whatever was buffered; a + # live one keeps feeding us the whole flood. + client.sock.settimeout(10) + drained = 0 + reaped = False + while drained < args.size: + try: + chunk = client.recv() + except Exception: + reaped = True + break + if not chunk: + reaped = True + break + drained += len(chunk) + finally: + client.close() + waited = time.monotonic() - started + print( + f"verdict={'pass' if reaped else 'FAIL'} " + f"{'reaped' if reaped else 'still open'} after {waited:.1f}s, drained {drained}B" + ) + return reaped + + +def probe_concurrent(args): + """Check many simultaneous transfers all arrive intact. + + Exercises the parts a single-connection test cannot: the pipe pool, the + per-core balancer, and whatever the gate does under real concurrency. + """ + import concurrent.futures + + def one(_): + try: + sock = connect(args) + request(sock, f"/bytes/{args.size}", args.sni) + body, how = read_response(sock) + sock.close() + return ( + hashlib.sha256(body).hexdigest() == origin.digest(args.size) + and how == "complete" + ) + except Exception: + return False + + with concurrent.futures.ThreadPoolExecutor(max_workers=args.count) as pool: + results = list(pool.map(one, range(args.count))) + ok = all(results) + print( + f"verdict={'pass' if ok else 'FAIL'} " + f"intact={sum(results)}/{args.count} size={args.size}" + ) + return ok + + +PROBES = { + "fetch": probe_fetch, + "halfclose": probe_halfclose, + "close-notify": probe_close_notify, + "idle": probe_idle, + "stalled-halfclose": probe_stalled_after_halfclose, + "stalled-write-halfclose": probe_stalled_write_after_halfclose, + "concurrent": probe_concurrent, +} + + +def main(): + """Run one probe and exit non-zero if it failed.""" + ap = argparse.ArgumentParser() + ap.add_argument("probe", choices=sorted(PROBES)) + ap.add_argument("--host", default="127.0.0.1") + ap.add_argument("--port", type=int, required=True) + ap.add_argument("--sni", required=True) + ap.add_argument("--size", type=int, default=1024) + ap.add_argument("--count", type=int, default=16) + ap.add_argument("--wait", type=float, default=12.0) + ap.add_argument("--expect", default="reaped", choices=["alive", "reaped"]) + ap.add_argument("--timeout", type=float, default=60.0) + args = ap.parse_args() + sys.exit(0 if PROBES[args.probe](args) else 1) + + +if __name__ == "__main__": + main() diff --git a/dstack/gateway/test-run/proxy/tlsclient.py b/dstack/gateway/test-run/proxy/tlsclient.py new file mode 100644 index 000000000..c4210fb68 --- /dev/null +++ b/dstack/gateway/test-run/proxy/tlsclient.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 +"""A TLS client that can half-close. + +`ssl.SSLSocket` cannot express "I have finished sending, keep sending to me". +Its only shutdown is `unwrap()`, which sends our `close_notify` and then blocks +for the peer's -- but the peer will not send its own until it has finished +replying, which is the very thing under test. Going one level lower and doing +`shutdown(SHUT_WR)` on the socket sends a bare FIN, which mid-TLS is a +truncation rather than an orderly half-close, so the peer is right to abandon +the connection. + +Driving the TLS state machine over memory BIOs solves it: `unwrap()` on an +`SSLObject` writes the `close_notify` record into the outgoing BIO before it +raises `SSLWantReadError` waiting for the reply. Flushing that BIO and then +declining to finish the handshake is exactly a half-close -- the peer sees an +orderly end of our stream and can keep writing to us. +""" + +import socket +import ssl + + +class HalfCloseTlsClient: + """A TLS connection whose write side can be closed independently.""" + + def __init__(self, host: str, port: int, sni: str, timeout: float = 30.0): + """Connect and complete the TLS handshake over memory BIOs.""" + self.sock = socket.create_connection((host, port), timeout=timeout) + self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + self._incoming = ssl.MemoryBIO() + self._outgoing = ssl.MemoryBIO() + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + ctx.maximum_version = ssl.TLSVersion.TLSv1_2 + self._tls = ctx.wrap_bio(self._incoming, self._outgoing, server_hostname=sni) + self._eof = False + self._run(self._tls.do_handshake) + + # --- BIO plumbing ------------------------------------------------------- + + def _flush(self) -> None: + """Send whatever the TLS engine has queued.""" + data = self._outgoing.read() + if data: + self.sock.sendall(data) + + def _fill(self) -> bool: + """Feed one chunk of wire data in. False once the peer is done.""" + data = self.sock.recv(65536) + if not data: + self._incoming.write_eof() + self._eof = True + return False + self._incoming.write(data) + return True + + def _run(self, op): + """Drive one TLS operation to completion, pumping both BIOs.""" + while True: + try: + result = op() + except ssl.SSLWantReadError: + # Always flush first: the engine often needs to *send* something + # before the peer will send what it is waiting for, and skipping + # that deadlocks the handshake. + self._flush() + if not self._fill(): + raise + except ssl.SSLWantWriteError: + self._flush() + else: + self._flush() + return result + + # --- the interesting part ---------------------------------------------- + + def close_write(self) -> None: + """Send `close_notify` without waiting for the peer's. + + This is the half-close the rest of the suite could not express. The + `SSLWantReadError` is the engine asking for the peer's reply; ignoring + it is the point, since the peer owes us a response first. The TCP FIN + goes out too, so the far side of a passthrough relay sees a genuine + half-close rather than just a TLS alert. + """ + try: + self._tls.unwrap() + except (ssl.SSLWantReadError, ssl.SSLWantWriteError): + pass + self._flush() + self.sock.shutdown(socket.SHUT_WR) + + # --- ordinary I/O ------------------------------------------------------- + + def send(self, data: bytes) -> None: + """Write application data.""" + self._run(lambda: self._tls.write(data)) + + def recv(self, size: int = 65536) -> bytes: + """Read application data. `b""` means the peer closed cleanly.""" + while True: + try: + return self._tls.read(size) + except ssl.SSLWantReadError: + if self._eof: + return b"" + if not self._fill(): + return b"" + except ssl.SSLZeroReturnError: + # The peer's own close_notify: an orderly end of its stream. + return b"" + + def read_http_response(self) -> tuple[bytes, str]: + """Read one response, returning (body, how_it_ended).""" + buf = b"" + while b"\r\n\r\n" not in buf: + chunk = self.recv() + if not chunk: + return b"", "eof_before_headers" + buf += chunk + head, body = buf.split(b"\r\n\r\n", 1) + length = None + for line in head.split(b"\r\n"): + if line.lower().startswith(b"content-length:"): + length = int(line.split(b":")[1]) + if length is None: + return body, "no_content_length" + while len(body) < length: + chunk = self.recv() + if not chunk: + return body, "truncated" + body += chunk + return body, "complete" + + def close(self) -> None: + """Drop the connection.""" + try: + self.sock.close() + except OSError: + pass diff --git a/dstack/gateway/test-run/test_proxy.sh b/dstack/gateway/test-run/test_proxy.sh new file mode 100755 index 000000000..18944a30a --- /dev/null +++ b/dstack/gateway/test-run/test_proxy.sh @@ -0,0 +1,463 @@ +#!/usr/bin/env bash +# SPDX-FileCopyrightText: © 2026 Phala Network +# +# SPDX-License-Identifier: Apache-2.0 +# +# Integration tests for the gateway's proxy data path. +# +# Covers what unit tests cannot: a real gateway process relaying real +# connections, across every combination of the two gated optimisations +# (`tcp_splice`, `ktls`) and both proxy paths (TLS terminate, TLS passthrough). +# +# Complements `test_suite.sh`, which covers the control plane, WaveKV and the +# handshake cache. This one is about bytes on the wire. +# +# Requirements: python3, openssl, ip, sudo (once, to create the WireGuard-named +# link the gateway expects at startup). No root for the gateway itself. +# +# ./test_proxy.sh # build and run everything +# GATEWAY_BIN=/path/to/dstack-gateway ./test_proxy.sh +# KEEP_LOGS=1 ./test_proxy.sh # leave the work dir behind on success +set -uo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd -P)" +CORE_DIR="$(cd "$HERE/../.." && pwd -P)" +PROXY_DIR="$HERE/proxy" +WORK="${WORK:-$(mktemp -d /tmp/dstack-gw-proxy-test.XXXXXX)}" + +# Ports are derived from one base so a busy CI machine can shift the whole set. +BASE_PORT="${BASE_PORT:-38400}" +PROXY_PORT=$((BASE_PORT + 0)) +RPC_PORT=$((BASE_PORT + 1)) +ADMIN_PORT=$((BASE_PORT + 2)) +WG_PORT=$((BASE_PORT + 3)) +ORIGIN_PLAIN=$((BASE_PORT + 4)) +ORIGIN_TLS=$((BASE_PORT + 5)) + +BASE_DOMAIN="gwtest.local" +# `localhost_enabled` routes `localhost-[s]` to 127.0.0.1:, which is +# what lets these tests run without registering a CVM. +SNI_TERMINATE="localhost-$ORIGIN_PLAIN.$BASE_DOMAIN" +SNI_PASSTHROUGH="localhost-${ORIGIN_TLS}s.$BASE_DOMAIN" +ADMIN_TOKEN="proxy-integration-test" +WG_IFACE="${WG_IFACE:-gwtest0}" + +GATEWAY_BIN="${GATEWAY_BIN:-$CORE_DIR/target/release/dstack-gateway}" + +PASS=0 +FAIL=0 +SKIP=0 +FAILED_NAMES=() + +say() { printf '%s\n' "$*"; } +group() { printf '\n\033[1m== %s ==\033[0m\n' "$*"; } + +# $1 = test name, remaining args = command. The command must exit 0 to pass. +check() { + local name="$1"; shift + local out rc + out=$("$@" 2>&1); rc=$? + if [ $rc -eq 0 ]; then + PASS=$((PASS + 1)) + printf ' \033[32mPASS\033[0m %-52s %s\n' "$name" "$(tail -1 <<<"$out")" + else + FAIL=$((FAIL + 1)) + FAILED_NAMES+=("$name") + printf ' \033[31mFAIL\033[0m %-52s %s\n' "$name" "$(tail -1 <<<"$out")" + printf ' %s\n' "${out//$'\n'/$'\n' }" | tail -12 + fi +} + +skip() { + SKIP=$((SKIP + 1)) + printf ' \033[33mSKIP\033[0m %-52s %s\n' "$1" "$2" +} + +probe() { python3 "$PROXY_DIR/probe.py" "$@"; } + +# --- setup ------------------------------------------------------------------ + +cleanup() { + stop_gateway + [ -n "${ORIGIN_PID:-}" ] && kill "$ORIGIN_PID" 2>/dev/null + if [ -n "${WG_CREATED:-}" ]; then + sudo ip link del "$WG_IFACE" 2>/dev/null + fi + if [ $FAIL -eq 0 ] && [ -z "${KEEP_LOGS:-}" ]; then + rm -rf "$WORK" + else + say "work dir kept at $WORK" + fi +} +trap cleanup EXIT + +require() { + command -v "$1" >/dev/null || { say "missing required tool: $1"; exit 1; } +} + +setup() { + require python3; require openssl; require ip; require ss + mkdir -p "$WORK/certs" "$WORK/logs" + + # Fail fast on a port that is already taken. A stale listener does not just + # break startup: the probes reach *it* instead, and its config is not the arm + # under test, so the run reports a scatter of unrelated assertion failures. + # That cost a long time to diagnose once. + local port + local busy="" + for port in "$PROXY_PORT" "$RPC_PORT" "$ADMIN_PORT" "$ORIGIN_PLAIN" "$ORIGIN_TLS"; do + if ss -ltn "sport = :$port" 2>/dev/null | grep -q LISTEN; then + busy="$busy $port" + fi + done + if [ -n "$busy" ]; then + say "ports already in use:$busy" + say "stop whatever holds them, or re-run with BASE_PORT set to a free range" + exit 1 + fi + + if [ ! -x "$GATEWAY_BIN" ]; then + say "building the gateway (set GATEWAY_BIN to skip)" + (cd "$CORE_DIR" && cargo build --release -p dstack-gateway) || exit 1 + fi + + openssl req -x509 -newkey rsa:2048 -nodes -days 2 \ + -keyout "$WORK/certs/key.pem" -out "$WORK/certs/cert.pem" \ + -subj "/CN=$BASE_DOMAIN" \ + -addext "subjectAltName=DNS:$BASE_DOMAIN,DNS:*.$BASE_DOMAIN,IP:127.0.0.1" \ + 2>/dev/null || { say "failed to generate a test certificate"; exit 1; } + + # The gateway insists on a link with the configured name at startup. A real + # WireGuard device also makes `wg show` work, which the Status RPC needs; a + # dummy link is enough for the data path, so fall back to one rather than + # skipping every test on a host without the wireguard module. + if ip link show "$WG_IFACE" >/dev/null 2>&1; then + WG_KIND=preexisting + elif sudo ip link add "$WG_IFACE" type wireguard 2>/dev/null; then + WG_KIND=wireguard; WG_CREATED=1 + elif sudo ip link add "$WG_IFACE" type dummy 2>/dev/null; then + WG_KIND=dummy; WG_CREATED=1 + else + say "cannot create the link '$WG_IFACE' the gateway needs at startup"; exit 1 + fi + [ -n "${WG_CREATED:-}" ] && sudo ip link set "$WG_IFACE" up + say "link $WG_IFACE: $WG_KIND" + + CERT="$WORK/certs/cert.pem" KEY="$WORK/certs/key.pem" \ + PLAIN_PORT="$ORIGIN_PLAIN" TLS_PORT="$ORIGIN_TLS" \ + python3 "$PROXY_DIR/origin.py" >"$WORK/logs/origin.log" 2>&1 & + ORIGIN_PID=$! + for _ in $(seq 50); do + grep -q "origin ready" "$WORK/logs/origin.log" 2>/dev/null && break + sleep 0.2 + done + grep -q "origin ready" "$WORK/logs/origin.log" || { say "origin failed to start"; exit 1; } +} + +# --- gateway lifecycle ------------------------------------------------------ + +stop_gateway() { + [ -n "${GW_PID:-}" ] || return 0 + kill "$GW_PID" 2>/dev/null + for _ in $(seq 50); do + kill -0 "$GW_PID" 2>/dev/null || break + sleep 0.1 + done + kill -9 "$GW_PID" 2>/dev/null + GW_PID="" + # The suite restarts the gateway ~25 times. Killing the process is not enough: + # its listeners linger briefly, and the next arm then dies with EADDRINUSE on + # the RPC port -- which showed up as unrelated assertions failing. Wait for the + # ports to actually come back before handing them to the next arm. + local port + for port in "$PROXY_PORT" "$RPC_PORT" "$ADMIN_PORT"; do + for _ in $(seq 100); do + ss -ltn "sport = :$port" 2>/dev/null | grep -q LISTEN || break + sleep 0.1 + done + done +} + +# start_gateway