diff --git a/CLAUDE.md b/CLAUDE.md index 584a95c..cd107e6 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,8 +37,9 @@ TCP accept (SO_REUSEPORT per worker thread) └─ [suspended.rs register, emit 'suspended', await oneshot] └─ tls.rs TlsAcceptor::accept() with handshake timeout (if terminate_tls) └─ upstream.rs connect(Destination, peer_ip) → UpstreamStream - └─ tokio::io::copy_bidirectional_with_sizes wrapped in idle_timeout - (per-direction buffers from readBufferSize / client|upstreamReadBufferSize) + └─ copy::copy_bidirectional_lazy wrapped in idle_timeout + (readBufferSize / client|upstreamReadBufferSize is a per-direction MAXIMUM, not a + permanent allocation — see src/copy.rs) └─ RAII drop: BalancerGuard, ActiveGuard — all counter decrements happen here ``` @@ -127,6 +128,67 @@ Each tokio worker thread gets its own listening socket on the same address via ` ### Suspended connections via oneshot channels Each suspended connection gets a `tokio::sync::oneshot::channel`. The sender is stored in a `DashMap`. `resolveConnection()` removes the sender and fires it — synchronous from the JS side (no async needed). `oneshot` is used rather than `mpsc` because exactly one resolution is possible per connection. If no resolution arrives within `suspendTimeoutMs`, the `timeout(rx.await)` in `proxy_conn.rs` returns an error and the TCP stream is dropped. +### Escalating copy buffer instead of a permanent per-connection allocation (`src/copy.rs`) +`tokio::io::copy_bidirectional_with_sizes` allocates its two per-direction `CopyBuffer`s once and +holds them for the connection's whole life, whether or not it is transferring — at a million +mostly-idle MQTT subscribers, `readBufferSize × 2` held forever per connection is dead weight. +`copy_bidirectional_lazy` is a direct port of tokio's own `copy_bidirectional_impl`/`CopyBuffer` +state machine (`transfer_one_direction` + `TransferState::Running/ShuttingDown/Done`, `poll_fn` +over both directions), with one addition: `LazyCopyBuffer` starts each direction at a small fixed +size (`PROBE_BUFFER_SIZE`, 512 B) and escalates straight to the full configured `max_buf_size` +only once **two consecutive** reads exactly saturate the current buffer — real evidence of a +sustained burst — dropping straight back to the floor once the direction actually **parks** with +nothing left to write, not on the first under-capacity read. One full read isn't enough escalation +evidence: a message that happens to exactly match the current (small) buffer size is a coincidence, +not a burst, and escalating on it would leave the (now oversized) buffer resident for however long +the connection then sits idle, however large `max_buf_size` is configured. Requiring a second +confirming read bounds that coincidence to the floor size at the cost of one extra small-buffer +round trip on every genuine burst — negligible. Symmetrically, shrinking is deferred to the actual +park rather than the first under-capacity read: a connection with continuously active but +variably-sized traffic (never actually idle) would otherwise reallocate on every undersized read +only to grow right back on the next burst. + +**None of it is unconditional.** The saving scales with connection count and the per-burst cost of +growing/releasing a buffer does not, so the behaviour is gated on the proxy's live +`GlobalMetrics::active_connections` (`copy::LazyBufferGate`, config `lazyCopyBufferThreshold`, +default 1000). Below the threshold each direction allocates its full configured buffer once and +never resizes — byte-for-byte the old `copy_bidirectional_with_sizes` behaviour — so the six-stream +replication port-set pays nothing for a memory problem it does not have, while the 100k-subscriber +MQTT port-set is far above the threshold and gets the full mechanism. `0` engages always; a value +above peak concurrency disables it. The gauge is re-read at **every** resize decision rather than +latched per connection: a connection established while the proxy was quiet would otherwise hold a +full-size buffer for its whole life however busy the proxy later became, and long-lived connections +accumulating while idle is exactly the shape this exists for. + +**A recreate does not drop established connections.** `stop()` sends the shutdown broadcast (ending +the accept loops) and sleeps 100 ms; it never aborts connection tasks, and the runtime lives in the +napi wrap until GC. So the construction-frozen proxy fields (`readBufferSize`, its two overrides, +`lazyCopyBufferThreshold`, `workerThreads`) apply to *new* connections only — existing sessions keep +running on the old values for as long as they stay open. `__test__/server.spec.ts` pins both halves +of this: that editing one of those fields forces a recreate (the signature does its job, with a +route-only control proving it is the signature and not "recreates on every write"), and that a held +connection survives that recreate and still proxies. + +Reusing tokio's own poll-based structure (rather than `tokio::io::split` plus independent +per-direction `async fn`s, tried first) matters for two reasons found by an independent review of +that version: `split` wraps each side in an `Arc>` — two heap allocations and a +lock/unlock on every read/write/flush/shutdown, *per connection*, which itself scales with +connection count, exactly what this change exists to avoid — and a hand-rolled `async fn` pump +that just calls `write_all` does not flush before parking on the next read the way tokio's +`poll_copy` does. That gap is real: `write_all` only guarantees the data reached the writer's +internal buffer, not the wire (`tokio-rustls` in particular defers sending encrypted records +until `poll_flush`) — a TLS client that sends one request and waits for the reply would never see +it, both sides sitting idle until `idle_timeout` cleaned up the session. A non-blocking +`poll_read`-with-`Waker::noop()` peek was also tried, to opportunistically drain "whatever's +already queued" without the extra confirming iteration; under load it silently stranded a +connection's wakeup (reproduced empirically — an increasing fraction of connections stopped +responding as concurrency grew) and was dropped in favor of the two-consecutive-full-reads scheme, +which uses only ordinary polling. `readBufferSize`/`client|upstreamReadBufferSize` are therefore a +per-transfer *maximum*, not a permanent allocation. Propagating a failed `poll_shutdown()` (not +swallowing it) is what makes `try_join!`-equivalent short-circuiting work: an error on either +side — including a shutdown failure — ends the whole copy immediately rather than waiting on the +other direction. + ### Per-route Arc deduplication (TlsConfigCache) Routes that share the same cert+mTLS combination share a single `Arc` allocation. The cache key is `(sha256(cert_pem + key_pem), sha256(mtls_ca_pem))`. Built at config-parse time in `tls.rs::TlsConfigCache`. Important for deployments where many routes share a wildcard cert. diff --git a/README.md b/README.md index 0d2bb2d..fc2674c 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,7 @@ console.log('proxy listening on :443'); | `readBufferSize` | `number` | `8192` | Per-direction copy buffer size in bytes, clamped to `[512, 1048576]`. See [Copy buffers and per-connection memory](#copy-buffers-and-per-connection-memory) | | `clientReadBufferSize` | `number` | `readBufferSize` | Overrides `readBufferSize` for the client→upstream direction only | | `upstreamReadBufferSize` | `number` | `readBufferSize` | Overrides `readBufferSize` for the upstream→client direction only | +| `lazyCopyBufferThreshold` | `number` | `1000` | Active connections at or above which copy buffers escalate and release instead of being held at full size. `0` always; above peak concurrency disables it. See [Copy buffers and per-connection memory](#copy-buffers-and-per-connection-memory) | ### `ListenerConfig` @@ -757,35 +758,69 @@ docker run --rm -v $(pwd):/build -w /build \ ## Copy buffers and per-connection memory -Each proxied connection holds two copy buffers — one per direction — for its entire life, whether or -not it is transferring anything. So buffer memory scales with *connection count*, not with traffic: - -``` -buffer bytes = (clientReadBufferSize + upstreamReadBufferSize) × connections -``` - -Each unset override falls back to `readBufferSize`, so with symmetric sizing that is just -`2 × readBufferSize × connections`. At the 8192-byte default it is 16 KiB per connection: 4.0 GiB at -262k connections, 5.1 GiB at 333k. The asymmetric MQTT setting recommended below is -`1024 + 4096` = 5 KiB per connection, which is where the general form matters. -The knob is per proxy, and the right value is the opposite for the two shapes of traffic symphony -carries: +`readBufferSize` (and its per-direction overrides) is a *maximum*, not a permanent allocation. Each +direction starts at a small fixed floor (1 KiB total across both directions) and escalates to the +configured maximum only once it observes a sustained burst — two consecutive reads that fill the +current buffer — dropping straight back to the floor once the direction actually parks with nothing +left to write, not on every single under-capacity read (that would reallocate a connection that is +still continuously active but simply has variably-sized traffic). So only the memory *above* that +floor scales with concurrently bursting transfers; the 1 KiB/connection floor itself still scales +with connection count, same as before. A million idle MQTT subscribers cost about 0.95 GiB in floors +(1 KiB × 1,000,000 connections), not `readBufferSize × 2 × 1,000,000`. + +### When escalation is active: `lazyCopyBufferThreshold` + +The escalate/release behaviour above is not unconditional, because its economics are one-sided. The +memory it saves scales with connection count, while its cost — two allocations plus the zeroing of +the new buffer, per burst — does not. A replication port-set carrying six bulk streams would pay +that on every burst to save a few hundred KiB it was never short of; a hundred thousand parked MQTT +subscribers are the entire reason the mechanism exists. + +So it engages only once the proxy is actually carrying enough connections for the saving to be +worth having. `lazyCopyBufferThreshold` (default `1000`) is the proxy-wide active connection count +at or above which buffers escalate and release. Below it, each direction gets its full configured +buffer once and never resizes — identical to symphony's behaviour before this mechanism existed, +with no resize churn at all. + +The default puts the two shapes symphony carries on opposite sides without anyone configuring it: a +replication port-set (~6 connections) stays static, an MQTT fan-out port-set (100k+) escalates. Set +it to `0` to engage always, or above the port-set's peak concurrency to disable it entirely. It is +per proxy, so those two port-sets can differ. + +The count is re-read at each resize decision rather than fixed when a connection is established. +That matters for the case this exists for: connections that arrive while the proxy is quiet would +otherwise hold a full-size buffer for their entire lives no matter how busy it later got, and +long-lived connections accumulating while idle is precisely the MQTT shape. Instead they begin +releasing at their next park once the proxy crosses the threshold. + +``` +worst-case buffer bytes = (clientReadBufferSize + upstreamReadBufferSize) × connections bursting right now +``` + +That's the ceiling if every connection happened to be mid-burst simultaneously — a useful number to +size against, but not the steady-state cost, which sits near the 1 KiB/connection floor regardless of +`readBufferSize`. Each unset override falls back to `readBufferSize`, so with symmetric sizing the +ceiling is `2 × readBufferSize × connections bursting right now`. At the 8192-byte default that is 16 +KiB per bursting connection. The asymmetric MQTT setting recommended below is `1024 + 4096` = 5 KiB +per bursting connection, which is where the general form matters. +The knob still bounds how large a *single* transfer's buffer may grow, and the right bound is the +opposite for the two shapes of traffic symphony carries: | Traffic | Connections | Payloads | Suggested | |---|---|---|---| | Native MQTT (`8883`) | 100k–1M | hundreds of bytes | `clientReadBufferSize: 1024`, `upstreamReadBufferSize: 4096` | | HTTPS (`443`) | thousands | mixed | leave at the default | | Operations API (`9925`) | tens | can be large | leave at the default | -| Replication (`9933`) | ~6 | bulk streams | leave, or raise — 64 KiB across 6 connections is 768 KB total | +| Replication (`9933`) | ~6 | bulk streams | leave, or raise — 64 KiB across 6 connections bursting at once is 768 KB total | MQTT is worth splitting by direction: after `SUBSCRIBE` a client sends almost nothing but `PINGREQ`, -while the broker carries the whole fan-out. `1024`/`4096` is 5 KiB per connection against the -default's 16 KiB — 3.5 GiB saved at 333k connections — and buys more downstream headroom than a -symmetric 2048 would. +while the broker carries the whole fan-out. `1024`/`4096` bounds a bursting connection to 5 KiB +against the default's 16 KiB ceiling, and buys more downstream headroom than a symmetric 2048 would. Going small costs CPU, not correctness: a payload larger than the buffer is simply copied in more -iterations. On a TLS-terminating listener those extra iterations are not even syscalls, since the -reads come out of rustls's already-decrypted buffer. +iterations, and a connection that keeps bursting re-escalates after two full reads. On a +TLS-terminating listener those extra iterations are not even syscalls, since the reads come out of +rustls's already-decrypted buffer. Two limits on where these settings apply: @@ -796,16 +831,22 @@ Two limits on where these settings apply: PROXY-protocol routes, including every UDS route, take the plain path and are governed normally. - **A config reload cannot change these.** They are frozen when the proxy is constructed, so changing one makes `symphony-server` recreate the proxy rather than hot-swap it. `SO_REUSEPORT` - means there is no *bind* gap, but established connections on the old proxy are **not** drained — - `stop()` waits 100 ms and connection tasks are detached — so they are all dropped. On a - high-connection-count listener, treat a buffer-size edit as a reconnect event, not a live tune. + means there is no *bind* gap. Established connections on the old proxy are **not** drained and + **not** dropped either: `stop()` sends the shutdown broadcast, which ends the accept loops, then + sleeps 100 ms — it never aborts connection tasks, and the tokio runtime lives inside the addon + until the old proxy is garbage-collected. Those sessions keep running on the *old* buffer sizes + for as long as they stay open. So a buffer-size edit applies to new connections only, and on a + listener whose connections are long-lived by design (MQTT subscribers) that can mean the change + reaches almost nothing until those clients reconnect. Plan the edit around a reconnect rather + than expecting it to take effect fleet-wide on reload. > **Upgrading:** before this setting was applied to the copy loop, `readBufferSize` had no effect — -> every connection got 8 KiB per direction regardless of what the config said, and the default -> documented here was `65536`. A config that leaves it unset is unaffected. A config that *sets* it -> explicitly now gets what it asked for, so a value copied from the old documented default becomes -> 128 KiB per connection instead of 16 KiB. Drop or remove such a value before upgrading a -> high-connection-count deployment. +> every connection got a fixed 8 KiB per direction regardless of what the config said (held for the +> connection's whole life, not escalating/shrinking), and the default documented here was `65536`. A +> config that leaves it unset is unaffected. A config that *sets* it explicitly now gets what it +> asked for as a per-transfer ceiling, so a value copied from the old documented default raises that +> ceiling to 128 KiB per bursting connection instead of 16 KiB. Drop or remove such a value before +> upgrading a high-connection-count deployment. Two caveats when sizing a node from this: diff --git a/__test__/bench-copy-burst-idle.ts b/__test__/bench-copy-burst-idle.ts new file mode 100644 index 0000000..b0cb2c2 --- /dev/null +++ b/__test__/bench-copy-burst-idle.ts @@ -0,0 +1,161 @@ +/** + * Repeated burst/idle cycles per connection — the workload the acceptance criteria describes + * (MQTT fan-out: connections that burst, go quiet, then burst again) and the one shape neither + * bench-copy-memory (one burst, then quiet forever) nor bench-copy-throughput (one sustained + * burst) exercises. Reviewer ask on PR #41: the design's escalate -> park -> shrink -> + * re-escalate mechanism costs two allocations plus a zeroing `vec![0u8; n]` per direction per + * cycle (src/copy.rs) — this is the only benchmark where that churn sits on the hot path, so + * it's the one that can actually show whether the cost is as negligible as the module docs + * argue, by reporting both throughput and RSS across many repeated cycles instead of a single + * before/after snapshot. + * + * Every connection, each cycle: write a `readBufferSize`-sized burst, wait for the full echo + * back, then sleep `idleMs` — long enough for both directions to actually park and shrink (see + * src/copy.rs's shrink-on-park design) before the next cycle's burst starts. All connections run + * each cycle in lockstep so RSS can be sampled between cycles: a flat curve across cycles means + * the per-cycle churn isn't accumulating; a climbing one would mean it is. + * + * Not part of `npm test` — a manual measurement tool. Compare by running it against the base + * commit (`tokio::io::copy_bidirectional_with_sizes`) and this branch. + * + * Run with: + * npm run build:debug + * node --expose-gc dist-test/__test__/bench-copy-burst-idle.js [connections] [readBufferSize] [cycles] [idleMs] [lazyCopyBufferThreshold] + */ +import * as tls from 'node:tls'; +import { SymphonyProxy } from '../ts/proxy.js'; +import { generateSelfSignedCert, getFreePort, startEchoServer, sleep } from './util.js'; + +const CONNECTIONS = Number(process.argv[2] ?? 2000); +const READ_BUFFER_SIZE = Number(process.argv[3] ?? 65536); +const CYCLES = Number(process.argv[4] ?? 15); +const IDLE_MS = Number(process.argv[5] ?? 250); +// Escalation is gated on active connections (see src/copy.rs LazyBufferGate). Pinned to 0 by +// default so the benchmark measures the escalating path whatever connection count it is given; +// pass a value above `connections` to measure the static path for comparison. +const LAZY_THRESHOLD = Number(process.argv[6] ?? 0); +const CONNECT_BATCH = 250; + +// See bench-copy-memory.ts: spreading client sockets across several loopback source addresses +// avoids exhausting the ~28k usable ephemeral ports on a single source address. +const SOURCE_ADDRESSES = Array.from({ length: 16 }, (_, i) => `127.0.0.${i + 1}`); + +function rssMb(): number { + if (global.gc) global.gc(); + return process.memoryUsage().rss / (1024 * 1024); +} + +async function main() { + const cert = generateSelfSignedCert('localhost'); + // A real echo, not a sink: each cycle's burst has to actually round-trip so the connection + // goes genuinely idle (nothing left to write on either side) before the next cycle — that's + // the "park" condition src/copy.rs shrinks on, and the whole point of this benchmark is to + // exercise that transition repeatedly rather than once. + const upstream = await startEchoServer(); + + const proxyPort = await getFreePort(); + const proxy = new SymphonyProxy({ + listeners: [{ host: '127.0.0.1', port: proxyPort }], + routes: [ + { + sni: 'localhost', + terminateTls: true, + cert: { certChain: cert.cert, privateKey: cert.key }, + upstreams: [{ kind: 'tcp', host: '127.0.0.1', port: upstream.port }], + }, + ], + readBufferSize: READ_BUFFER_SIZE, + lazyCopyBufferThreshold: LAZY_THRESHOLD, + }); + await proxy.start(); + + console.log(`connections=${CONNECTIONS} readBufferSize=${READ_BUFFER_SIZE}B cycles=${CYCLES} idleMs=${IDLE_MS}`); + + const sockets: tls.TLSSocket[] = []; + for (let i = 0; i < CONNECTIONS; i += CONNECT_BATCH) { + const n = Math.min(CONNECT_BATCH, CONNECTIONS - i); + const batchStart = Date.now(); + await Promise.all( + Array.from({ length: n }, (_, j) => new Promise((resolve, reject) => { + const localAddress = SOURCE_ADDRESSES[(i + j) % SOURCE_ADDRESSES.length]; + const s = tls.connect( + { port: proxyPort, host: '127.0.0.1', servername: 'localhost', rejectUnauthorized: false, localAddress } as tls.ConnectionOptions, + () => resolve(), + ); + s.on('error', reject); + sockets.push(s); + })), + ); + console.log(` connected ${i + n}/${CONNECTIONS} (batch took ${Date.now() - batchStart}ms)`); + } + console.log(`connected ${sockets.length} connections`); + + await sleep(200); + const baselineMb = rssMb(); + console.log(`baseline RSS (connected, before first burst): ${baselineMb.toFixed(1)} MiB`); + + const burst = Buffer.alloc(READ_BUFFER_SIZE, 7); + let totalBytes = 0; + const start = Date.now(); + const rssSamples: number[] = []; + + for (let cycle = 0; cycle < CYCLES; cycle++) { + await Promise.all( + sockets.map((s) => new Promise((resolve) => { + let received = 0; + const cleanup = () => { + s.off('data', onData); + s.off('error', onDone); + s.off('close', onDone); + resolve(); + }; + const onData = (chunk: Buffer) => { + received += chunk.length; + if (received >= burst.length) cleanup(); + }; + // A socket that errors or closes mid-cycle (loopback flakiness at thousands of + // connections) must still resolve this promise — otherwise that one socket's + // `Promise.all` never settles and the whole run wedges with no diagnostic. + const onDone = () => cleanup(); + s.on('data', onData); + s.once('error', onDone); + s.once('close', onDone); + s.write(burst); + })), + ); + totalBytes += burst.length * sockets.length * 2; // client->upstream and upstream->client + + // Long enough for both copy directions to actually park (see src/copy.rs's + // shrink-on-park design) before the next cycle's burst — this is what makes the + // escalate/shrink churn under test happen at all; too short a gap would just look like + // one continuing sustained burst and never trigger a shrink. + await sleep(IDLE_MS); + const sampleMb = rssMb(); + rssSamples.push(sampleMb); + console.log(` cycle ${cycle + 1}/${CYCLES}: RSS ${sampleMb.toFixed(1)} MiB`); + } + + const elapsedS = (Date.now() - start) / 1000; + const mibps = totalBytes / (1024 * 1024) / elapsedS; + const minRss = Math.min(...rssSamples); + const maxRss = Math.max(...rssSamples); + + console.log( + `throughput: ${mibps.toFixed(1)} MiB/s over ${elapsedS.toFixed(1)}s ` + + `(${CYCLES} burst/idle cycles x ${sockets.length} connections, both directions)`, + ); + console.log( + `RSS across cycles: min ${minRss.toFixed(1)} MiB, max ${maxRss.toFixed(1)} MiB, ` + + `baseline ${baselineMb.toFixed(1)} MiB, peak delta from baseline ${(maxRss - baselineMb).toFixed(1)} MiB`, + ); + console.log('A flat RSS curve across cycles (no upward trend from the first sample to the last) means the escalate/park/shrink churn is not accumulating.'); + + // See bench-copy-memory.ts: a graceful teardown of thousands of sockets is slow and + // unnecessary for a one-shot measurement process. + process.exit(0); +} + +main().catch((err) => { + console.error(err); + process.exit(1); +}); diff --git a/__test__/bench-copy-memory.ts b/__test__/bench-copy-memory.ts new file mode 100644 index 0000000..2d2e004 --- /dev/null +++ b/__test__/bench-copy-memory.ts @@ -0,0 +1,169 @@ +/** + * Measures process RSS growth from parked connections — the memory curve issue #37 is about. + * `tokio::io::copy_bidirectional`'s `CopyBuffer` allocates once and holds its buffer for the + * connection's whole life, whether or not it is transferring; the fix (src/copy.rs) holds only a + * small fixed floor while idle and escalates to the full `readBufferSize` only for the duration + * of an actual sustained burst, so a parked connection (the MQTT shape: idle between publishes) + * should cost a small fixed amount instead of `readBufferSize × 2` forever. + * + * Every connection sends one burst sized to `readBufferSize` (an MQTT PUBLISH shape: some + * payload at least that large, at some point) before settling back to parked. This matters: a + * read below the buffer's capacity only dirties the one page it actually writes into, so a + * static buffer above the allocator's mmap threshold (glibc: 128 KiB) mostly sits on + * lazily-faulted, never-touched pages regardless of whether the code holds it eagerly — that + * would understate the OLD code's cost and make old-vs-new meaningless. Forcing a real burst + * through is what makes the OLD static buffer become fully resident (and stay that way for the + * connection's life) while the NEW buffer escalates for the burst and drops back down once + * traffic quiets — demonstrating the actual point of the fix: resident memory should track + * *peak concurrent transfers*, not total connection count or the configured maximum. + * + * TLS termination is used (matching __test__/copy-buffers.spec.ts) with an echoing TCP upstream + * — the client's real TLS handshake terminates at symphony. rustls session state is an + * unrelated, unchanged memory cost; the before/after delta on the same rig isolates the + * copy-buffer question regardless. + * + * Not part of `npm test` — a manual measurement tool. Compare by running it against the base + * commit (`tokio::io::copy_bidirectional_with_sizes`) and this branch; see the PR description + * for the actual before/after numbers. + * + * Run with: + * npm run build:debug + * node --expose-gc dist-test/__test__/bench-copy-memory.js [connectionCount] [readBufferSize] [lazyCopyBufferThreshold] + */ +import * as tls from 'node:tls'; +import { SymphonyProxy } from '../ts/proxy.js'; +import { generateSelfSignedCert, getFreePort, startEchoServer, sleep } from './util.js'; + +const CONNECTIONS = Number(process.argv[2] ?? 30000); +// Deliberately large: the whole point of the fix is that a big configured buffer no longer +// means every idle connection pays for it. A small default would hide a regression. +const READ_BUFFER_SIZE = Number(process.argv[3] ?? 65536); +// Escalation is gated on active connections (src/copy.rs LazyBufferGate). 0 forces the escalating +// path whatever the connection count; pass a value above `connectionCount` to measure the static +// path — that pair is what isolates what the gate is actually buying. +const LAZY_THRESHOLD = Number(process.argv[4] ?? 0); +const BATCH = 250; + +// A single loopback source address only has ~28k usable ephemeral ports +// (net.ipv4.ip_local_port_range), well under the connection counts this benchmark needs. +// Loopback supports the whole 127.0.0.0/8 range, so spreading client sockets across several +// source addresses multiplies the available (srcIP, srcPort) tuples instead. +const SOURCE_ADDRESSES = Array.from({ length: 16 }, (_, i) => `127.0.0.${i + 1}`); + +function rssMb(): number { + if (global.gc) global.gc(); + return process.memoryUsage().rss / (1024 * 1024); +} + +async function main() { + const cert = generateSelfSignedCert('localhost'); + // Idle upstream: accepts and holds the connection, never sends or expects data — the + // "parked MQTT subscriber between publishes" shape this issue is about. + const upstream = await startEchoServer(); + + const proxyPort = await getFreePort(); + const proxy = new SymphonyProxy({ + listeners: [{ host: '127.0.0.1', port: proxyPort }], + routes: [ + { + sni: 'localhost', + terminateTls: true, + cert: { certChain: cert.cert, privateKey: cert.key }, + upstreams: [{ kind: 'tcp', host: '127.0.0.1', port: upstream.port }], + }, + ], + readBufferSize: READ_BUFFER_SIZE, + lazyCopyBufferThreshold: LAZY_THRESHOLD, + }); + await proxy.start(); + + await sleep(200); + const baselineMb = rssMb(); + console.log(`readBufferSize=${READ_BUFFER_SIZE}B/direction (${READ_BUFFER_SIZE * 2}B/connection if held permanently)`); + console.log(`baseline RSS (proxy + upstream server, 0 connections): ${baselineMb.toFixed(1)} MiB`); + + const sockets: tls.TLSSocket[] = []; + for (let i = 0; i < CONNECTIONS; i += BATCH) { + const n = Math.min(BATCH, CONNECTIONS - i); + const batchStart = Date.now(); + await Promise.all( + Array.from({ length: n }, (_, j) => new Promise((resolve, reject) => { + const localAddress = SOURCE_ADDRESSES[(i + j) % SOURCE_ADDRESSES.length]; + const s = tls.connect( + { port: proxyPort, host: '127.0.0.1', servername: 'localhost', rejectUnauthorized: false, localAddress } as tls.ConnectionOptions, + () => resolve(), + ); + s.on('error', reject); + // Keep flowing so the echoed burst (see below) actually drains instead of piling + // up unread — a paused socket has no bearing on symphony's own memory, but an + // unread echo sitting in the kernel receive buffer would stall the proxy's + // upstream→client write, which would in turn distort the very thing under test. + s.on('data', () => {}); + sockets.push(s); + })), + ); + console.log(` connected ${i + n}/${CONNECTIONS} (batch took ${Date.now() - batchStart}ms)`); + } + console.log(`connected ${sockets.length} parked connections`); + + // One burst per connection sized to the configured buffer, fired in batches (no + // per-connection acknowledgment tracking — with tens of thousands of sockets in one process, + // thousands of individual listeners/timers is itself a source of event-loop pressure separate + // from anything under test). This has to be at least `readBufferSize` bytes, not a single + // byte: a read below the buffer's capacity only ever dirties the one page it actually writes + // into, so a static buffer well above the allocator's mmap threshold (glibc: 128 KiB) mostly + // sits on lazily-faulted, never-touched pages regardless of whether the code holds it + // eagerly or not — that would understate the OLD code's cost and make the comparison + // meaningless. Sized to the buffer, the burst forces the OLD static buffer to become fully + // resident (matching a real client that at some point publishes a payload at least that + // large), while the NEW code escalates for the burst and then drops back down once the + // connection goes quiet again. A fixed settle period afterward gives the whole batch time to + // round-trip through the echo upstream — see the file header for why this step (not just + // connecting) is what makes the OLD/NEW difference observable in RSS. A handful of + // connections not completing their round trip in time doesn't materially change a delta + // computed across tens of thousands of connections. + const burst = Buffer.alloc(READ_BUFFER_SIZE, 7); + // Small on purpose: each write is up to `readBufferSize` (possibly 1 MiB), and this loop + // waits for the kernel/TLS layer to actually accept it (backpressure-aware) before moving on + // — a fire-and-forget write here would pile the unsent bytes up in *this test process's own* + // write buffers, dwarfing anything happening on the symphony side and making the measurement + // meaningless. + const PING_BATCH = 200; + let stalled = 0; + for (let i = 0; i < sockets.length; i += PING_BATCH) { + await Promise.all( + sockets.slice(i, i + PING_BATCH).map((s) => new Promise((resolve) => { + // A per-write timeout backstop: at tens of thousands of sockets in one test + // process a handful occasionally don't drain promptly (event-loop/socket + // pressure in the test harness itself, reproduced identically against the base + // commit — not a symphony behavior). Moving on rather than blocking the whole + // batch keeps that noise from taking down the entire run. + const timer = setTimeout(() => { stalled++; resolve(); }, 3000); + if (s.write(burst)) { clearTimeout(timer); resolve(); } + else s.once('drain', () => { clearTimeout(timer); resolve(); }); + })), + ); + if ((i / PING_BATCH) % 10 === 0) console.log(` burst-sent ${i + PING_BATCH}/${sockets.length}`); + } + console.log(`sent ${sockets.length} ${READ_BUFFER_SIZE}-byte bursts (${stalled} stalled past 3s and were skipped)`); + + // Let accept/handshake/ping bookkeeping settle so we're measuring steady-state idle, not + // transient allocations. + await sleep(5000); + const loadedMb = rssMb(); + const deltaMb = loadedMb - baselineMb; + const perConnBytes = (deltaMb * 1024 * 1024) / CONNECTIONS; + + console.log(`loaded RSS (${CONNECTIONS} parked connections): ${loadedMb.toFixed(1)} MiB`); + console.log(`delta: ${deltaMb.toFixed(1)} MiB total, ${perConnBytes.toFixed(0)} bytes/connection`); + + // The numbers we care about are already printed; at tens of thousands of sockets a graceful + // teardown (destroying each socket, waiting for the proxy/upstream to notice) is slow and + // unnecessary for a one-shot measurement process, so exit immediately rather than hang here. + process.exit(0); +} + +main().catch((err) => { + console.error(err); + process.exit(1); +}); diff --git a/__test__/bench-copy-throughput.ts b/__test__/bench-copy-throughput.ts new file mode 100644 index 0000000..4f27ca1 --- /dev/null +++ b/__test__/bench-copy-throughput.ts @@ -0,0 +1,101 @@ +/** + * Bulk throughput check for issue #37: a few connections pushing sustained high-volume data + * through the copy loop — the replication profile on port 9933 (few connections, high volume), + * the opposite shape from the many-parked-MQTT-subscribers case the memory fix targets. The + * lazy/released buffer in src/copy.rs must not cost throughput on this path relative to + * `tokio::io::copy_bidirectional_with_sizes`. + * + * TLS termination (matching __test__/copy-buffers.spec.ts) with a plain TCP sink upstream — + * client→upstream is the direction under measurement. + * + * Not part of `npm test` — a manual measurement tool. Compare by running it against the base + * commit and this branch; see the PR description for the actual before/after numbers. + * + * Run with: + * npm run build:debug + * node dist-test/__test__/bench-copy-throughput.js [connections] [durationMs] + */ +import * as net from 'node:net'; +import * as tls from 'node:tls'; +import { SymphonyProxy } from '../ts/proxy.js'; +import { generateSelfSignedCert, getFreePort } from './util.js'; + +const CONNECTIONS = Number(process.argv[2] ?? 4); +const DURATION_MS = Number(process.argv[3] ?? 8000); +const CHUNK_SIZE = 256 * 1024; + +async function main() { + const cert = generateSelfSignedCert('localhost'); + let totalReceived = 0; + const upstreamPort = await getFreePort(); + // Sink: reads and discards, so the client→upstream direction is the one under measurement + // (matches a replication log-shipping / bulk-write shape). + const upstream = net.createServer((socket) => { + socket.on('data', (chunk: Buffer) => { + totalReceived += chunk.length; + }); + socket.on('error', () => {}); + }); + await new Promise((resolve) => upstream.listen(upstreamPort, '127.0.0.1', resolve)); + + const proxyPort = await getFreePort(); + const proxy = new SymphonyProxy({ + listeners: [{ host: '127.0.0.1', port: proxyPort }], + routes: [ + { + sni: 'localhost', + terminateTls: true, + cert: { certChain: cert.cert, privateKey: cert.key }, + upstreams: [{ kind: 'tcp', host: '127.0.0.1', port: upstreamPort }], + }, + ], + // This benchmark runs a handful of sustained streams, which is well below the default + // `lazyCopyBufferThreshold` — so without pinning it here it would measure the static path + // and report nothing about the escalating buffers it exists to check for regressions in. + lazyCopyBufferThreshold: 0, + }); + await proxy.start(); + + const chunk = Buffer.alloc(CHUNK_SIZE, 7); + const sockets: tls.TLSSocket[] = []; + await Promise.all( + Array.from({ length: CONNECTIONS }, () => new Promise((resolve, reject) => { + const s = tls.connect( + { port: proxyPort, host: '127.0.0.1', servername: 'localhost', rejectUnauthorized: false }, + () => resolve(), + ); + s.on('error', reject); + sockets.push(s); + })), + ); + + function pump(socket: tls.TLSSocket) { + function write() { + let ok = true; + while (ok) ok = socket.write(chunk); + } + socket.on('drain', write); + write(); + } + + const start = Date.now(); + for (const s of sockets) pump(s); + await new Promise((r) => setTimeout(r, DURATION_MS)); + const elapsedS = (Date.now() - start) / 1000; + + const mibps = totalReceived / (1024 * 1024) / elapsedS; + console.log( + `connections=${CONNECTIONS} duration=${elapsedS.toFixed(1)}s ` + + `received=${(totalReceived / 1024 / 1024).toFixed(1)}MiB throughput=${mibps.toFixed(1)} MiB/s`, + ); + + for (const s of sockets) s.destroy(); + await proxy.stop(1000).catch(() => {}); + await new Promise((resolve) => upstream.close(() => resolve())); + process.exit(0); +} + +main().catch((err) => { + console.error(err); + process.exit(1); +}); diff --git a/__test__/copy-buffers.spec.ts b/__test__/copy-buffers.spec.ts index 9810ee4..ca37049 100644 --- a/__test__/copy-buffers.spec.ts +++ b/__test__/copy-buffers.spec.ts @@ -44,7 +44,7 @@ describe('copy buffers', () => { async function startProxy( buffers: Pick< ProxyConfig, - 'readBufferSize' | 'clientReadBufferSize' | 'upstreamReadBufferSize' + 'readBufferSize' | 'clientReadBufferSize' | 'upstreamReadBufferSize' | 'lazyCopyBufferThreshold' >, ): Promise { const port = await getFreePort(); @@ -99,4 +99,26 @@ describe('copy buffers', () => { const port = await startProxy({ readBufferSize: 2048, upstreamReadBufferSize: 16384 }); await assertRoundTrip(port, patterned(LARGE_PAYLOAD)); }); + + // The gate picks the buffer strategy but must never change what comes out the other end. + // Both sides of the threshold are exercised on the real addon: a single test connection is + // always below any positive threshold, so `0` is the only way to reach the escalating path + // from here, and a high value is the only way to reach the static one. + + it('round-trips with escalating buffers forced on (threshold 0)', async () => { + const port = await startProxy({ readBufferSize: 4096, lazyCopyBufferThreshold: 0 }); + await assertRoundTrip(port, patterned(LARGE_PAYLOAD)); + }); + + it('round-trips with escalating buffers disabled (threshold above peak concurrency)', async () => { + const port = await startProxy({ readBufferSize: 4096, lazyCopyBufferThreshold: 1_000_000 }); + await assertRoundTrip(port, patterned(LARGE_PAYLOAD)); + }); + + it('round-trips a payload far larger than the buffer with the gate disengaged', async () => { + // The static path has to keep working for the large-payload case too: with no escalation + // the buffer never grows, so every one of these bytes crosses in 512-byte reads. + const port = await startProxy({ readBufferSize: 512, lazyCopyBufferThreshold: 1_000_000 }); + await assertRoundTrip(port, patterned(LARGE_PAYLOAD)); + }); }); diff --git a/__test__/server.spec.ts b/__test__/server.spec.ts index 165e15e..22b6325 100644 --- a/__test__/server.spec.ts +++ b/__test__/server.spec.ts @@ -14,6 +14,7 @@ import { spawn, ChildProcess } from 'node:child_process'; import * as fs from 'node:fs'; import * as os from 'node:os'; import * as path from 'node:path'; +import * as tls from 'node:tls'; import { generateSelfSignedCert, getFreePort, startEchoServer, tlsRoundTrip, sleep } from './util.js'; // server.js sits next to this compiled spec's sibling ts/ dir: dist-test/ts/server.js @@ -41,20 +42,24 @@ async function waitFor(predicate: () => boolean | Promise, timeoutMs = interface RunningServer { child: ChildProcess; getStderr: () => string; + /** stdout carries the server's own lifecycle log lines (log(), not logErr()). */ + getStdout: () => string; markShutdown: () => void; } function spawnServer(configPath: string, statusPath?: string): RunningServer { let stderr = ''; + let stdout = ''; let shuttingDown = false; const args = [SERVER_JS, '--config', configPath]; if (statusPath) args.push('--status', statusPath); const child = spawn(process.execPath, args, { stdio: ['ignore', 'pipe', 'pipe'] }); child.stderr?.on('data', (d) => (stderr += d.toString())); + child.stdout?.on('data', (d) => (stdout += d.toString())); child.on('exit', (code, sig) => { if (!shuttingDown) stderr += `\n[child exited early code=${code} sig=${sig}]`; }); - return { child, getStderr: () => stderr, markShutdown: () => (shuttingDown = true) }; + return { child, getStderr: () => stderr, getStdout: () => stdout, markShutdown: () => (shuttingDown = true) }; } async function killServer(server: RunningServer): Promise { @@ -705,3 +710,142 @@ describe('symphony-server (status.json ownership guard)', () => { ); }); }); + +describe('symphony-server (construction-frozen proxy fields force a recreate)', () => { + // readBufferSize, its two per-direction overrides, and lazyCopyBufferThreshold are all frozen + // in SymphonyProxyWrap at construction — updateConfig() reaches none of them. If they were + // missing from the reconcile's construction signature, editing one would leave the signature + // unchanged, take the route-only hot-swap branch, and report a successful reload while the + // proxy kept running the old value. Silent, and only ever visible as "the setting we shipped + // didn't do anything". + // + // The observable is the server's own "proxy listening on ports" line, which the recreate + // branch emits and the hot-swap branch does not. Note it is NOT connection loss: stop() ends + // the accept loops but in-flight connection tasks run to completion, so established + // connections survive a recreate and keep the old buffer sizes — which is exactly why the + // README calls a buffer-size edit a reconnect event. + // + // The route-only case is the control. Without it this would pass equally well against a + // server that recreated on every config write, which would prove nothing about the signature. + const cert = generateSelfSignedCert('localhost'); + let dir: string; + let configPath: string; + let statusPath: string; + let proxyPort: number; + let echo: Awaited>; + let server: RunningServer; + + const LISTENING_LINE = /proxy listening on ports/g; + const listenCount = () => (server.getStdout().match(LISTENING_LINE) ?? []).length; + + const baseConfig = (extra: Record, routeSnis: string[] = ['localhost']) => ({ + version: 1, + proxies: [ + { + listeners: [{ host: '127.0.0.1', port: proxyPort }], + routes: routeSnis.map((sni) => ({ + sni, + upstreams: [{ kind: 'tcp', host: '127.0.0.1', port: echo.port }], + terminateTls: true, + cert: { certChain: cert.cert, privateKey: cert.key }, + })), + ...extra, + }, + ], + }); + + before(async () => { + dir = fs.mkdtempSync(path.join(os.tmpdir(), 'symphony-recreate-')); + configPath = path.join(dir, 'config.json'); + statusPath = path.join(dir, 'status.json'); + echo = await startEchoServer(); + proxyPort = await getFreePort(); + writeConfigAtomic(configPath, baseConfig({ readBufferSize: 4096, lazyCopyBufferThreshold: 0 })); + server = spawnServer(configPath); + await waitFor(() => fs.existsSync(statusPath)); + await waitFor(() => listenCount() >= 1); + }); + + after(async () => { + await killServer(server); + await echo.close().catch(() => {}); + fs.rmSync(dir, { recursive: true, force: true }); + }); + + it('recreates the proxy when lazyCopyBufferThreshold changes', async () => { + const before = listenCount(); + writeConfigAtomic(configPath, baseConfig({ readBufferSize: 4096, lazyCopyBufferThreshold: 5000 })); + await waitFor(() => listenCount() > before, 8000, 100); + + // Healthy on the new value, not merely torn down and rebuilt into nothing. + const fresh = await tlsRoundTrip({ port: proxyPort, servername: 'localhost', caCert: cert.cert, data: Buffer.from('after-threshold'), rejectUnauthorized: false }); + assert.equal(fresh.toString(), 'after-threshold'); + }); + + it('recreates the proxy when readBufferSize changes', async () => { + const before = listenCount(); + writeConfigAtomic(configPath, baseConfig({ readBufferSize: 16384, lazyCopyBufferThreshold: 5000 })); + await waitFor(() => listenCount() > before, 8000, 100); + + const fresh = await tlsRoundTrip({ port: proxyPort, servername: 'localhost', caCert: cert.cert, data: Buffer.from('after-bufsize'), rejectUnauthorized: false }); + assert.equal(fresh.toString(), 'after-bufsize'); + }); + + it('leaves established connections running on the old proxy across a recreate', async () => { + // Documents what a recreate actually does to in-flight sessions, which is NOT what the + // mechanism suggests at a glance: stop() sends the shutdown broadcast (ending the accept + // loops) and sleeps 100ms, but it never aborts connection tasks, and the tokio runtime + // lives inside the napi wrap until JS garbage-collects it. So established sessions keep + // running — on the OLD buffer settings — rather than being dropped. + // + // This matters most for exactly the deployment the setting targets: long-lived MQTT + // subscribers would keep their old buffers indefinitely after an operator lowered the + // value to reclaim memory. + const held = await new Promise((resolve, reject) => { + const s = tls.connect( + { port: proxyPort, host: '127.0.0.1', servername: 'localhost', ca: cert.cert, rejectUnauthorized: false }, + () => resolve(s), + ); + s.on('error', reject); + }); + const echoOn = (payload: string) => + new Promise((resolve, reject) => { + const t = setTimeout(() => reject(new Error('no echo on held connection')), 3000); + held.once('data', (d: Buffer) => { + clearTimeout(t); + resolve(d.toString()); + }); + held.write(payload); + }); + + assert.equal(await echoOn('pre-recreate'), 'pre-recreate'); + + const before = listenCount(); + writeConfigAtomic(configPath, baseConfig({ readBufferSize: 32768, lazyCopyBufferThreshold: 5000 })); + await waitFor(() => listenCount() > before, 8000, 100); + + assert.equal(held.destroyed, false, 'the held connection must survive the recreate'); + assert.equal(await echoOn('post-recreate'), 'post-recreate', 'and must still proxy on the old proxy'); + held.destroy(); + }); + + it('hot-swaps instead of recreating for a route-only change (control)', async () => { + const before = listenCount(); + // Same listeners, same proxy-level fields — only the route table grows, which is what the + // hot-swap path exists for. + writeConfigAtomic(configPath, baseConfig({ readBufferSize: 16384, lazyCopyBufferThreshold: 5000 }, ['localhost', 'other.localhost'])); + + // The added route proving the reload really applied — otherwise "no recreate" would also + // be satisfied by the server having ignored the edit entirely. + await waitFor(async () => { + try { + const r = await tlsRoundTrip({ port: proxyPort, servername: 'other.localhost', caCert: cert.cert, data: Buffer.from('hi'), rejectUnauthorized: false }); + return r.toString() === 'hi'; + } catch { + return false; + } + }, 8000, 100); + + assert.equal(listenCount(), before, `a route-only edit must hot-swap, not recreate. stdout:\n${server.getStdout()}`); + }); +}); diff --git a/package-lock.json b/package-lock.json index c76cca0..8a90772 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,17 +1,20 @@ { "name": "@harperfast/symphony", - "version": "0.1.0", + "version": "1.1.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@harperfast/symphony", - "version": "0.1.0", + "version": "1.1.0", "license": "Elastic-2.0", "os": [ "linux", "darwin" ], + "bin": { + "symphony-server": "dist/server.js" + }, "devDependencies": { "@harperfast/integration-testing": "^0.2.0", "@napi-rs/cli": "^2.18.0", diff --git a/package.json b/package.json index 532dcb6..e5b7970 100644 --- a/package.json +++ b/package.json @@ -35,6 +35,9 @@ "test:integration": "tsc -p tsconfig.test.json && node --test --test-force-exit dist-test/__test__/harper-integration.spec.js", "benchmark": "tsc -p tsconfig.test.json && node dist-test/__test__/benchmark.js", "benchmark:throughput": "tsc -p tsconfig.test.json && node dist-test/__test__/benchmark-throughput.js", + "benchmark:copy-memory": "tsc -p tsconfig.test.json && node --expose-gc dist-test/__test__/bench-copy-memory.js", + "benchmark:copy-throughput": "tsc -p tsconfig.test.json && node dist-test/__test__/bench-copy-throughput.js", + "benchmark:copy-burst-idle": "tsc -p tsconfig.test.json && node --expose-gc dist-test/__test__/bench-copy-burst-idle.js", "diagnostic:keepalive": "tsc -p tsconfig.test.json && node dist-test/__test__/diagnostic-uds-keepalive.js" }, "devDependencies": { diff --git a/src/copy.rs b/src/copy.rs new file mode 100644 index 0000000..b3025c5 --- /dev/null +++ b/src/copy.rs @@ -0,0 +1,1039 @@ +//! Hand-rolled replacement for `tokio::io::copy_bidirectional_with_sizes` that does not hold a +//! per-direction buffer for the connection's whole life. +//! +//! `tokio::io::CopyBuffer` allocates its buffer once, before the first read, and keeps it for as +//! long as the copy future exists — i.e. for the whole proxied session, whether or not it is +//! actively transferring. For a mostly-idle connection (an MQTT subscriber parked between +//! publishes) that is dead weight: `readBufferSize` × 2 held in memory for a session that spends +//! nearly all its time doing nothing. `LazyCopyBuffer` below holds only a small fixed floor +//! (`PROBE_BUFFER_SIZE`) while idle or exchanging small discrete messages, escalating straight to +//! the full `max_buf_size` only once *two consecutive* reads land at capacity, and dropping +//! straight back to the floor once the direction actually parks with nothing left to write — not +//! on every single under-capacity read, which would shrink (and then immediately re-grow) a +//! connection that is still continuously active but simply has variably-sized traffic. +//! `readBufferSize` (and its per-direction overrides) becomes the *maximum* per-transfer buffer +//! size rather than a permanent allocation. +//! +//! None of this is unconditional. The saving scales with connection count and the cost does not, +//! so the escalate/shrink behavior is gated on the proxy's live active-connection count +//! (`LazyBufferGate`, configured by `lazyCopyBufferThreshold`). Below the threshold each direction +//! gets its full configured buffer once and never resizes — byte-for-byte the +//! `tokio::io::copy_bidirectional_with_sizes` behavior this module replaced — so a proxy carrying +//! a few bulk replication streams pays nothing for a memory problem it does not have. Above it, +//! everything below applies. +//! +//! Why two consecutive full reads, not one: escalating on a single full read means a message +//! that happens to exactly fill the current (small) buffer — coincidence, not evidence of a +//! burst — jumps straight to the configured maximum and then sits there for however long the +//! connection is next idle, which can be indefinite. Requiring a second consecutive full read +//! before jumping costs one extra small-buffer round trip on every genuine burst (negligible) and +//! bounds that single-message coincidence to the floor size instead of the configured maximum, +//! however large that maximum is configured. +//! +//! This is structured as a direct port of `tokio::io::util::copy::CopyBuffer` and +//! `copy_bidirectional`'s `transfer_one_direction`/`TransferState` (see the tokio source), with +//! the buffer-resize decision spliced into the point where a fully-drained buffer is reset for +//! the next read. Reusing that proven state machine — rather than `tokio::io::split` plus +//! independent per-direction `async fn`s — matters for two reasons: `split` wraps each side in an +//! `Arc>` (two heap allocations and a lock/unlock on every read/write/flush/shutdown, +//! *per connection*, which itself scales with connection count — exactly what this change exists +//! to avoid), and a hand-rolled `async fn` pump does not flush before parking on the next read the +//! way tokio's `poll_copy` does. That gap is real: `write_all` only guarantees the data reached +//! the writer's internal buffer, not the wire (`tokio-rustls` in particular defers sending +//! encrypted records until `poll_flush`). A TLS client that sends one request and then waits for +//! the reply would never see it — both sides would sit idle until `idle_timeout` cleaned up the +//! session — because the pump had already moved on to (blocking on) the next read before the +//! previous write was actually flushed. + +use std::future::{poll_fn, Future}; +use std::io; +use std::pin::Pin; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::task::{ready, Context, Poll}; +use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; + +use crate::metrics::GlobalMetrics; + +/// Decides, at each resize point, whether escalating buffers are worth their cost *right now*. +/// +/// The saving scales with connection count; the cost does not. Growing and releasing a buffer +/// costs two allocations plus the zeroing of `vec![0u8; n]` per burst, and a proxy carrying four +/// bulk replication streams pays that on every burst while saving a few hundred KiB it was never +/// short of. The same behavior across a hundred thousand parked MQTT subscribers is the whole +/// point of the module. So the behavior is gated on how busy the proxy actually is rather than +/// chosen once, fleet-wide, by whoever wrote the config. +/// +/// `threshold` is the proxy-wide active connection count at or above which escalation engages. +/// `0` engages it always; a value above peak concurrency disables it, leaving each direction on a +/// full-size buffer allocated once — byte-for-byte the `tokio::io::copy_bidirectional_with_sizes` +/// behavior this module replaced, with no resize churn at all. +/// +/// Deliberately re-read at every resize point rather than latched per connection. A connection +/// established while the proxy was quiet would otherwise hold a full-size buffer for its entire +/// life however busy the proxy later became — and long-lived connections accumulating while idle +/// is exactly the shape that motivates this. Re-reading means those connections start releasing +/// their buffers at their next park once the proxy crosses the threshold. +#[derive(Clone)] +pub struct LazyBufferGate { + metrics: Arc, + threshold: u64, +} + +impl LazyBufferGate { + pub fn new(metrics: Arc, threshold: u64) -> Self { + Self { metrics, threshold } + } + + /// `Relaxed` matches how the gauge is maintained and is all this needs: the result only picks + /// a buffer size, so a read that is momentarily stale costs one connection one resize + /// decision, never correctness. + fn engaged(&self) -> bool { + self.threshold == 0 || self.metrics.active_connections.load(Ordering::Relaxed) >= self.threshold + } +} + +#[cfg(test)] +impl LazyBufferGate { + /// Engaged regardless of connection count — the escalate/release behavior under test. + pub fn always() -> Self { + Self::new(Arc::new(GlobalMetrics::default()), 0) + } + + /// Never engaged — one full-size buffer per direction, held for the connection's life. + pub fn never() -> Self { + Self::new(Arc::new(GlobalMetrics::default()), u64::MAX) + } + + /// A gate over a counter the test drives directly, for the threshold-crossing cases. + pub fn with_metrics(metrics: Arc, threshold: u64) -> Self { + Self::new(metrics, threshold) + } +} + +/// Resident buffer size while a direction is idle or only exchanging small, discrete messages +/// (PINGREQ, a short request). Deliberately tiny and fixed, not zero: a zero-length read buffer +/// makes `Ok(0)` ambiguous with EOF (see `proxy_conn::MIN_COPY_BUFFER_SIZE`), and a small constant +/// floor is a trivial, bounded cost — 1 KiB/connection total across both directions — next to a +/// configured max that can run up to 2 MiB/connection if held permanently the way +/// `tokio::io::CopyBuffer` holds it. +const PROBE_BUFFER_SIZE: usize = 512; + +/// Copies bytes in both directions between `client` and `upstream`, mirroring +/// `tokio::io::copy_bidirectional_with_sizes`'s observable behavior: +/// - EOF on one reader shuts down the corresponding writer on the other stream ("half-close"); +/// the other direction keeps running until it too reaches EOF or errors. +/// - An error on *either* direction — including a failed `shutdown()`, which tokio also +/// propagates — ends the whole copy immediately rather than waiting for the other direction. +/// This is what makes an RST on one leg end the session instead of hanging until idle_timeout. +/// +/// `client_buf_size`/`upstream_buf_size` bound the buffer each direction may grow to; whether +/// either is held at that size permanently depends on `gate` (see `LazyBufferGate` and the module +/// docs). +pub async fn copy_bidirectional_lazy( + client: &mut C, + upstream: &mut U, + client_buf_size: usize, + upstream_buf_size: usize, + gate: LazyBufferGate, +) -> io::Result<()> +where + C: AsyncRead + AsyncWrite + Unpin, + U: AsyncRead + AsyncWrite + Unpin, +{ + let mut client_to_upstream = TransferState::Running(LazyCopyBuffer::new(client_buf_size, gate.clone())); + let mut upstream_to_client = TransferState::Running(LazyCopyBuffer::new(upstream_buf_size, gate)); + poll_fn(|cx| { + let a = transfer_one_direction(cx, &mut client_to_upstream, client, upstream)?; + let b = transfer_one_direction(cx, &mut upstream_to_client, upstream, client)?; + // It is not a problem if `ready!` returns early here: `transfer_one_direction` for the + // side that already reached `Done` keeps returning `Poll::Ready(Ok(()))` on every future + // call, so the other side's completion is picked up on a later poll. + ready!(a); + ready!(b); + Poll::Ready(Ok(())) + }) + .await +} + +enum TransferState { + Running(LazyCopyBuffer), + ShuttingDown, + Done, +} + +fn transfer_one_direction( + cx: &mut Context<'_>, + state: &mut TransferState, + r: &mut A, + w: &mut B, +) -> Poll> +where + A: AsyncRead + AsyncWrite + Unpin + ?Sized, + B: AsyncRead + AsyncWrite + Unpin + ?Sized, +{ + let mut r = Pin::new(r); + let mut w = Pin::new(w); + loop { + match state { + TransferState::Running(buf) => { + ready!(buf.poll_copy(cx, r.as_mut(), w.as_mut()))?; + *state = TransferState::ShuttingDown; + } + TransferState::ShuttingDown => { + // Propagated, matching tokio: a shutdown failure (e.g. the peer already reset) + // must end the copy the same as any other I/O error, not be silently ignored — + // ignoring it would leave the *other* direction's `try`-equivalent waiting + // forever for a completion that this side will never report. + ready!(w.as_mut().poll_shutdown(cx))?; + *state = TransferState::Done; + } + TransferState::Done => return Poll::Ready(Ok(())), + } + } +} + +/// A single direction's copy buffer. Structurally a port of `tokio::io::util::copy::CopyBuffer` +/// with one addition: `buf` is resized (never merely indexed into a smaller slice) at the point +/// where a fully-drained buffer is reset for the next read, based on whether the read that just +/// filled it reached capacity. +struct LazyCopyBuffer { + /// Floor size for this direction — `max_buf_size` itself if that is already ≤ + /// `PROBE_BUFFER_SIZE` (keeps tiny configured buffers, e.g. in tests, correct without special + /// casing). + small_size: usize, + max_buf_size: usize, + read_done: bool, + need_flush: bool, + pos: usize, + cap: usize, + buf: Vec, + /// Consecutive read cycles that exactly saturated `buf`. Escalation requires two, not one — + /// see the module docs for why a single full read isn't enough evidence of a sustained burst. + full_streak: u32, + /// Consulted at each resize point, not latched, so a connection follows the proxy's load + /// rather than the conditions it happened to be established under. See `LazyBufferGate`. + gate: LazyBufferGate, +} + +impl LazyCopyBuffer { + fn new(max_buf_size: usize, gate: LazyBufferGate) -> Self { + let small_size = max_buf_size.min(PROBE_BUFFER_SIZE); + // Start at the floor only if the gate is engaged. Disengaged, this allocates the full + // buffer once and (with the shrink below equally gated) never resizes it again, which is + // exactly what tokio's `CopyBuffer` did — so a proxy under the threshold pays none of + // this module's churn, not a reduced amount of it. + let initial_size = if gate.engaged() { small_size } else { max_buf_size }; + Self { + small_size, + max_buf_size, + full_streak: 0, + read_done: false, + need_flush: false, + pos: 0, + cap: 0, + buf: vec![0u8; initial_size], + gate, + } + } + + fn poll_fill_buf(&mut self, cx: &mut Context<'_>, reader: Pin<&mut R>) -> Poll> + where + R: AsyncRead + ?Sized, + { + let me = &mut *self; + let mut buf = ReadBuf::new(&mut me.buf); + buf.set_filled(me.cap); + let res = reader.poll_read(cx, &mut buf); + if let Poll::Ready(Ok(())) = res { + let filled_len = buf.filled().len(); + // No new bytes were added by this call — `AsyncRead`'s contract for that is EOF. + me.read_done = me.cap == filled_len; + me.cap = filled_len; + } + res + } + + fn poll_write_buf( + &mut self, + cx: &mut Context<'_>, + mut reader: Pin<&mut R>, + mut writer: Pin<&mut W>, + ) -> Poll> + where + R: AsyncRead + ?Sized, + W: AsyncWrite + ?Sized, + { + let me = &mut *self; + match writer.as_mut().poll_write(cx, &me.buf[me.pos..me.cap]) { + Poll::Pending => { + // Top up the buffer towards full if we can read a bit more data while the write + // is blocked — improves the chances of a large write once it can proceed. + if !me.read_done && me.cap < me.buf.len() { + ready!(me.poll_fill_buf(cx, reader.as_mut()))?; + } + Poll::Pending + } + res => res, + } + } + + fn poll_copy(&mut self, cx: &mut Context<'_>, mut reader: Pin<&mut R>, mut writer: Pin<&mut W>) -> Poll> + where + R: AsyncRead + ?Sized, + W: AsyncWrite + ?Sized, + { + // Mirror tokio's own `CopyBuffer::poll_copy`: consume one unit of the task's + // cooperative-scheduling budget on entry. Without this, a direction whose read and + // write never return `Pending` (a fast loopback pair, or two directions that keep + // topping each other off) can spin through this loop indefinitely inside one `poll` + // call and monopolize its worker thread instead of yielding back to the scheduler. + // + // One deliberate divergence from tokio's own accounting: tokio's `poll_proceed` only + // commits a budget unit when the poll goes on to make progress (`RestoreOnPending` + // gives the unit back otherwise); this always commits on entry. A connection parked on + // many no-progress wakeups therefore burns budget — and forces a coop yield — faster + // here than tokio would. That's simpler (no progress-tracking to thread through this + // port) and not unsafe: `consume_budget()` registers the waker before returning + // `Pending` (`poll_proceed` → `register_waker`), so an early return here can't strand a + // wakeup. It just yields somewhat more eagerly than tokio's own loop would. + if std::pin::pin!(tokio::task::coop::consume_budget()).poll(cx).is_pending() { + return Poll::Pending; + } + loop { + if self.cap < self.buf.len() && !self.read_done { + match self.poll_fill_buf(cx, reader.as_mut()) { + Poll::Ready(Ok(())) => {} + Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), + Poll::Pending => { + // Ignore a pending read when there's still buffered data to write. + if self.pos == self.cap { + // Flush before parking so a writer that only sends on flush (e.g. + // tokio-rustls, which defers ciphertext records) doesn't leave the + // last write sitting unsent while this side waits for more input — + // the other direction may be blocked waiting for exactly that reply. + if self.need_flush { + ready!(writer.as_mut().poll_flush(cx))?; + self.need_flush = false; + } + // The buffer is fully drained and no more data is ready right now — + // shrink to the floor before parking. Without this, a burst whose size + // happens to land exactly on a buffer boundary ends by parking right + // here still holding the escalated (max-sized) buffer: the shrink below + // only runs after a *completed* under-capacity read, and if the + // connection now goes idle — precisely the case this buffer exists to + // keep cheap — that under-capacity read may never come. Shrinking here + // costs the same one-round-trip re-escalation the design already pays + // for growth, and a real sustained burst won't reach this branch (more + // input is already buffered in the kernel, so the next read is Ready). + // + // Gated: below the threshold the buffer is left where it is, so a proxy + // carrying a handful of bulk streams never pays a resize. Re-read here + // rather than latched at construction, so a connection established while + // the proxy was quiet does start releasing once it gets busy. + if self.gate.engaged() && self.buf.len() > self.small_size { + self.buf = vec![0u8; self.small_size]; + } + self.full_streak = 0; + return Poll::Pending; + } + } + } + } + + while self.pos < self.cap { + let i = ready!(self.poll_write_buf(cx, reader.as_mut(), writer.as_mut()))?; + if i == 0 { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::WriteZero, "write zero byte into writer"))); + } + self.pos += i; + self.need_flush = true; + } + + // All data written — the buffer is empty again. Capture whether this cycle's read(s) + // exactly saturated it *before* resetting, since that's the escalate/de-escalate + // signal, then decide the next buffer size. + let was_full = self.cap == self.buf.len(); + self.pos = 0; + self.cap = 0; + + if self.read_done { + ready!(writer.as_mut().poll_flush(cx))?; + return Poll::Ready(Ok(())); + } + + if was_full { + self.full_streak += 1; + if self.full_streak >= 2 && self.buf.len() < self.max_buf_size { + // Two consecutive reads have now saturated the buffer — real evidence of a + // sustained burst, not a single message that happened to match the current + // size. Jump straight to the configured max: by this point more escalation + // steps would only delay reaching full efficiency for no added safety. + self.buf = vec![0u8; self.max_buf_size]; + } + } else { + // Came back under capacity — this cycle's burst evidence is gone, so require two + // fresh full reads before escalating again. Deliberately does NOT shrink the + // buffer here: a connection with continuously active but variably-sized traffic + // (never actually idle) would otherwise reallocate on every undersized read only + // to grow right back on the next burst — heap-thrashing a connection that never + // stopped transferring. The buffer only shrinks once the direction actually parks + // (see the `Poll::Pending` branch above), which is the one point that reliably + // distinguishes "genuinely idle" from "momentarily between packets." + self.full_streak = 0; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + /// Distinguishable per byte, so a dropped or reordered chunk fails the assert. + fn patterned(size: usize) -> Vec { + (0..size).map(|i| (i % 251) as u8).collect() + } + + #[tokio::test] + async fn large_payload_integrity_through_a_small_buffer() { + let (mut client, client_peer) = tokio::io::duplex(4096); + let (mut upstream, mut upstream_peer) = tokio::io::duplex(4096); + + // 512× the 512-byte copy buffer configured below, so the loop must iterate many times. + let payload = patterned(256 * 1024); + + let echo = tokio::spawn(async move { + let mut buf = vec![0u8; 8192]; + loop { + match upstream_peer.read(&mut buf).await { + Ok(0) | Err(_) => break, + Ok(n) => upstream_peer.write_all(&buf[..n]).await.unwrap(), + } + } + }); + + let (mut client_peer_read, mut client_peer_write) = tokio::io::split(client_peer); + let to_send = payload.clone(); + let write_task = tokio::spawn(async move { + client_peer_write.write_all(&to_send).await.unwrap(); + client_peer_write.shutdown().await.unwrap(); + }); + let read_task = tokio::spawn(async move { + let mut buf = Vec::new(); + client_peer_read.read_to_end(&mut buf).await.unwrap(); + buf + }); + + let copy = copy_bidirectional_lazy(&mut client, &mut upstream, 512, 512, LazyBufferGate::always()); + let (copy_result, write_result, received) = tokio::join!(copy, write_task, read_task); + copy_result.unwrap(); + write_result.unwrap(); + echo.await.unwrap(); + + let received = received.unwrap(); + assert_eq!(received.len(), payload.len(), "no bytes dropped or duplicated"); + assert_eq!(received, payload, "payload round-trips byte-for-byte, in order"); + } + + struct ErroringReader { + err_after: usize, + read: usize, + } + + impl AsyncRead for ErroringReader { + fn poll_read( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + if self.read >= self.err_after { + return Poll::Ready(Err(io::Error::new(io::ErrorKind::ConnectionReset, "simulated RST"))); + } + buf.put_slice(b"x"); + self.read += 1; + Poll::Ready(Ok(())) + } + } + + /// Composes a reader half and a writer half into one `AsyncRead + AsyncWrite` type, for + /// tests that need independent control over each side. + struct RW { + r: R, + w: W, + } + impl AsyncRead for RW { + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.r).poll_read(cx, buf) + } + } + impl AsyncWrite for RW { + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.w).poll_write(cx, buf) + } + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.w).poll_flush(cx) + } + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.w).poll_shutdown(cx) + } + } + + #[tokio::test] + async fn rst_on_one_leg_ends_the_copy_immediately() { + // The client leg errors after a few bytes; the upstream leg would otherwise read forever + // (never returns Ready). The copy must still resolve promptly with the error, rather + // than waiting on the upstream leg (that's what the idle timeout wrapping `forward()` is + // for in a real hang — this test asserts the copy itself doesn't need it for an outright + // I/O error). + use tokio::io::duplex; + + let mut client = ErroringReader { err_after: 4, read: 0 }; + let (mut client_write_sink, _keep_alive) = duplex(64); + let (mut upstream, _never_closes) = duplex(64); + + let mut client_rw = RW { r: &mut client, w: &mut client_write_sink }; + + let result = tokio::time::timeout( + std::time::Duration::from_secs(5), + copy_bidirectional_lazy(&mut client_rw, &mut upstream, 512, 512, LazyBufferGate::always()), + ) + .await + .expect("copy must resolve promptly on RST, not hang for the upstream leg"); + + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::ConnectionReset); + } + + #[tokio::test] + async fn half_close_from_client_side_lets_upstream_direction_finish() { + let (mut client, mut client_peer) = tokio::io::duplex(1024); + let (mut upstream, mut upstream_peer) = tokio::io::duplex(1024); + + // Client sends nothing and closes immediately (write-half EOF); upstream still has data + // queued for the client and must be allowed to deliver it before the copy completes. + client_peer.shutdown().await.unwrap(); + let upstream_write = tokio::spawn(async move { + upstream_peer.write_all(b"late data").await.unwrap(); + upstream_peer.shutdown().await.unwrap(); + }); + let read_task = tokio::spawn(async move { + let mut buf = Vec::new(); + client_peer.read_to_end(&mut buf).await.unwrap(); + buf + }); + + copy_bidirectional_lazy(&mut client, &mut upstream, 512, 512, LazyBufferGate::always()).await.unwrap(); + upstream_write.await.unwrap(); + let received = read_task.await.unwrap(); + assert_eq!(received, b"late data"); + } + + #[tokio::test] + async fn half_close_from_upstream_side_lets_client_direction_finish() { + let (mut client, mut client_peer) = tokio::io::duplex(1024); + let (mut upstream, mut upstream_peer) = tokio::io::duplex(1024); + + // Upstream sends nothing and closes immediately; the client still has data queued to send + // upstream and must be allowed to deliver it (recorded by upstream_peer) before completion. + upstream_peer.shutdown().await.unwrap(); + let client_write = tokio::spawn(async move { + client_peer.write_all(b"queued request").await.unwrap(); + client_peer.shutdown().await.unwrap(); + }); + let read_task = tokio::spawn(async move { + let mut buf = Vec::new(); + upstream_peer.read_to_end(&mut buf).await.unwrap(); + buf + }); + + copy_bidirectional_lazy(&mut client, &mut upstream, 512, 512, LazyBufferGate::always()).await.unwrap(); + client_write.await.unwrap(); + let received = read_task.await.unwrap(); + assert_eq!(received, b"queued request"); + } + + #[tokio::test] + async fn single_byte_message_with_nothing_queued_behind_it_is_not_held_waiting_for_more() { + // A lone byte (an MQTT PINGREQ, a short request) with no further data queued behind it + // must be forwarded immediately, never held waiting for more that will never come. + let (mut client, mut client_peer) = tokio::io::duplex(1024); + let (mut upstream, mut upstream_peer) = tokio::io::duplex(1024); + + let echo = tokio::spawn(async move { + let mut byte = [0u8; 1]; + upstream_peer.read_exact(&mut byte).await.unwrap(); + upstream_peer.write_all(&byte).await.unwrap(); + upstream_peer.shutdown().await.unwrap(); + }); + + let driver = tokio::spawn(async move { + client_peer.write_all(b"p").await.unwrap(); + let mut reply = [0u8; 1]; + client_peer.read_exact(&mut reply).await.unwrap(); + assert_eq!(&reply, b"p"); + client_peer.shutdown().await.unwrap(); + }); + + tokio::time::timeout( + std::time::Duration::from_secs(5), + copy_bidirectional_lazy(&mut client, &mut upstream, 65536, 65536, LazyBufferGate::always()), + ) + .await + .expect("a single queued byte must be forwarded promptly, not held waiting for more") + .unwrap(); + + driver.await.unwrap(); + echo.await.unwrap(); + } + + /// A writer that only actually delivers bytes to `inner` on `poll_flush`, and counts flushes. + /// Models a buffering transport like tokio-rustls, which documents that `poll_write` may not + /// send all data and `poll_flush` must be called to guarantee it does. + struct FlushGatedWriter { + inner: W, + pending: Vec, + flushes: std::sync::Arc, + } + impl AsyncRead for FlushGatedWriter { + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_read(cx, buf) + } + } + impl AsyncWrite for FlushGatedWriter { + fn poll_write(self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + let this = self.get_mut(); + this.pending.extend_from_slice(buf); + Poll::Ready(Ok(buf.len())) + } + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.pending.is_empty() { + return Pin::new(&mut this.inner).poll_flush(cx); + } + match Pin::new(&mut this.inner).poll_write(cx, &this.pending) { + Poll::Ready(Ok(n)) => { + this.pending.drain(..n); + this.flushes.fetch_add(1, Ordering::SeqCst); + if this.pending.is_empty() { + Pin::new(&mut this.inner).poll_flush(cx) + } else { + Poll::Pending + } + } + Poll::Ready(Err(e)) => Poll::Ready(Err(e)), + Poll::Pending => Poll::Pending, + } + } + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_shutdown(cx) + } + } + + #[tokio::test] + async fn write_is_flushed_before_parking_on_the_next_read() { + // Regression test for the tokio-rustls-shaped deadlock: client sends one request, + // upstream sends one short reply and then waits for the next request. If the + // client-facing writer defers delivery until an explicit flush, and the pump doesn't + // flush before blocking on the next client read, the reply sits undelivered forever. + let (mut client, mut client_peer) = tokio::io::duplex(1024); + let (upstream_side, mut upstream_peer) = tokio::io::duplex(1024); + let flushes = std::sync::Arc::new(AtomicUsize::new(0)); + let mut upstream = FlushGatedWriter { inner: upstream_side, pending: Vec::new(), flushes: flushes.clone() }; + + // Both peers are returned (not dropped) at the end of their task so the duplex stays + // open — otherwise the task ending would drop its half, the copy would see a real EOF, + // and the race below would be flaky depending on which finishes first. + let echo = tokio::spawn(async move { + let mut req = [0u8; 1]; + upstream_peer.read_exact(&mut req).await.unwrap(); + upstream_peer.write_all(b"reply").await.unwrap(); + // Upstream now waits indefinitely for the next request — nothing more is ever sent. + upstream_peer + }); + + let driver = tokio::spawn(async move { + client_peer.write_all(b"r").await.unwrap(); + let mut reply = [0u8; 5]; + client_peer.read_exact(&mut reply).await.unwrap(); + assert_eq!(&reply, b"reply"); + client_peer + }); + + // `upstream_peer` never sends anything more, so the copy itself runs forever; race it + // against `driver` instead. The real assertion is that `driver` completes at all — its + // `read_exact` only succeeds if the reply was actually flushed to the wire rather than + // left buffered while the pump parked on the next (never-arriving) client read. + let copy_fut = copy_bidirectional_lazy(&mut client, &mut upstream, 512, 512, LazyBufferGate::always()); + tokio::pin!(copy_fut); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + tokio::select! { + result = &mut copy_fut => panic!("copy must not complete while driver is still waiting: {result:?}"), + driver_result = driver => driver_result.unwrap(), + } + }) + .await + .expect("driver must complete — the reply must have been flushed to it"); + + assert!(flushes.load(Ordering::SeqCst) > 0, "poll_flush must have been called to deliver the reply"); + echo.abort(); + } + + /// A writer whose `shutdown()` always fails, to prove shutdown errors are propagated rather + /// than swallowed. + struct FailingShutdownWriter { + inner: W, + } + impl AsyncRead for FailingShutdownWriter { + fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_read(cx, buf) + } + } + impl AsyncWrite for FailingShutdownWriter { + fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_write(cx, buf) + } + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.inner).poll_flush(cx) + } + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Err(io::Error::other("simulated shutdown failure"))) + } + } + + #[tokio::test] + async fn shutdown_failure_is_propagated_not_swallowed() { + // The client leg reaches EOF, so its shutdown of the upstream writer runs — and that + // shutdown fails. The whole copy must end with that error immediately; if the failure + // were swallowed, this would instead hang waiting on the (otherwise-idle) upstream leg. + let (mut client, client_peer) = tokio::io::duplex(64); + drop(client_peer); // client reader hits EOF right away + let (upstream_side, _never_closes) = tokio::io::duplex(64); + let mut upstream = FailingShutdownWriter { inner: upstream_side }; + + let result = tokio::time::timeout( + std::time::Duration::from_secs(5), + copy_bidirectional_lazy(&mut client, &mut upstream, 512, 512, LazyBufferGate::always()), + ) + .await + .expect("a failed shutdown must end the copy promptly, not hang waiting on the other direction"); + + assert!(result.is_err()); + assert_eq!(result.unwrap_err().kind(), io::ErrorKind::Other); + } + + #[tokio::test] + async fn buffer_escalates_on_a_sustained_burst_and_shrinks_back_once_it_ends() { + // max_buf_size well above PROBE_BUFFER_SIZE so both branches are actually reachable + // (with max_buf_size <= 512 the "small" and "max" sizes are identical, and this test + // would pass even with the escalate/de-escalate logic deleted entirely). + const MAX: usize = 8192; + let (mut client, mut client_peer) = tokio::io::duplex(MAX * 2); + let (mut upstream, mut upstream_peer) = tokio::io::duplex(MAX * 2); + + let echo = tokio::spawn(async move { + let mut buf = vec![0u8; MAX]; + loop { + match upstream_peer.read(&mut buf).await { + Ok(0) => break, + Ok(n) => upstream_peer.write_all(&buf[..n]).await.unwrap(), + Err(_) => break, + } + } + }); + + let driver = tokio::spawn(async move { + // A burst several times larger than PROBE_BUFFER_SIZE, forcing escalation past two + // consecutive full reads. This proves round-trip correctness *despite* the buffer + // resizing mid-transfer, not that it actually resized — `buf_len_transitions_through_a_burst_and_back` + // below inspects the buffer directly for that. + let burst = vec![7u8; MAX * 4]; + client_peer.write_all(&burst).await.unwrap(); + let mut readback = vec![0u8; burst.len()]; + client_peer.read_exact(&mut readback).await.unwrap(); + assert_eq!(readback, burst, "burst round-trips byte-for-byte despite the escalating buffer"); + + // Then a single small message — after a burst this size, an un-shrunk buffer would + // still work, but the point of the fix is that it doesn't stay at the burst's size. + client_peer.write_all(&[9u8]).await.unwrap(); + let mut one = [0u8; 1]; + client_peer.read_exact(&mut one).await.unwrap(); + assert_eq!(one[0], 9); + + client_peer.shutdown().await.unwrap(); + }); + + tokio::time::timeout( + std::time::Duration::from_secs(10), + copy_bidirectional_lazy(&mut client, &mut upstream, MAX, MAX, LazyBufferGate::always()), + ) + .await + .expect("burst then small message must complete without stalling") + .unwrap(); + driver.await.unwrap(); + echo.await.unwrap(); + } + + /// Drives `LazyCopyBuffer::poll_copy` directly (no sockets, no tokio scheduler) so the two + /// state transitions the escalate/shrink design depends on can be asserted directly instead + /// of inferred from end-to-end timing: `full_streak` reaching 2 grows `buf` to `max_buf_size`, + /// and parking with nothing left to write (the burst has ended and the reader has nothing + /// more ready) shrinks it back to `small_size`. + #[test] + fn buf_len_transitions_through_a_burst_and_back() { + /// Fills every requested read fully (whatever the buffer's current capacity is) for the + /// first `full_reads` calls, then a single byte (deliberately under capacity) to end the + /// burst, then parks forever. + struct BurstThenOneByte { + full_reads: usize, + calls: usize, + } + impl AsyncRead for BurstThenOneByte { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + self.calls += 1; + if self.calls <= self.full_reads { + let n = buf.remaining(); + buf.put_slice(&vec![7u8; n]); + Poll::Ready(Ok(())) + } else if self.calls == self.full_reads + 1 { + buf.put_slice(&[9u8]); + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + } + + /// Records the length of every write it receives and always accepts immediately. + struct RecordingWriter(std::sync::Arc>>); + impl AsyncWrite for RecordingWriter { + fn poll_write(self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + self.0.lock().unwrap().push(buf.len()); + Poll::Ready(Ok(buf.len())) + } + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + const MAX: usize = 8192; + let writes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + // Reads 1–2 (at the floor size) drive `full_streak` to 2 and trigger escalation; read 3 + // is the first one issued against the now-`MAX`-sized buffer, so it's the one that must + // actually observe (and write out) the escalated size. + let mut reader = BurstThenOneByte { full_reads: 3, calls: 0 }; + let mut writer = RecordingWriter(writes.clone()); + let mut lazy_buf = LazyCopyBuffer::new(MAX, LazyBufferGate::always()); + let small_size = lazy_buf.small_size; + + let waker = std::task::Waker::noop(); + let mut cx = Context::from_waker(waker); + // A single call: `poll_copy`'s inner loop only ever returns on `Pending` or completion, so + // it runs the two full reads, the escalation, the under-capacity read, and the shrink + // entirely within this one call before parking on the reader's subsequent `Pending`. + let result = lazy_buf.poll_copy(&mut cx, Pin::new(&mut reader), Pin::new(&mut writer)); + + assert!(result.is_pending(), "reader parks after the burst ends, so this call must not resolve"); + assert_eq!( + *writes.lock().unwrap(), + vec![small_size, small_size, MAX, 1], + "two floor-sized writes, then escalation to MAX on the third, then the under-capacity write that ends the burst" + ); + assert_eq!(lazy_buf.buf.len(), small_size, "buffer must have shrunk back to the floor after the burst ended"); + assert_eq!(lazy_buf.full_streak, 0, "the under-capacity read must reset the streak"); + } + + /// Fills every read to the buffer's current capacity `full_reads` times, then returns one + /// under-capacity byte to end the burst, then parks. Shared by the gate tests below. + struct BurstThenPark { + full_reads: usize, + calls: usize, + } + impl AsyncRead for BurstThenPark { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + self.calls += 1; + if self.calls <= self.full_reads { + let n = buf.remaining(); + buf.put_slice(&vec![7u8; n]); + Poll::Ready(Ok(())) + } else if self.calls == self.full_reads + 1 { + buf.put_slice(&[9u8]); + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + } + } + + /// Accepts every write immediately and records its length. + struct SizeRecordingWriter(std::sync::Arc>>); + impl AsyncWrite for SizeRecordingWriter { + fn poll_write(self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + self.0.lock().unwrap().push(buf.len()); + Poll::Ready(Ok(buf.len())) + } + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + /// Below the threshold the module must behave exactly like the `tokio::io::CopyBuffer` it + /// replaced: one full-size allocation up front, no probing at the floor, and no shrink on + /// park. This is the property that makes the gate worth having — a proxy carrying a few bulk + /// streams pays *none* of the churn, not a reduced amount of it. + #[test] + fn a_disengaged_gate_allocates_full_size_once_and_never_resizes() { + const MAX: usize = 8192; + let writes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let mut reader = BurstThenPark { full_reads: 3, calls: 0 }; + let mut writer = SizeRecordingWriter(writes.clone()); + let mut lazy_buf = LazyCopyBuffer::new(MAX, LazyBufferGate::never()); + + assert_eq!(lazy_buf.buf.len(), MAX, "a disengaged gate must allocate the full buffer up front, not the floor"); + + let waker = std::task::Waker::noop(); + let mut cx = Context::from_waker(waker); + let result = lazy_buf.poll_copy(&mut cx, Pin::new(&mut reader), Pin::new(&mut writer)); + + assert!(result.is_pending(), "the reader parks after the burst"); + assert_eq!( + *writes.lock().unwrap(), + vec![MAX, MAX, MAX, 1], + "every read should have been served at full size — no floor-sized probing" + ); + assert_eq!(lazy_buf.buf.len(), MAX, "a disengaged gate must not shrink on park"); + } + + /// The gate is re-read at each resize point rather than latched at construction. A connection + /// established while the proxy was quiet must start releasing its buffer once the proxy gets + /// busy — otherwise long-lived connections that accumulate while idle, the exact shape this + /// module exists for, would each keep a full-size buffer forever. + #[test] + fn crossing_the_threshold_makes_an_already_established_connection_release() { + const MAX: usize = 8192; + const THRESHOLD: u64 = 10; + let metrics = Arc::new(GlobalMetrics::default()); + let gate = LazyBufferGate::with_metrics(metrics.clone(), THRESHOLD); + + // Established while quiet: full-size buffer, gate disengaged. + let mut lazy_buf = LazyCopyBuffer::new(MAX, gate); + assert_eq!(lazy_buf.buf.len(), MAX, "established below the threshold, so it starts full-size"); + let small_size = lazy_buf.small_size; + + // The proxy fills up. Exactly at the threshold, which must count as engaged. + metrics.active_connections.store(THRESHOLD, Ordering::Relaxed); + + let writes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let mut reader = BurstThenPark { full_reads: 3, calls: 0 }; + let mut writer = SizeRecordingWriter(writes.clone()); + let waker = std::task::Waker::noop(); + let mut cx = Context::from_waker(waker); + let result = lazy_buf.poll_copy(&mut cx, Pin::new(&mut reader), Pin::new(&mut writer)); + + assert!(result.is_pending(), "the reader parks after the burst"); + assert_eq!( + lazy_buf.buf.len(), + small_size, + "once active connections reach the threshold, the next park must release the full-size buffer" + ); + } + + /// One below the threshold is still disengaged — the boundary is `>=`, and a test that only + /// checked 0-vs-huge would not catch an off-by-one there. + #[test] + fn just_under_the_threshold_stays_disengaged() { + const MAX: usize = 8192; + const THRESHOLD: u64 = 10; + let metrics = Arc::new(GlobalMetrics::default()); + metrics.active_connections.store(THRESHOLD - 1, Ordering::Relaxed); + let lazy_buf = LazyCopyBuffer::new(MAX, LazyBufferGate::with_metrics(metrics.clone(), THRESHOLD)); + assert_eq!(lazy_buf.buf.len(), MAX, "one connection below the threshold must still be disengaged"); + + metrics.active_connections.store(THRESHOLD, Ordering::Relaxed); + let engaged_buf = LazyCopyBuffer::new(MAX, LazyBufferGate::with_metrics(metrics, THRESHOLD)); + assert_eq!(engaged_buf.buf.len(), engaged_buf.small_size, "at the threshold it must engage"); + } + + /// A single under-capacity read must NOT shrink the buffer while the connection is still + /// actively transferring (more data already ready right after) — only parking does. Without + /// this, a connection with continuously active but variably-sized traffic would reallocate on + /// every undersized read only to grow right back on the next burst: heap-thrashing a + /// connection that never actually went idle. + #[test] + fn a_single_undersized_read_mid_burst_does_not_shrink_the_buffer() { + /// full reads, then one under-capacity read, then full reads again, then parks forever — + /// modeling a connection that dips below capacity once but never actually goes idle. + struct DipThenResumeBurst { + calls: usize, + } + impl AsyncRead for DipThenResumeBurst { + fn poll_read(mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll> { + self.calls += 1; + match self.calls { + 1 | 2 | 3 | 5 => { + let n = buf.remaining(); + buf.put_slice(&vec![7u8; n]); + Poll::Ready(Ok(())) + } + 4 => { + buf.put_slice(&[9u8]); + Poll::Ready(Ok(())) + } + _ => Poll::Pending, + } + } + } + + struct RecordingWriter(std::sync::Arc>>); + impl AsyncWrite for RecordingWriter { + fn poll_write(self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8]) -> Poll> { + self.0.lock().unwrap().push(buf.len()); + Poll::Ready(Ok(buf.len())) + } + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + const MAX: usize = 8192; + let writes = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let mut reader = DipThenResumeBurst { calls: 0 }; + let mut writer = RecordingWriter(writes.clone()); + let mut lazy_buf = LazyCopyBuffer::new(MAX, LazyBufferGate::always()); + let small_size = lazy_buf.small_size; + + let waker = std::task::Waker::noop(); + let mut cx = Context::from_waker(waker); + let result = lazy_buf.poll_copy(&mut cx, Pin::new(&mut reader), Pin::new(&mut writer)); + + assert!(result.is_pending(), "reader parks at the end, so this call must not resolve"); + assert_eq!( + *writes.lock().unwrap(), + vec![small_size, small_size, MAX, 1, MAX], + "the write right after the under-capacity read must still be MAX-sized — the buffer must not have shrunk from the single dip" + ); + // Only parking (the reader's final Pending) shrinks it — proven separately by the + // previous test; here the point is that call 4 alone did not. + assert_eq!(lazy_buf.buf.len(), small_size, "buffer shrinks once the reader actually parks"); + } +} diff --git a/src/lib.rs b/src/lib.rs index 157d98a..ee25921 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,6 +4,7 @@ extern crate napi_derive; mod balancer; +mod copy; mod error; mod http_listener; mod http_proxy; diff --git a/src/proxy.rs b/src/proxy.rs index 81bd213..fb4417c 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -3,7 +3,8 @@ use crate::listener::spawn_listeners; use crate::metrics::{total_of, GlobalMetrics, ListenerMetrics}; use crate::protection::ProtectionState; use crate::proxy_conn::{ - ConnContext, JsEvent, DEFAULT_COPY_BUFFER_SIZE, MAX_COPY_BUFFER_SIZE, MIN_COPY_BUFFER_SIZE, + ConnContext, JsEvent, DEFAULT_COPY_BUFFER_SIZE, DEFAULT_LAZY_COPY_BUFFER_THRESHOLD, MAX_COPY_BUFFER_SIZE, + MIN_COPY_BUFFER_SIZE, }; use crate::router::{ build_route_table, ForwardFingerprint, ListenerTlsSpec, LiveRouteTable, RouteSpec, @@ -131,15 +132,22 @@ pub struct JsProxyConfig { pub listeners: Vec, pub routes: Vec, pub worker_threads: Option, - /// Per-direction copy buffer, in bytes (default 8192). One buffer per direction is held for - /// the whole life of every proxied connection, idle or not, so these are a direct multiplier - /// on per-connection memory: `(client + upstream) × connections`, i.e. - /// `2 × readBufferSize × connections` when both directions use this value. + /// Per-direction copy buffer *maximum*, in bytes (default 8192). Each direction starts at a + /// small fixed floor and escalates to this value only for a sustained burst, dropping back to + /// the floor once the burst ends — it is not a permanent per-connection allocation. It bounds + /// how large a single bursty transfer's buffer may grow, which is what to weigh against many + /// concurrent bursty transfers held at that size at once. pub read_buffer_size: Option, /// Overrides `readBufferSize` for the client→upstream direction only. pub client_read_buffer_size: Option, /// Overrides `readBufferSize` for the upstream→client direction only. pub upstream_read_buffer_size: Option, + /// Active connections at or above which the copy buffers escalate and release rather than + /// being held at full size for each connection's whole life (default 1000). `0` engages that + /// always; a value above this proxy's peak concurrency disables it, giving each direction one + /// full-size buffer with no resize churn. Per proxy, so a replication port-set and an MQTT + /// fan-out port-set can differ. + pub lazy_copy_buffer_threshold: Option, } #[napi(object)] @@ -262,6 +270,7 @@ pub struct SymphonyProxyWrap { idle_timeout: Duration, client_read_buffer_size: usize, upstream_read_buffer_size: usize, + lazy_copy_buffer_threshold: u64, // Shared runtime state route_table: Arc, suspended_registry: Arc, @@ -326,6 +335,8 @@ impl SymphonyProxyWrap { Some(v) => resolve_copy_buffer_size(Some(v), "upstreamReadBufferSize"), None => base_read_buffer_size, }; + let lazy_copy_buffer_threshold = + u64::from(config.lazy_copy_buffer_threshold.unwrap_or(DEFAULT_LAZY_COPY_BUFFER_THRESHOLD)); let mut internal_listeners = Vec::new(); let mut listener_states = Vec::new(); @@ -424,6 +435,7 @@ impl SymphonyProxyWrap { idle_timeout, client_read_buffer_size, upstream_read_buffer_size, + lazy_copy_buffer_threshold, route_table: Arc::new(LiveRouteTable(arc_swap::ArcSwap::new(Arc::new(table)))), suspended_registry: SuspendedRegistry::new(), global_metrics: Arc::new(GlobalMetrics::default()), @@ -461,6 +473,7 @@ impl SymphonyProxyWrap { upstream_connect_timeout, client_read_buffer_size: self.client_read_buffer_size, upstream_read_buffer_size: self.upstream_read_buffer_size, + lazy_copy_buffer_threshold: self.lazy_copy_buffer_threshold, js_emit: self.js_emit.clone(), }); diff --git a/src/proxy_conn.rs b/src/proxy_conn.rs index 2c36f32..6a2a7a3 100644 --- a/src/proxy_conn.rs +++ b/src/proxy_conn.rs @@ -9,18 +9,16 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; use std::marker::Unpin; -use tokio::io::copy_bidirectional_with_sizes; use tokio::net::TcpStream; use tokio::time::timeout; use tokio_rustls::TlsAcceptor; -/// Default per-direction copy buffer. Matches what `tokio::io::copy_bidirectional` uses -/// internally, so wiring the previously-inert `readBufferSize` through leaves any config that -/// does not set it on exactly the footprint it already had. Do not raise this default: both -/// buffers are held for the whole life of every connection whether or not it is transferring, so -/// the number multiplies straight into per-connection memory (`(client + upstream) × connections`) -/// — across a million mostly-idle MQTT subscribers, 64 KiB per direction instead of 8 KiB is -/// ~107 GiB of buffers that are never read. +/// Default per-direction copy buffer *maximum*. Matches what `tokio::io::copy_bidirectional` used +/// to hold permanently, so wiring the previously-inert `readBufferSize` through leaves any config +/// that does not set it on exactly the footprint it already had. `copy::LazyCopyBuffer` only grows +/// to this ceiling for a direction observed mid-burst, and shrinks back to a small fixed floor once +/// it parks — so this default is not a permanent per-connection cost, but it does bound how large a +/// buffer every route that doesn't override it may grow to while actively bursting. /// /// Applies to the plain proxying path. A route that injects HTTP headers takes /// `http_proxy::proxy_http1_rewriting` instead, which frames with its own fixed buffers. @@ -32,6 +30,17 @@ pub const MIN_COPY_BUFFER_SIZE: usize = 512; /// 1 MiB per direction is already 2 MiB per connection; beyond this a config value is far more /// likely to be a units mistake than an intent. pub const MAX_COPY_BUFFER_SIZE: usize = 1024 * 1024; +/// Proxy-wide active connections at or above which `copy::LazyCopyBuffer` starts escalating and +/// releasing its buffers instead of holding one full-size buffer per direction (`copy::LazyBufferGate`). +/// +/// Chosen so the two shapes symphony actually carries land on opposite sides without anyone +/// configuring it. A replication port-set runs a handful of bulk streams: it stays under this, +/// keeps a static buffer, and pays no resize churn for a per-connection saving it would never +/// notice. An MQTT fan-out port-set runs tens or hundreds of thousands of mostly-parked +/// subscribers: it is far above this, and 2 x `readBufferSize` per parked connection is the cost +/// that matters. At the threshold itself the held-buffer cost is ~128 MiB at a 64 KiB +/// `readBufferSize` — enough to be worth reclaiming, low enough that nothing under it is at risk. +pub const DEFAULT_LAZY_COPY_BUFFER_THRESHOLD: u32 = 1000; /// JS event types emitted from connection tasks back to Node. #[derive(Debug)] @@ -75,6 +84,10 @@ pub struct ConnContext { /// MQTT is strongly asymmetric: after SUBSCRIBE a client sends almost nothing but PINGREQ, /// while the broker carries the whole fan-out. pub upstream_read_buffer_size: usize, + /// Active connections at or above which the copy buffers escalate/release rather than being + /// held at full size. `0` always, above peak concurrency never. See + /// `DEFAULT_LAZY_COPY_BUFFER_THRESHOLD` and `copy::LazyBufferGate`. + pub lazy_copy_buffer_threshold: u64, pub js_emit: Arc>, } @@ -331,8 +344,14 @@ where write_connection_prefix(upstream, sf).await?; let rewrites = header_rewrites(sf, l7_http1); if rewrites.is_empty() { - copy_both_ways(client, upstream, ctx.client_read_buffer_size, ctx.upstream_read_buffer_size) - .await + copy_both_ways( + client, + upstream, + ctx.client_read_buffer_size, + ctx.upstream_read_buffer_size, + crate::copy::LazyBufferGate::new(ctx.global_metrics.clone(), ctx.lazy_copy_buffer_threshold), + ) + .await } else { crate::http_proxy::proxy_http1_rewriting(client, upstream, &rewrites).await } @@ -353,21 +372,22 @@ where /// and observe which configured size reached which direction — a swapped mapping here is invisible /// to a round-trip test, since both directions still deliver every byte. /// -/// `copy_bidirectional_with_sizes` takes `(a, b, a_to_b, b_to_a)`, so with `a` = the client the -/// client buffer sizes reads *from* the client and the upstream buffer sizes the fan-out half. +/// Delegates to `copy::copy_bidirectional_lazy` (escalating per-direction buffers, see +/// `src/copy.rs`) rather than `tokio::io::copy_bidirectional_with_sizes`; the argument order is +/// the same either way — `client_read_buffer_size` sizes reads *from* the client, `upstream_read_buffer_size` +/// sizes the fan-out half. async fn copy_both_ways( client: &mut C, upstream: &mut U, client_read_buffer_size: usize, upstream_read_buffer_size: usize, + gate: crate::copy::LazyBufferGate, ) -> std::io::Result<()> where C: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, U: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, { - copy_bidirectional_with_sizes(client, upstream, client_read_buffer_size, upstream_read_buffer_size) - .await - .map(|_| ()) + crate::copy::copy_bidirectional_lazy(client, upstream, client_read_buffer_size, upstream_read_buffer_size, gate).await } /// Per-connection source-address + fingerprint forwarding parameters. All fields are `Copy` @@ -635,7 +655,7 @@ mod tests { let from_upstream = Arc::new(Mutex::new(Vec::new())); let mut client = CapacityRecorder::new(64 * 1024, from_client.clone()); let mut upstream = CapacityRecorder::new(64 * 1024, from_upstream.clone()); - copy_both_ways(&mut client, &mut upstream, client_size, upstream_size).await.unwrap(); + copy_both_ways(&mut client, &mut upstream, client_size, upstream_size, crate::copy::LazyBufferGate::always()).await.unwrap(); let max_of = |v: &Arc>>| *v.lock().unwrap().iter().max().unwrap(); (max_of(&from_client), max_of(&from_upstream)) } diff --git a/ts/addon.d.ts b/ts/addon.d.ts index b82d693..2e458b2 100644 --- a/ts/addon.d.ts +++ b/ts/addon.d.ts @@ -104,16 +104,25 @@ export interface JsProxyConfig { routes: Array workerThreads?: number /** - * Per-direction copy buffer, in bytes (default 8192). One buffer per direction is held for - * the whole life of every proxied connection, idle or not, so these are a direct multiplier - * on per-connection memory: `(client + upstream) × connections`, i.e. - * `2 × readBufferSize × connections` when both directions use this value. + * Per-direction copy buffer *maximum*, in bytes (default 8192). Each direction starts at a + * small fixed floor and escalates to this value only for a sustained burst, dropping back to + * the floor once the burst ends — it is not a permanent per-connection allocation. It bounds + * how large a single bursty transfer's buffer may grow, which is what to weigh against many + * concurrent bursty transfers held at that size at once. */ readBufferSize?: number /** Overrides `readBufferSize` for the client→upstream direction only. */ clientReadBufferSize?: number /** Overrides `readBufferSize` for the upstream→client direction only. */ upstreamReadBufferSize?: number + /** + * Active connections at or above which the copy buffers escalate and release rather than + * being held at full size for each connection's whole life (default 1000). `0` engages that + * always; a value above this proxy's peak concurrency disables it, giving each direction one + * full-size buffer with no resize churn. Per proxy, so a replication port-set and an MQTT + * fan-out port-set can differ. + */ + lazyCopyBufferThreshold?: number } export interface JsListenerProtectionHotConfig { /** Port of the listener to update. Must match a listener configured at start. */ diff --git a/ts/server.ts b/ts/server.ts index 02921f2..5ed733c 100644 --- a/ts/server.ts +++ b/ts/server.ts @@ -55,6 +55,7 @@ interface FileProxyConfig { readBufferSize?: number; clientReadBufferSize?: number; upstreamReadBufferSize?: number; + lazyCopyBufferThreshold?: number; } interface ConfigFile { @@ -135,6 +136,7 @@ function toProxyConfig(spec: FileProxyConfig, baseDir: string): ProxyConfig { readBufferSize: spec.readBufferSize, clientReadBufferSize: spec.clientReadBufferSize, upstreamReadBufferSize: spec.upstreamReadBufferSize, + lazyCopyBufferThreshold: spec.lazyCopyBufferThreshold, }; } @@ -329,6 +331,7 @@ class ServerState { readBufferSize: proxyConfig.readBufferSize, clientReadBufferSize: proxyConfig.clientReadBufferSize, upstreamReadBufferSize: proxyConfig.upstreamReadBufferSize, + lazyCopyBufferThreshold: proxyConfig.lazyCopyBufferThreshold, }); if (existing && existing.constructionSig === constructionSig) { // Same listeners (presence-unchanged) → hot-swap routes and protection contents. diff --git a/ts/types.ts b/ts/types.ts index 2702b2a..0deec8d 100644 --- a/ts/types.ts +++ b/ts/types.ts @@ -232,20 +232,35 @@ export interface ProxyConfig { /** Number of tokio worker threads. Defaults to available CPU count. */ workerThreads?: number; /** - * Per-direction copy buffer size in bytes. Default: 8192. Clamped to [512, 1048576]. + * Per-direction copy buffer *maximum* in bytes. Default: 8192. Clamped to [512, 1048576]. * - * One buffer per direction is allocated per proxied connection and held for its whole life, - * transferring or not, so these are a direct multiplier on per-connection memory: - * `(client + upstream) x connections`, which is `2 x readBufferSize x connections` only when - * both directions use this value. Raising it buys throughput on a handful of bulk streams (replication) and - * costs gigabytes on hundreds of thousands of mostly-idle ones (MQTT subscribers). Since the - * value is per proxy, tune it per port-set rather than fleet-wide. + * Each direction starts at a small fixed floor and escalates to this value only once a + * sustained burst is observed, dropping back to the floor as soon as the burst ends — it is + * not a permanent per-connection allocation, so raising it does not cost this much memory + * per idle connection. It does bound how large a buffer a genuinely bursty transfer (e.g. + * replication) may grow to, which trades off against the per-transfer memory a proxy running + * many concurrent bursty transfers at once will hold. Since the value is per proxy, tune it + * per port-set rather than fleet-wide. */ readBufferSize?: number; /** Overrides `readBufferSize` for the client -> upstream direction only. */ clientReadBufferSize?: number; /** Overrides `readBufferSize` for the upstream -> client direction only. */ upstreamReadBufferSize?: number; + /** + * Proxy-wide active connections at or above which the copy buffers escalate and release + * rather than each connection holding its full buffer for its whole life. Default 1000. + * + * The memory saving scales with connection count; the per-burst cost of growing and releasing + * a buffer does not. A port-set carrying a few bulk streams (replication) therefore stays + * below this and keeps one full-size buffer per direction — no resize churn at all — while an + * MQTT fan-out port-set with tens of thousands of parked subscribers is far above it, where + * `2 x readBufferSize` per parked connection is the cost that matters. + * + * `0` engages it always. A value above this proxy's peak concurrency disables it, which is + * exactly the pre-existing behaviour (one full-size buffer per direction, allocated once). + */ + lazyCopyBufferThreshold?: number; } // ── Hot-swap config ───────────────────────────────────────────────────────────