diff --git a/DESIGN.md b/DESIGN.md index 53aa4ceb..57f64812 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -1104,11 +1104,28 @@ The `webui::streaming` module provides: `write` returns a typed error (`HandlerError::ClientDisconnected` / `HandlerError::StreamTimeout`) so the handler aborts the render rather than waste CPU producing bytes that have nowhere to go. + The coalescing target is not a byte cap: writes are appended whole before + checking the target, so a large write can produce a larger chunk. The channel + bounds chunk count, not payload bytes or total in-flight memory. - **`ChunkPool`** — lock-free shared pool of chunk buffers. Used via `StreamingWriter::new_pooled` to recycle the per-flush `Vec` - across requests, eliminating per-flush heap allocation in - steady-state high-RPS workloads. + across requests. Returned buffers are cleared and retained only if their + capacity is at most the configured `chunk_size` and the pool has space. + Oversized buffers are dropped, not shrunk. Summed idle `Vec` capacity is + bounded by `max_pool.max(1) × chunk_size`; allocator, queue, and owner + metadata are excluded. Active producer buffers, pending sends, queued chunks, + and consumer-held references are additional. A chunk becomes eligible for + return only after its final `Bytes` reference (including clones and slices) + drops, on whichever thread owns that reference. `Bytes::from_owner` still + allocates owner metadata per chunk; pooling avoids buffer allocations on + suitably sized hits, not all per-flush heap allocations. + Acquired buffers are empty with capacity at least `chunk_size`; any smaller + returned buffer is reserved relative to its cleared length, not its capacity. + Default writer pools remain sized at `StreamingWriter::CHUNK_TARGET + 1024` + (5 KiB), including the writer's existing 1 KiB headroom. Custom targets need + the same headroom. Mismatched sizes and recurring oversized writes can trade + extra allocations for bounded idle retention. ### Progressive Response API diff --git a/crates/webui/src/streaming.rs b/crates/webui/src/streaming.rs index 4bca7401..f12d1501 100644 --- a/crates/webui/src/streaming.rs +++ b/crates/webui/src/streaming.rs @@ -16,11 +16,12 @@ //! //! * [`StreamingWriter`] — coalesces small writes into ~4 KB chunks and //! pushes them through a **bounded** [`tokio::sync::mpsc::Sender`]. The -//! bound (`DEFAULT_CHANNEL_CAPACITY = 4` chunks ≈ 16 KB) provides +//! bound (`DEFAULT_CHANNEL_CAPACITY = 4` chunks) provides //! backpressure: when a slow client cannot keep up, the producer parks //! on the channel until the receiver drains, instead of queuing the -//! entire response in memory. A configurable flush deadline (via -//! [`with_flush_timeout`](StreamingWriter::with_flush_timeout)) caps +//! entire response in memory. The coalescing target is not a byte cap: +//! a large write can produce a larger chunk. A configurable flush deadline +//! (via [`with_flush_timeout`](StreamingWriter::with_flush_timeout)) caps //! the maximum time a producer thread can be parked, bounding the //! slow-loris DoS surface to `timeout × concurrent_renders`. When the //! receiver is dropped (client disconnect) or the deadline elapses, @@ -29,8 +30,8 @@ //! //! * [`ChunkPool`] — lock-free shared pool of chunk buffers. Used via //! [`StreamingWriter::new_pooled`] to recycle the per-flush `Vec` -//! across requests, eliminating per-flush heap allocation in -//! steady-state high-RPS workloads. +//! across requests, avoiding chunk-buffer allocations on suitably sized +//! pool hits. //! //! Hot-path allocation profile: //! @@ -39,9 +40,9 @@ //! when `len < cap`; when `len == cap`, `Bytes::from(Vec)` is still a //! move via `into_boxed_slice`). Plus one small `Box` for the //! refcount metadata. -//! * `StreamingWriter::new_pooled()`: zero per-flush heap allocation in -//! steady state — chunk buffers come from the pool and return on -//! `Bytes` drop. Single atomic CAS per acquire/release. +//! * `StreamingWriter::new_pooled()`: eligible chunk buffers return to +//! the pool on the last `Bytes` drop. `Bytes::from_owner` still allocates +//! owner metadata per chunk; pooling does not eliminate all heap allocations. //! //! # Per-render HTML injection //! @@ -71,41 +72,44 @@ use webui_handler::{FlushWriter, HandlerError, ResponseWriter, Result}; /// allocations across `StreamingWriter` instances. /// /// Backed by a [`crossbeam_queue::ArrayQueue`] (MPMC, lock-free, fixed -/// capacity). Acquiring a buffer is a single atomic CAS; releasing is -/// the same. When the pool is empty, `acquire` allocates a fresh -/// `Vec`. When the pool is full, `release` drops the buffer. +/// capacity). Acquiring and releasing use atomic queue operations. When the +/// pool is empty, a fresh `Vec` is allocated. When the pool is full or a +/// returned buffer's capacity exceeds `chunk_size`, the buffer is dropped, +/// not shrunk. /// /// # Lifetime model /// /// A buffer leaves the pool on `acquire`, gets handed to -/// [`bytes::Bytes::from_owner`] wrapped in a [`PooledChunk`] owner, -/// and is released back to the pool when **the last `Bytes` reference -/// is dropped** — typically after the HTTP framework has flushed the -/// chunk to the wire. Because `Bytes` may be dropped on any thread -/// (the actix worker that wrote the chunk to the socket, not the +/// [`bytes::Bytes::from_owner`] wrapped in an owner, +/// and becomes eligible for reuse only when **the last `Bytes` reference +/// is dropped**, including clones and slices - typically after the HTTP +/// framework has flushed the chunk to the wire. Because `Bytes` may be dropped +/// on any thread (the actix worker that wrote the chunk to the socket, not the /// `spawn_blocking` worker that produced it), the pool MUST be /// thread-safe — `ArrayQueue` is. /// /// # Sizing /// /// `max_pool` should match the expected concurrent in-flight chunk -/// count: `concurrent_renders × channel_capacity` in the worst case. -/// For the production setup (4-chunk channels, ~100 concurrent -/// renders), `max_pool = 512` covers the working set; surplus buffers -/// are dropped when full so memory cannot grow unboundedly. +/// count, including queued chunks, producer buffers, pending sends, and +/// consumer-held chunks, not just `concurrent_renders × channel_capacity`. +/// Size the pool for the measured working set. This bounds idle retention, +/// not active allocations. /// -/// `chunk_size` should match `StreamingWriter::CHUNK_TARGET + -/// BUF_HEADROOM`. When acquiring, the writer always grows the buffer -/// if the pool returned a smaller one (host code that mixes pool -/// sizes pays a one-time grow per buffer). +/// For the default writer, use `StreamingWriter::CHUNK_TARGET + 1024` +/// (5 KiB), including the writer's 1 KiB buffer headroom. For a custom +/// coalescing target, include the same headroom. This is a sizing recommendation, +/// not a hard chunk limit: large writes can grow active buffers. Capacities +/// above `chunk_size` are rejected on return, so mismatched sizes or recurring +/// oversized writes can cause repeated allocations. /// /// # Cost /// -/// * `acquire`: 1 atomic CAS (~10 ns on x86) + an `unwrap_or_else` -/// that allocates only on miss. -/// * `release`: 1 atomic CAS + drop-on-overflow. -/// * Pool storage: `max_pool * size_of::>>` = -/// ~32 bytes per slot, i.e. 512 slots = 16 KiB pool overhead. +/// * Idle buffers: summed `Vec` capacity is at most +/// `max_pool.max(1) × chunk_size`, excluding allocator overhead. +/// * Queue storage and pool metadata are additional. +/// * Active, queued, pending-send, and consumer-held buffers are additional. +/// `Bytes::from_owner` still allocates owner metadata per chunk. /// /// # Example /// @@ -114,7 +118,7 @@ use webui_handler::{FlushWriter, HandlerError, ResponseWriter, Result}; /// use webui::streaming::{ChunkPool, StreamingWriter}; /// /// // Construct ONE pool at server startup: -/// let pool = Arc::new(ChunkPool::new(512, StreamingWriter::CHUNK_TARGET)); +/// let pool = Arc::new(ChunkPool::new(512, StreamingWriter::CHUNK_TARGET + 1024)); /// /// // Each request: /// let (tx, rx) = tokio::sync::mpsc::channel(StreamingWriter::DEFAULT_CHANNEL_CAPACITY); @@ -131,12 +135,15 @@ impl ChunkPool { /// chunk buffers. /// /// `max_pool` is the maximum number of buffers held idle at once. - /// Surplus buffers are dropped (returned to the allocator) — this - /// caps total pool memory at `max_pool × chunk_size`. + /// Zero is raised to one. Surplus or oversized buffers are dropped, + /// not shrunk, bounding summed idle `Vec` capacity to + /// `max_pool.max(1) × chunk_size`. Allocator, queue, and owner metadata, + /// as well as active and consumer-held buffers, are additional. /// /// `chunk_size` is the initial capacity used when allocating a - /// fresh buffer on a pool miss. Pre-sizing avoids a Vec-grow on - /// the hot path. + /// fresh buffer on a pool miss and the maximum capacity accepted on + /// return. For the default writer, use `StreamingWriter::CHUNK_TARGET + 1024` + /// to include its buffer headroom. #[must_use] pub fn new(max_pool: usize, chunk_size: usize) -> Self { Self { @@ -146,15 +153,8 @@ impl ChunkPool { } } - /// Acquire a buffer from the pool, or allocate a fresh one if the - /// pool is empty. The returned `Vec` is empty (`len == 0`); its - /// capacity is at least `chunk_size` (may be larger if a previous - /// caller grew it). - /// - /// Trusts that callers (only [`PooledChunk::drop`] in this crate) - /// have already cleared the buffer before release. In debug builds - /// we assert the invariant; release builds skip the check to keep - /// `acquire` to a single CAS + capacity check. + // Acquired buffers are empty with capacity at least `chunk_size`. + // `release` clears returned buffers; assert rather than double-clear. fn acquire(&self) -> Vec { match self.queue.pop() { Some(mut buf) => { @@ -163,7 +163,8 @@ impl ChunkPool { "ChunkPool invariant violation: pool returned non-empty buffer" ); if buf.capacity() < self.chunk_size { - buf.reserve(self.chunk_size - buf.capacity()); + // Reserve is additional to len (zero), not capacity. + buf.reserve_exact(self.chunk_size); } buf } @@ -171,11 +172,12 @@ impl ChunkPool { } } - /// Release a buffer back to the pool. The buffer is `clear()`-ed - /// here (cheap — sets `len` to 0, no deallocation), so `acquire` - /// can trust the invariant and skip a defensive clear on the hot - /// path. Drops the buffer if the pool is full. + // Clear retained buffers so `acquire` can skip a defensive clear. + // Drop oversized returns rather than shrinking them. fn release(&self, mut buf: Vec) { + if buf.capacity() > self.chunk_size { + return; + } buf.clear(); // ArrayQueue::push returns Err with the value if full; we // simply drop in that case. @@ -189,7 +191,7 @@ impl ChunkPool { self.queue.len() } - /// Maximum buffers the pool can hold idle. + /// Maximum buffers the pool can hold idle (at least one). #[must_use] pub fn capacity(&self) -> usize { self.queue.capacity() @@ -206,8 +208,7 @@ struct PooledChunk { /// `Option` so we can `take()` the `Vec` in `Drop` and return /// it to the pool — Drop receives `&mut self`, so we can't move /// out of the field directly. Using `Option` keeps the impl - /// safe (no `ManuallyDrop` / `unsafe`) at the cost of one - /// 8-byte tag per chunk-in-flight; negligible vs the chunk size. + /// safe (no `ManuallyDrop` / `unsafe`). buf: Option>, pool: Arc, } @@ -254,7 +255,8 @@ impl Drop for PooledChunk { /// Streaming `ResponseWriter` backed by a **bounded** tokio mpsc channel /// of [`Bytes`]. /// -/// Coalesces small writes into ~4 KB chunks before flushing. The +/// Coalesces small writes into ~4 KB chunks before flushing. This is a +/// target, not a byte cap: a large write can produce a larger chunk. The /// underlying channel has a small bound /// ([`DEFAULT_CHANNEL_CAPACITY`](Self::DEFAULT_CHANNEL_CAPACITY)) so a /// slow consumer naturally backpressures the producer — the render @@ -328,9 +330,9 @@ impl StreamingWriter { /// Tunable via [`with_chunk_size`](Self::with_chunk_size). pub const CHUNK_TARGET: usize = 4 * 1024; - /// Default bounded-channel capacity in chunks. With - /// `CHUNK_TARGET = 4 KB`, this caps in-flight memory at ~16 KB per - /// in-progress request. + /// Default bounded-channel capacity in chunks, not bytes. + /// Large writes can exceed `CHUNK_TARGET`, so this does not cap + /// queued payload bytes or total in-flight memory. pub const DEFAULT_CHANNEL_CAPACITY: usize = 4; /// Minimum allowed chunk size. Below this the per-flush channel @@ -361,16 +363,17 @@ impl StreamingWriter { } /// Wrap a tokio mpsc sender, drawing chunk buffers from the - /// shared `pool`. Recycled buffers eliminate per-flush allocation - /// in steady-state high-RPS workloads. The pool is shared via - /// `Arc` and is safe to use from any number of concurrent + /// shared `pool`. Reusing suitably sized buffers avoids chunk-buffer + /// allocations; `Bytes::from_owner` still allocates owner metadata per flush. + /// The pool is shared via `Arc` and is safe to use from any number of concurrent /// `StreamingWriter` instances; release happens when the consumer /// drops the `Bytes`, on whichever thread held the last reference. /// /// `chunk_target` defaults to [`CHUNK_TARGET`](Self::CHUNK_TARGET); /// override with [`with_chunk_size`](Self::with_chunk_size). When - /// the pool's chunk size disagrees with the writer's target, the - /// writer grows the acquired buffer on first use (one-time cost). + /// sizing the pool, include the writer's 1 KiB headroom as well as + /// its target. Buffers grown beyond the pool's chunk size are dropped + /// on return, so mismatched sizes can cause repeated allocations. #[must_use] pub fn new_pooled(tx: Sender, pool: Arc) -> Self { let buf = pool.acquire(); @@ -667,6 +670,9 @@ impl FlushWriter for StreamingWriter { // ── Tests ────────────────────────────────────────────────────────── +#[cfg(test)] +mod pool_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/webui/src/streaming/pool_tests.rs b/crates/webui/src/streaming/pool_tests.rs new file mode 100644 index 00000000..1a5ca591 --- /dev/null +++ b/crates/webui/src/streaming/pool_tests.rs @@ -0,0 +1,149 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT license. + +#![allow(clippy::disallowed_methods)] + +use super::*; +use tokio::sync::mpsc::channel; + +#[test] +fn pool_reuses_normal_buffer_without_retaining_content() { + let pool = ChunkPool::new(2, 64); + let mut buf = pool.acquire(); + let allocation = buf.as_ptr(); + buf.extend_from_slice(b"previous response"); + pool.release(buf); + assert_eq!(pool.idle_count(), 1); + + let reused = pool.acquire(); + assert!(reused.is_empty()); + assert_eq!(reused.as_ptr(), allocation); + assert_eq!(reused.capacity(), 64); + assert_eq!(pool.idle_count(), 0); +} + +#[test] +fn pool_rejects_oversized_capacity_instead_of_shrinking() { + let pool = ChunkPool::new(4, 4096); + for capacity in [4097, 64 * 1024] { + let mut buf = Vec::with_capacity(capacity); + buf.extend_from_slice(b"small payload"); + pool.release(buf); + assert_eq!(pool.idle_count(), 0, "oversized capacity {capacity}"); + } + + let exact = Vec::with_capacity(4096); + let allocation = exact.as_ptr(); + pool.release(exact); + assert_eq!(pool.idle_count(), 1, "exact capacity remains reusable"); + let reused = pool.acquire(); + assert_eq!(reused.as_ptr(), allocation); +} + +#[test] +fn pool_enforces_idle_count_limit_including_zero_max_pool() { + for max_pool in [0, 1, 3] { + let pool = ChunkPool::new(max_pool, 64); + let limit = max_pool.max(1); + assert_eq!(pool.capacity(), limit); + for returned in 1..=limit + 1 { + pool.release(Vec::with_capacity(64)); + assert_eq!(pool.idle_count(), returned.min(limit)); + } + } +} + +#[test] +fn pool_acquire_reserves_from_length_after_small_return() { + let pool = ChunkPool::new(1, 4096); + for capacity in [0, 3072, 4095] { + pool.release(Vec::with_capacity(capacity)); + assert_eq!(pool.idle_count(), 1); + let acquired = pool.acquire(); + assert!(acquired.is_empty()); + assert_eq!(acquired.capacity(), 4096); + } +} + +#[test] +fn pool_bounds_sum_of_idle_capacities() { + let chunk_size = 4096; + let pool = ChunkPool::new(3, chunk_size); + for capacity in [ + chunk_size + 1, + chunk_size / 2, + chunk_size * 8, + chunk_size, + chunk_size - 1, + chunk_size, + ] { + pool.release(Vec::with_capacity(capacity)); + } + assert_eq!(pool.idle_count(), pool.capacity()); + + let mut retained = 0; + while let Some(buf) = pool.queue.pop() { + assert!(buf.is_empty()); + assert!(buf.capacity() <= chunk_size); + retained += buf.capacity(); + } + assert!(retained <= pool.capacity() * chunk_size); +} + +#[test] +fn pooled_chunk_recycles_only_after_final_clone_or_slice_drop() { + let pool = Arc::new(ChunkPool::new(2, 64)); + let mut buf = pool.acquire(); + buf.extend_from_slice(b"consumer-owned"); + let allocation = buf.as_ptr(); + let original = Bytes::from_owner(PooledChunk::new(buf, Arc::clone(&pool))); + let clone = original.clone(); + let slice = clone.slice(9..); + assert_eq!(pool.idle_count(), 0); + drop(original); + assert_eq!(pool.idle_count(), 0); + drop(clone); + assert_eq!(pool.idle_count(), 0); + assert_eq!(slice.as_ref(), b"owned"); + + std::thread::spawn(move || drop(slice)).join().unwrap(); + assert_eq!(pool.idle_count(), 1); + let recycled = pool.acquire(); + assert_eq!(recycled.as_ptr(), allocation); + assert!(recycled.is_empty()); +} + +#[test] +fn large_write_remains_one_chunk_and_is_not_retained_after_last_clone() { + let chunk_size = StreamingWriter::CHUNK_TARGET + StreamingWriter::BUF_HEADROOM; + let pool = Arc::new(ChunkPool::new(4, chunk_size)); + let (tx, mut rx) = channel(8); + let mut writer = StreamingWriter::new_pooled(tx, Arc::clone(&pool)); + assert_eq!(writer.buf.capacity(), chunk_size); + let large = "x".repeat(chunk_size * 3); + writer.write("prefix").unwrap(); + writer.write(&large).unwrap(); + writer.end().unwrap(); + assert_eq!(rx.len(), 1, "the coalescing target is not a byte cap"); + + let original = rx.try_recv().unwrap(); + assert_eq!(&original[..6], b"prefix"); + assert_eq!(&original[6..], large.as_bytes()); + assert!(original.len() > chunk_size); + let clone = original.clone(); + drop(original); + assert_eq!(pool.idle_count(), 0); + assert_eq!(&clone[6..], large.as_bytes()); + drop(clone); + assert_eq!(pool.idle_count(), 0, "grown allocation must be dropped"); + + writer.write("small").unwrap(); + writer.end().unwrap(); + let small = rx.try_recv().unwrap(); + assert_eq!(small.as_ref(), b"small"); + drop(small); + assert_eq!(pool.idle_count(), 1, "normal chunks still recycle"); + drop(writer); + assert_eq!(pool.idle_count(), 2, "writer's active buffer also returns"); + assert!(rx.try_recv().is_err()); +} diff --git a/docs/guide/concepts/performance.md b/docs/guide/concepts/performance.md index ed966c99..d48a52df 100644 --- a/docs/guide/concepts/performance.md +++ b/docs/guide/concepts/performance.md @@ -153,6 +153,17 @@ Node, Bun, Deno, Python, and other bindings remain useful when integration cost matters more than the last increment of throughput. Measure with the deployment host you intend to run. +## Bound idle streaming buffers + +For Rust transport streaming, a shared `ChunkPool` reuses eligible chunk buffers. +Its summed idle `Vec` capacity is bounded by `max_pool.max(1) * chunk_size`; +oversized returns are dropped rather than shrunk. Active and consumer-held +buffers, allocator/queue metadata, and the owner metadata allocated by +`Bytes::from_owner` are additional. Repeated oversized writes can therefore +trade extra allocations for bounded idle retention. Keep the default pool size +at the 4 KiB coalescing target plus 1 KiB headroom; the target is not a hard +chunk limit. See [Rust chunk pool sizing](/guide/integrations/rust#chunk-pool-sizing). + ## Stream slow data, not already-fast pages Progressive boundaries help when data dependencies complete at different diff --git a/docs/guide/integrations/rust.md b/docs/guide/integrations/rust.md index 1ad2b35f..af10cc2f 100644 --- a/docs/guide/integrations/rust.md +++ b/docs/guide/integrations/rust.md @@ -211,7 +211,7 @@ use webui::{WebUIHandler, RenderOptions, ResponseWriter}; // One shared pool per server (constructed at startup, lives forever). let chunk_pool = Arc::new(ChunkPool::new( - 256, // ~1.25 MiB peak pool memory + 256, // at most 1.25 MiB idle buffer capacity StreamingWriter::CHUNK_TARGET + 1024, )); let render_permits = Arc::new(Semaphore::new(4)); @@ -252,6 +252,25 @@ HttpResponse::Ok() .streaming(tokio_stream::wrappers::ReceiverStream::new(rx).map(Ok::<_, actix_web::Error>)) ``` +### Chunk pool sizing + +`ChunkPool::new(max_pool, chunk_size)` retains at most `max_pool.max(1)` idle +buffers whose summed `Vec` capacity is at most `max_pool.max(1) * chunk_size`. +Oversized returned buffers are dropped, not shrunk. The example's 5 KiB size +includes the default 4 KiB coalescing target plus 1 KiB of writer headroom; +include the same headroom when selecting a custom target. + +This is an idle-capacity bound, not a total server-memory limit. Allocator, +queue, and owner metadata are extra, as are active writer buffers, pending +sends, queued chunks, and consumer-held buffers. A buffer is eligible for reuse +only after the last `Bytes` reference (including clones and slices) drops. +`Bytes::from_owner` still allocates owner metadata per chunk. Mismatched pool +sizes or recurring oversized writes can cause repeated buffer allocations. + +The writer's coalescing target is not a byte cap: large writes remain whole +and can produce larger chunks. The bounded channel limits chunk count, not +queued payload bytes. + ### Host-driven boundaries and state updates `stream_response` returns a synchronous session that discovers runtime