Skip to main content

nntp_proxy/pool/
buffer.rs

1use crate::types::BufferSize;
2use bytes::{BufMut, Bytes, BytesMut};
3use crossbeam::queue::ArrayQueue;
4use smallvec::SmallVec;
5use std::future::poll_fn;
6use std::ops::{Deref, Range};
7use std::pin::Pin;
8use std::sync::Arc;
9use std::sync::OnceLock;
10use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
11use std::time::Instant;
12use tokio::io::{AsyncRead, AsyncWriteExt, ReadBuf};
13use tracing::{debug, info, warn};
14
15/// A pooled buffer that automatically returns to the pool when dropped
16///
17/// # Safety and Initialized Bytes
18///
19/// The backing allocation is kept logically empty until bytes are read or copied
20/// into it. `BytesMut::len()` is the initialized length, and immutable access
21/// only exposes that initialized portion.
22///
23/// ## Usage
24/// ```ignore
25/// let mut buffer = pool.acquire();
26/// let n = buffer.read_from(&mut stream).await?;  // Automatic tracking
27/// process(&*buffer);  // Deref returns only &buffer[..n]
28/// ```
29pub struct PooledBuffer {
30    buffer: BytesMut,
31    pool: Arc<ArrayQueue<BytesMut>>,
32    pool_size: Arc<AtomicUsize>,
33    allocated_count: Arc<AtomicUsize>,
34    max_pool_size: usize,
35    expected_capacity: usize,
36    writable_len: usize,
37    acquired_at: Instant,
38    source: PooledBufferSource,
39    counts_toward_pool: bool,
40}
41
42#[derive(Clone, Copy, Debug, PartialEq, Eq)]
43enum PooledBufferKind {
44    Regular,
45    Capture,
46}
47
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49struct PooledBufferSource {
50    kind: PooledBufferKind,
51    fallback: bool,
52}
53
54impl PooledBufferSource {
55    const fn regular(fallback: bool) -> Self {
56        Self {
57            kind: PooledBufferKind::Regular,
58            fallback,
59        }
60    }
61
62    const fn capture(fallback: bool) -> Self {
63        Self {
64            kind: PooledBufferKind::Capture,
65            fallback,
66        }
67    }
68}
69
70impl std::fmt::Debug for PooledBuffer {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        f.debug_struct("PooledBuffer")
73            .field("initialized", &self.buffer.len())
74            .field("capacity", &self.buffer.capacity())
75            .finish_non_exhaustive()
76    }
77}
78
79impl PooledBuffer {
80    /// Get the backing allocation capacity of the buffer.
81    #[must_use]
82    #[inline]
83    pub fn capacity(&self) -> usize {
84        self.buffer.capacity()
85    }
86
87    /// Get the number of initialized bytes
88    #[must_use]
89    #[inline]
90    pub fn initialized(&self) -> usize {
91        self.buffer.len()
92    }
93
94    #[inline]
95    fn read_limit(&self) -> usize {
96        if self.writable_len == 0 {
97            self.buffer.capacity()
98        } else {
99            self.writable_len
100        }
101    }
102
103    /// Read from an `AsyncRead` source, automatically tracking initialized bytes
104    ///
105    /// # Errors
106    /// Returns any read error produced by the underlying async reader.
107    pub async fn read_from<R>(&mut self, reader: &mut R) -> std::io::Result<usize>
108    where
109        R: AsyncRead + Unpin,
110    {
111        self.buffer.clear();
112        self.read_into_spare(reader, self.read_limit()).await
113    }
114
115    /// Read more data at the current initialized offset, accumulating bytes.
116    ///
117    /// Unlike `read_from` which resets the buffer, this appends to existing data.
118    /// Used when a higher-level reader needs more bytes in the same allocation.
119    ///
120    /// Returns the number of NEW bytes read (not total).
121    ///
122    /// # Errors
123    /// Returns any read error produced by the underlying async reader.
124    pub async fn read_more<R>(&mut self, reader: &mut R) -> std::io::Result<usize>
125    where
126        R: AsyncRead + Unpin,
127    {
128        let limit = self.read_limit().saturating_sub(self.buffer.len());
129        self.read_into_spare(reader, limit).await
130    }
131
132    async fn read_into_spare<R>(&mut self, reader: &mut R, limit: usize) -> std::io::Result<usize>
133    where
134        R: AsyncRead + Unpin,
135    {
136        let spare = self.buffer.spare_capacity_mut();
137        let read_len = limit.min(spare.len());
138        if read_len == 0 {
139            return Ok(0);
140        }
141
142        let mut read_buf = ReadBuf::uninit(&mut spare[..read_len]);
143        poll_fn(|cx| Pin::new(&mut *reader).poll_read(cx, &mut read_buf)).await?;
144        let read = read_buf.filled().len();
145        // SAFETY: `poll_read` initialized exactly `read` bytes in the spare slice.
146        unsafe {
147            self.buffer.advance_mut(read);
148        }
149        Ok(read)
150    }
151
152    /// Copy data into buffer and mark as initialized
153    ///
154    /// # Panics
155    /// Panics if `data.len()` > capacity
156    #[inline]
157    pub fn copy_from_slice(&mut self, data: &[u8]) {
158        assert!(
159            data.len() <= self.buffer.capacity(),
160            "data exceeds buffer capacity"
161        );
162        self.buffer.clear();
163        self.buffer.extend_from_slice(data);
164    }
165
166    /// Reset the logical contents of the buffer without changing its backing allocation.
167    ///
168    /// This is safe for both fixed-size I/O buffers and accumulator-style capture buffers:
169    /// callers see an empty initialized slice afterwards, while the underlying allocation
170    /// stays available for reuse.
171    #[inline]
172    pub fn clear(&mut self) {
173        self.buffer.clear();
174    }
175
176    /// Append data to the buffer (accumulator mode)
177    ///
178    /// Used when `PooledBuffer` is acquired from capture pool for accumulating
179    /// streaming data. Unlike `copy_from_slice` which overwrites, this extends.
180    ///
181    /// # Performance
182    ///
183    /// Capture buffers are allocated with sufficient capacity (1 MiB by default).
184    /// As long as total accumulated data fits within capacity, this operation performs
185    /// no Rust heap allocation.
186    ///
187    /// If data exceeds capacity, the `BytesMut` will grow automatically (with allocation).
188    /// This ensures complete data is always captured rather than truncating, which
189    /// would corrupt cached responses.
190    ///
191    /// # Note on length
192    ///
193    /// After calling this, `.len()` (via Deref) returns the total accumulated length.
194    #[inline]
195    pub fn extend_from_slice(&mut self, data: &[u8]) {
196        let needed = self.buffer.len() + data.len();
197        // Accumulator/capture mode may grow; resized buffers are discarded on drop.
198        if needed > self.buffer.capacity() {
199            tracing::warn!(
200                "Capture buffer growing: {} + {} > {} capacity (will allocate)",
201                self.buffer.len(),
202                data.len(),
203                self.buffer.capacity()
204            );
205        }
206        self.buffer.extend_from_slice(data);
207    }
208
209    /// Freeze the initialized bytes into a shared immutable buffer.
210    ///
211    /// The backing allocation is intentionally detached from the pool because
212    /// ownership escapes through `Bytes`.
213    #[must_use]
214    pub fn freeze(mut self) -> Bytes {
215        std::mem::take(&mut self.buffer).freeze()
216    }
217}
218
219impl Deref for PooledBuffer {
220    type Target = [u8];
221
222    #[inline]
223    fn deref(&self) -> &Self::Target {
224        &self.buffer
225    }
226}
227
228// Intentionally NO DerefMut or AsMut - forces explicit use of read_from()/read_more()
229
230impl AsRef<[u8]> for PooledBuffer {
231    #[inline]
232    fn as_ref(&self) -> &[u8] {
233        &self.buffer
234    }
235}
236
237impl Drop for PooledBuffer {
238    fn drop(&mut self) {
239        let held_micros = duration_micros_u64(self.acquired_at.elapsed());
240        record_buffer_hold(self.source, held_micros);
241        maybe_log_buffer_hold(self.source, held_micros, self.buffer.capacity());
242
243        if self.buffer.capacity() != self.expected_capacity {
244            debug!(
245                "Dropping resized pooled buffer instead of returning it to the pool (capacity {} != expected {})",
246                self.buffer.capacity(),
247                self.expected_capacity
248            );
249            if self.counts_toward_pool {
250                self.allocated_count.fetch_sub(1, Ordering::Relaxed);
251            }
252            return;
253        }
254
255        self.buffer.clear();
256
257        if !self.counts_toward_pool {
258            return;
259        }
260
261        // Atomically return buffer to pool if pool is not full
262        let mut current_size = self.pool_size.load(Ordering::Relaxed);
263        while current_size < self.max_pool_size {
264            match self.pool_size.compare_exchange_weak(
265                current_size,
266                current_size + 1,
267                Ordering::Relaxed,
268                Ordering::Relaxed,
269            ) {
270                Ok(_) => {
271                    let buffer = std::mem::take(&mut self.buffer);
272                    match self.pool.push(buffer) {
273                        Ok(()) => return,
274                        Err(buffer) => {
275                            self.buffer = buffer;
276                            self.pool_size.fetch_sub(1, Ordering::Relaxed);
277                            self.allocated_count.fetch_sub(1, Ordering::Relaxed);
278                            return;
279                        }
280                    }
281                }
282                Err(new_size) => {
283                    current_size = new_size;
284                }
285            }
286        }
287        // If pool is full, buffer is dropped
288        self.allocated_count.fetch_sub(1, Ordering::Relaxed);
289    }
290}
291
292/// Buffered response assembled from one or more pooled capture buffers.
293///
294/// This avoids reallocating or growing a single `Vec` on the hot path when
295/// a multiline response is larger than the typical capture size.
296///
297/// `ChunkedResponse` is storage only. It intentionally does not expose response
298/// boundary helpers such as prefix or terminator checks; callers that need to
299/// interpret backend bytes must go through the session framer/backend facade
300/// before putting data here.
301#[derive(Debug, Default)]
302pub struct ChunkedResponse {
303    chunks: SmallVec<[ResponseChunk; 16]>,
304    len: usize,
305}
306
307#[derive(Debug, Default)]
308struct ResponseWriteMetrics {
309    responses: AtomicUsize,
310    single_chunk_responses: AtomicUsize,
311    multi_chunk_responses: AtomicUsize,
312    chunks_written: AtomicUsize,
313    bytes_written: AtomicUsize,
314    tiny_chunks: AtomicUsize,
315    tiny_chunk_bytes: AtomicUsize,
316    small_chunks: AtomicUsize,
317    small_chunk_bytes: AtomicUsize,
318    max_chunks_per_response: AtomicUsize,
319}
320
321#[derive(Debug, Default)]
322struct HotPathAllocationMetrics {
323    regular_pool_fallback_allocations: AtomicUsize,
324    regular_pool_exhaustions: AtomicUsize,
325    capture_pool_fallback_allocations: AtomicUsize,
326    chunked_response_metadata_spills: AtomicUsize,
327    pending_backend_byte_heap_fallbacks: AtomicUsize,
328    non_owned_response_write_chunks: AtomicUsize,
329    non_owned_response_write_bytes: AtomicUsize,
330    regular_pool_buffer_holds: AtomicUsize,
331    regular_pool_buffer_hold_micros_total: AtomicU64,
332    regular_pool_buffer_hold_micros_max: AtomicU64,
333    capture_pool_buffer_holds: AtomicUsize,
334    capture_pool_buffer_hold_micros_total: AtomicU64,
335    capture_pool_buffer_hold_micros_max: AtomicU64,
336}
337
338/// Snapshot of direct client write-path activity from `ChunkedResponse::write_all_to`.
339#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
340pub struct ResponseWriteMetricsSnapshot {
341    pub responses: usize,
342    pub single_chunk_responses: usize,
343    pub multi_chunk_responses: usize,
344    pub chunks_written: usize,
345    pub bytes_written: usize,
346    pub tiny_chunks: usize,
347    pub tiny_chunk_bytes: usize,
348    pub small_chunks: usize,
349    pub small_chunk_bytes: usize,
350    pub max_chunks_per_response: usize,
351}
352
353/// Snapshot of allocation-sensitive hot-path activity.
354#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
355pub struct HotPathAllocationMetricsSnapshot {
356    pub regular_pool_fallback_allocations: usize,
357    pub regular_pool_exhaustions: usize,
358    pub capture_pool_fallback_allocations: usize,
359    pub chunked_response_metadata_spills: usize,
360    pub pending_backend_byte_heap_fallbacks: usize,
361    pub non_owned_response_write_chunks: usize,
362    pub non_owned_response_write_bytes: usize,
363    pub regular_pool_buffer_holds: usize,
364    pub regular_pool_buffer_hold_micros_total: u64,
365    pub regular_pool_buffer_hold_micros_max: u64,
366    pub capture_pool_buffer_holds: usize,
367    pub capture_pool_buffer_hold_micros_total: u64,
368    pub capture_pool_buffer_hold_micros_max: u64,
369}
370
371fn response_write_metrics_enabled() -> bool {
372    static ENABLED: OnceLock<bool> = OnceLock::new();
373    *ENABLED.get_or_init(|| std::env::var_os("NNTP_PROXY_RESPONSE_WRITE_METRICS_SECS").is_some())
374}
375
376fn response_write_metrics() -> &'static ResponseWriteMetrics {
377    static METRICS: OnceLock<ResponseWriteMetrics> = OnceLock::new();
378    METRICS.get_or_init(ResponseWriteMetrics::default)
379}
380
381fn hot_path_allocation_metrics() -> &'static HotPathAllocationMetrics {
382    static METRICS: OnceLock<HotPathAllocationMetrics> = OnceLock::new();
383    METRICS.get_or_init(HotPathAllocationMetrics::default)
384}
385
386fn duration_micros_u64(duration: std::time::Duration) -> u64 {
387    u64::try_from(duration.as_micros()).unwrap_or(u64::MAX)
388}
389
390fn buffer_hold_log_threshold_micros() -> Option<u64> {
391    static THRESHOLD: OnceLock<Option<u64>> = OnceLock::new();
392    *THRESHOLD.get_or_init(|| {
393        std::env::var("NNTP_PROXY_BUFFER_HOLD_LOG_MS")
394            .ok()
395            .and_then(|value| {
396                let millis = value.parse::<u64>().ok()?;
397                if millis == 0 {
398                    None
399                } else {
400                    Some(millis.saturating_mul(1000))
401                }
402            })
403    })
404}
405
406fn record_buffer_hold(source: PooledBufferSource, held_micros: u64) {
407    let metrics = hot_path_allocation_metrics();
408    match source.kind {
409        PooledBufferKind::Regular => {
410            metrics
411                .regular_pool_buffer_holds
412                .fetch_add(1, Ordering::Relaxed);
413            metrics
414                .regular_pool_buffer_hold_micros_total
415                .fetch_add(held_micros, Ordering::Relaxed);
416            metrics
417                .regular_pool_buffer_hold_micros_max
418                .fetch_max(held_micros, Ordering::Relaxed);
419        }
420        PooledBufferKind::Capture => {
421            metrics
422                .capture_pool_buffer_holds
423                .fetch_add(1, Ordering::Relaxed);
424            metrics
425                .capture_pool_buffer_hold_micros_total
426                .fetch_add(held_micros, Ordering::Relaxed);
427            metrics
428                .capture_pool_buffer_hold_micros_max
429                .fetch_max(held_micros, Ordering::Relaxed);
430        }
431    }
432}
433
434fn maybe_log_buffer_hold(source: PooledBufferSource, held_micros: u64, capacity: usize) {
435    let Some(threshold) = buffer_hold_log_threshold_micros() else {
436        return;
437    };
438    if held_micros < threshold {
439        return;
440    }
441
442    let pool = match source.kind {
443        PooledBufferKind::Regular => "regular",
444        PooledBufferKind::Capture => "capture",
445    };
446    debug!(
447        pool,
448        fallback = source.fallback,
449        held_ms = held_micros / 1000,
450        capacity,
451        "Pooled buffer held longer than threshold"
452    );
453}
454
455pub(crate) fn record_pending_backend_byte_heap_fallback() {
456    hot_path_allocation_metrics()
457        .pending_backend_byte_heap_fallbacks
458        .fetch_add(1, Ordering::Relaxed);
459}
460
461pub(crate) fn record_non_owned_response_write_chunk(bytes: usize) {
462    let metrics = hot_path_allocation_metrics();
463    metrics
464        .non_owned_response_write_chunks
465        .fetch_add(1, Ordering::Relaxed);
466    metrics
467        .non_owned_response_write_bytes
468        .fetch_add(bytes, Ordering::Relaxed);
469}
470
471#[must_use]
472pub fn hot_path_allocation_metrics_snapshot() -> HotPathAllocationMetricsSnapshot {
473    let metrics = hot_path_allocation_metrics();
474    HotPathAllocationMetricsSnapshot {
475        regular_pool_fallback_allocations: metrics
476            .regular_pool_fallback_allocations
477            .load(Ordering::Relaxed),
478        regular_pool_exhaustions: metrics.regular_pool_exhaustions.load(Ordering::Relaxed),
479        capture_pool_fallback_allocations: metrics
480            .capture_pool_fallback_allocations
481            .load(Ordering::Relaxed),
482        chunked_response_metadata_spills: metrics
483            .chunked_response_metadata_spills
484            .load(Ordering::Relaxed),
485        pending_backend_byte_heap_fallbacks: metrics
486            .pending_backend_byte_heap_fallbacks
487            .load(Ordering::Relaxed),
488        non_owned_response_write_chunks: metrics
489            .non_owned_response_write_chunks
490            .load(Ordering::Relaxed),
491        non_owned_response_write_bytes: metrics
492            .non_owned_response_write_bytes
493            .load(Ordering::Relaxed),
494        regular_pool_buffer_holds: metrics.regular_pool_buffer_holds.load(Ordering::Relaxed),
495        regular_pool_buffer_hold_micros_total: metrics
496            .regular_pool_buffer_hold_micros_total
497            .load(Ordering::Relaxed),
498        regular_pool_buffer_hold_micros_max: metrics
499            .regular_pool_buffer_hold_micros_max
500            .load(Ordering::Relaxed),
501        capture_pool_buffer_holds: metrics.capture_pool_buffer_holds.load(Ordering::Relaxed),
502        capture_pool_buffer_hold_micros_total: metrics
503            .capture_pool_buffer_hold_micros_total
504            .load(Ordering::Relaxed),
505        capture_pool_buffer_hold_micros_max: metrics
506            .capture_pool_buffer_hold_micros_max
507            .load(Ordering::Relaxed),
508    }
509}
510
511pub fn reset_hot_path_allocation_metrics() {
512    let metrics = hot_path_allocation_metrics();
513    metrics
514        .regular_pool_fallback_allocations
515        .store(0, Ordering::Relaxed);
516    metrics.regular_pool_exhaustions.store(0, Ordering::Relaxed);
517    metrics
518        .capture_pool_fallback_allocations
519        .store(0, Ordering::Relaxed);
520    metrics
521        .chunked_response_metadata_spills
522        .store(0, Ordering::Relaxed);
523    metrics
524        .pending_backend_byte_heap_fallbacks
525        .store(0, Ordering::Relaxed);
526    metrics
527        .non_owned_response_write_chunks
528        .store(0, Ordering::Relaxed);
529    metrics
530        .non_owned_response_write_bytes
531        .store(0, Ordering::Relaxed);
532    metrics
533        .regular_pool_buffer_holds
534        .store(0, Ordering::Relaxed);
535    metrics
536        .regular_pool_buffer_hold_micros_total
537        .store(0, Ordering::Relaxed);
538    metrics
539        .regular_pool_buffer_hold_micros_max
540        .store(0, Ordering::Relaxed);
541    metrics
542        .capture_pool_buffer_holds
543        .store(0, Ordering::Relaxed);
544    metrics
545        .capture_pool_buffer_hold_micros_total
546        .store(0, Ordering::Relaxed);
547    metrics
548        .capture_pool_buffer_hold_micros_max
549        .store(0, Ordering::Relaxed);
550}
551
552fn record_response_write_metrics(
553    chunk_count: usize,
554    bytes_written: usize,
555    tiny_chunks: usize,
556    tiny_chunk_bytes: usize,
557    small_chunks: usize,
558    small_chunk_bytes: usize,
559) {
560    if !response_write_metrics_enabled() {
561        return;
562    }
563
564    let metrics = response_write_metrics();
565    metrics.responses.fetch_add(1, Ordering::Relaxed);
566    metrics
567        .chunks_written
568        .fetch_add(chunk_count, Ordering::Relaxed);
569    metrics
570        .bytes_written
571        .fetch_add(bytes_written, Ordering::Relaxed);
572    metrics
573        .tiny_chunks
574        .fetch_add(tiny_chunks, Ordering::Relaxed);
575    metrics
576        .tiny_chunk_bytes
577        .fetch_add(tiny_chunk_bytes, Ordering::Relaxed);
578    metrics
579        .small_chunks
580        .fetch_add(small_chunks, Ordering::Relaxed);
581    metrics
582        .small_chunk_bytes
583        .fetch_add(small_chunk_bytes, Ordering::Relaxed);
584    metrics
585        .max_chunks_per_response
586        .fetch_max(chunk_count, Ordering::Relaxed);
587    if chunk_count <= 1 {
588        metrics
589            .single_chunk_responses
590            .fetch_add(1, Ordering::Relaxed);
591    } else {
592        metrics
593            .multi_chunk_responses
594            .fetch_add(1, Ordering::Relaxed);
595    }
596}
597
598#[must_use]
599pub fn response_write_metrics_snapshot() -> ResponseWriteMetricsSnapshot {
600    let metrics = response_write_metrics();
601    ResponseWriteMetricsSnapshot {
602        responses: metrics.responses.load(Ordering::Relaxed),
603        single_chunk_responses: metrics.single_chunk_responses.load(Ordering::Relaxed),
604        multi_chunk_responses: metrics.multi_chunk_responses.load(Ordering::Relaxed),
605        chunks_written: metrics.chunks_written.load(Ordering::Relaxed),
606        bytes_written: metrics.bytes_written.load(Ordering::Relaxed),
607        tiny_chunks: metrics.tiny_chunks.load(Ordering::Relaxed),
608        tiny_chunk_bytes: metrics.tiny_chunk_bytes.load(Ordering::Relaxed),
609        small_chunks: metrics.small_chunks.load(Ordering::Relaxed),
610        small_chunk_bytes: metrics.small_chunk_bytes.load(Ordering::Relaxed),
611        max_chunks_per_response: metrics.max_chunks_per_response.load(Ordering::Relaxed),
612    }
613}
614
615#[derive(Debug)]
616struct ResponseChunk {
617    storage: ResponseChunkStorage,
618    range: Range<usize>,
619}
620
621#[derive(Debug)]
622enum ResponseChunkStorage {
623    Pooled(PooledBuffer),
624    #[cfg(test)]
625    SharedPooled(Arc<PooledBuffer>),
626}
627
628impl ResponseChunk {
629    fn as_slice(&self) -> &[u8] {
630        match &self.storage {
631            ResponseChunkStorage::Pooled(buffer) => &buffer.as_ref()[self.range.clone()],
632            #[cfg(test)]
633            ResponseChunkStorage::SharedPooled(buffer) => &buffer.as_ref()[self.range.clone()],
634        }
635    }
636
637    fn initialized(&self) -> usize {
638        match &self.storage {
639            ResponseChunkStorage::Pooled(buffer) => buffer.initialized(),
640            #[cfg(test)]
641            ResponseChunkStorage::SharedPooled(buffer) => buffer.initialized(),
642        }
643    }
644
645    fn capacity(&self) -> usize {
646        match &self.storage {
647            ResponseChunkStorage::Pooled(buffer) => buffer.capacity(),
648            #[cfg(test)]
649            ResponseChunkStorage::SharedPooled(buffer) => buffer.capacity(),
650        }
651    }
652
653    fn can_append(&self) -> bool {
654        matches!(self.storage, ResponseChunkStorage::Pooled(_))
655            && self.range.end == self.initialized()
656            && self.initialized() < self.capacity()
657    }
658
659    fn extend_from_slice(&mut self, data: &[u8]) {
660        match &mut self.storage {
661            ResponseChunkStorage::Pooled(buffer) => buffer.extend_from_slice(data),
662            #[cfg(test)]
663            ResponseChunkStorage::SharedPooled(_) => {
664                panic!("cannot extend shared response chunk")
665            }
666        }
667    }
668}
669
670impl ChunkedResponse {
671    /// Total buffered length across all chunks.
672    #[must_use]
673    #[inline]
674    pub const fn len(&self) -> usize {
675        self.len
676    }
677
678    /// Returns true when no bytes are buffered.
679    #[must_use]
680    #[inline]
681    pub const fn is_empty(&self) -> bool {
682        self.len == 0
683    }
684
685    /// Remove all buffered data, returning chunks to the pool on drop.
686    #[inline]
687    pub fn clear(&mut self) {
688        self.chunks.clear();
689        self.len = 0;
690    }
691
692    /// Append bytes using one or more pooled capture buffers.
693    ///
694    /// # Panics
695    /// Panics if the internal chunk list becomes unexpectedly empty after a
696    /// new capture buffer is acquired.
697    pub fn extend_from_slice(&mut self, pool: &BufferPool, mut data: &[u8]) {
698        while !data.is_empty() {
699            let need_new_chunk = self.chunks.last().is_none_or(|chunk| !chunk.can_append());
700
701            if need_new_chunk {
702                if self.chunks.len() == self.chunks.inline_size() {
703                    hot_path_allocation_metrics()
704                        .chunked_response_metadata_spills
705                        .fetch_add(1, Ordering::Relaxed);
706                }
707                self.chunks.push(ResponseChunk {
708                    storage: ResponseChunkStorage::Pooled(pool.acquire_capture_now()),
709                    range: 0..0,
710                });
711            }
712
713            let chunk = self
714                .chunks
715                .last_mut()
716                .expect("chunk just pushed or already existed");
717            let available = chunk.capacity().saturating_sub(chunk.initialized());
718            debug_assert!(available > 0, "chunk must have space after allocation");
719
720            let take = available.min(data.len());
721            chunk.extend_from_slice(&data[..take]);
722            chunk.range.end += take;
723            self.len += take;
724            data = &data[take..];
725        }
726    }
727
728    /// Move a byte range from an initialized pooled buffer into this response.
729    ///
730    /// This avoids copying large backend reads into separate capture buffers.
731    /// The caller must pass only a range already produced by response framing.
732    /// Later appends will allocate a fresh pooled chunk when this range does not
733    /// end at the initialized length, so skipped bytes are never exposed.
734    /// # Panics
735    /// Panics if `range` is outside the buffer's initialized bytes.
736    pub(crate) fn push_buffer_range(&mut self, buffer: PooledBuffer, range: Range<usize>) {
737        assert!(
738            range.start <= range.end && range.end <= buffer.initialized(),
739            "response range must be inside initialized buffer"
740        );
741        if self.chunks.len() == self.chunks.inline_size() {
742            hot_path_allocation_metrics()
743                .chunked_response_metadata_spills
744                .fetch_add(1, Ordering::Relaxed);
745        }
746        self.len += range.end - range.start;
747        self.chunks.push(ResponseChunk {
748            storage: ResponseChunkStorage::Pooled(buffer),
749            range,
750        });
751    }
752
753    /// Add a range from a shared pooled buffer without copying.
754    ///
755    /// # Panics
756    /// Panics if `range` is outside the pooled buffer's initialized bytes.
757    #[cfg(test)]
758    pub(crate) fn push_shared_pooled_range(
759        &mut self,
760        buffer: Arc<PooledBuffer>,
761        range: Range<usize>,
762    ) {
763        assert!(
764            range.start <= range.end && range.end <= buffer.initialized(),
765            "response range must be inside initialized pooled buffer"
766        );
767        if self.chunks.len() == self.chunks.inline_size() {
768            hot_path_allocation_metrics()
769                .chunked_response_metadata_spills
770                .fetch_add(1, Ordering::Relaxed);
771        }
772        self.len += range.end - range.start;
773        self.chunks.push(ResponseChunk {
774            storage: ResponseChunkStorage::SharedPooled(buffer),
775            range,
776        });
777    }
778
779    /// First chunk of buffered data, if any.
780    #[cfg(test)]
781    #[must_use]
782    pub fn first_chunk(&self) -> Option<&[u8]> {
783        self.chunks.first().map(ResponseChunk::as_slice)
784    }
785
786    /// Iterate buffered chunks without flattening.
787    pub fn iter_chunks(&self) -> impl Iterator<Item = &[u8]> {
788        self.chunks.iter().map(ResponseChunk::as_slice)
789    }
790
791    /// Copy up to `len` bytes from the front of the response into a small stack-backed buffer.
792    pub fn copy_prefix_into(&self, len: usize, out: &mut SmallVec<[u8; 128]>) {
793        out.clear();
794        let mut remaining = len.min(self.len);
795        for chunk in self.iter_chunks() {
796            if remaining == 0 {
797                break;
798            }
799            let take = remaining.min(chunk.len());
800            out.extend_from_slice(&chunk[..take]);
801            remaining -= take;
802        }
803    }
804
805    /// Copy the buffered response into a contiguous `Vec<u8>`.
806    #[must_use]
807    pub fn to_vec(&self) -> Vec<u8> {
808        let mut out = Vec::with_capacity(self.len);
809        for chunk in self.iter_chunks() {
810            out.extend_from_slice(chunk);
811        }
812        out
813    }
814
815    /// Write all buffered chunks to a sink in order.
816    ///
817    /// # Errors
818    /// Returns any write error produced by the underlying async sink.
819    pub async fn write_all_to<W>(&self, writer: &mut W) -> std::io::Result<()>
820    where
821        W: AsyncWriteExt + Unpin,
822    {
823        if response_write_metrics_enabled() {
824            let mut chunk_count = 0usize;
825            let mut tiny_chunks = 0usize;
826            let mut tiny_chunk_bytes = 0usize;
827            let mut small_chunks = 0usize;
828            let mut small_chunk_bytes = 0usize;
829            for chunk in self.iter_chunks() {
830                chunk_count += 1;
831                let len = chunk.len();
832                if len <= 256 {
833                    tiny_chunks += 1;
834                    tiny_chunk_bytes += len;
835                }
836                if len <= 4096 {
837                    small_chunks += 1;
838                    small_chunk_bytes += len;
839                }
840            }
841            record_response_write_metrics(
842                chunk_count,
843                self.len,
844                tiny_chunks,
845                tiny_chunk_bytes,
846                small_chunks,
847                small_chunk_bytes,
848            );
849        }
850
851        for chunk in self.iter_chunks() {
852            writer.write_all(chunk).await?;
853        }
854        Ok(())
855    }
856}
857
858/// Lock-free buffer pool for reusing large I/O buffers.
859///
860/// Uses bounded `ArrayQueue`s so queue slots are allocated once at startup and
861/// buffer return does not allocate linked queue blocks on the forwarding path.
862#[derive(Debug, Clone)]
863pub struct BufferPool {
864    pool: Arc<ArrayQueue<BytesMut>>,
865    buffer_size: BufferSize,
866    max_pool_size: usize,
867    pool_size: Arc<AtomicUsize>,
868    allocated_count: Arc<AtomicUsize>,
869    // Capture buffer pool (for accumulating streaming data)
870    capture_pool: Arc<ArrayQueue<BytesMut>>,
871    capture_capacity: usize,
872    max_capture_pool_size: usize,
873    capture_pool_size: Arc<AtomicUsize>,
874    capture_allocated_count: Arc<AtomicUsize>,
875}
876
877impl BufferPool {
878    /// Create a page-aligned buffer for optimal DMA performance
879    ///
880    /// Returns a raw `BytesMut` that will be wrapped in `PooledBuffer` by `acquire()`.
881    /// The buffer has a logical length of zero before it enters the pool.
882    ///
883    /// # Safety
884    ///
885    /// **INTERNAL USE ONLY.** This function is not exposed publicly and is only used
886    /// within the buffer pool implementation where the safety contract is guaranteed.
887    ///
888    /// The returned buffer has a logical length of zero. Only bytes added to the
889    /// `BytesMut` logical length are exposed through the public initialized slice APIs.
890    fn create_aligned_buffer(size: usize) -> BytesMut {
891        // Align to page boundaries (4KB) for better memory performance
892        let page_size = 4096;
893        let aligned_size = size.div_ceil(page_size) * page_size;
894
895        BytesMut::with_capacity(aligned_size)
896    }
897
898    /// Create a new lazy buffer pool
899    ///
900    /// # Arguments
901    /// * `buffer_size` - Size of each buffer in bytes (must be non-zero)
902    /// * `max_pool_size` - Maximum number of buffers to pool
903    ///
904    /// # Design Philosophy
905    ///
906    /// **Buffer pool queue capacity is allocated once at application boot**. Backing
907    /// buffers are allocated lazily as clients/backend work first need them, then
908    /// returned to the pool for reuse. This is a critical performance
909    /// optimization:
910    ///
911    /// - **Boot time**: Queue slots only; idle RSS stays small
912    /// - **Warm runtime**: Zero allocations in hot path (acquire/release from pool)
913    /// - **Per-connection**: Buffers are borrowed and returned, never owned
914    ///
915    /// **IMPORTANT**: Do NOT create a `BufferPool` per-client, per-connection, or
916    /// per-request. Create ONE pool at application startup and share it across
917    /// all operations via Arc or static reference.
918    ///
919    /// Buffers are allocated on first use up to `max_pool_size`; allocations beyond
920    /// that remain visible through fallback metrics.
921    #[must_use]
922    pub fn new(buffer_size: BufferSize, max_pool_size: usize) -> Self {
923        let pool = Arc::new(ArrayQueue::new(max_pool_size.max(1)));
924        let pool_size = Arc::new(AtomicUsize::new(0));
925        let allocated_count = Arc::new(AtomicUsize::new(0));
926
927        info!(
928            "Creating lazy buffer pool for up to {} buffers of {}KB each ({}MB max resident after use)",
929            max_pool_size,
930            buffer_size.get() / 1024,
931            (max_pool_size * buffer_size.get()) / (1024 * 1024)
932        );
933
934        Self {
935            pool,
936            buffer_size,
937            max_pool_size,
938            pool_size,
939            allocated_count,
940            // Initialize capture pool as empty (will be configured via with_capture_pool)
941            capture_pool: Arc::new(ArrayQueue::new(1)),
942            capture_capacity: 0,
943            max_capture_pool_size: 0,
944            capture_pool_size: Arc::new(AtomicUsize::new(0)),
945            capture_allocated_count: Arc::new(AtomicUsize::new(0)),
946        }
947    }
948
949    /// Configure the capture buffer pool for accumulator buffers
950    ///
951    /// Capture buffers are used for accumulating streaming data (e.g., caching).
952    /// # Arguments
953    /// * `capacity` - Size of each capture buffer in bytes (e.g., 1 MiB by default)
954    /// * `count` - Maximum number of capture buffers to pool
955    #[must_use]
956    pub fn with_capture_pool(mut self, capacity: usize, count: usize) -> Self {
957        self.capture_pool = Arc::new(ArrayQueue::new(count.max(1)));
958        self.capture_pool_size = Arc::new(AtomicUsize::new(0));
959        self.capture_allocated_count = Arc::new(AtomicUsize::new(0));
960
961        info!(
962            "Creating lazy capture buffer pool for up to {} buffers of {}KB each ({}MB max resident after use)",
963            count,
964            capacity / 1024,
965            (count * capacity) / (1024 * 1024)
966        );
967
968        self.capture_capacity = capacity;
969        self.max_capture_pool_size = count;
970
971        self
972    }
973
974    /// Create a buffer pool suitable for testing
975    ///
976    /// Uses sensible defaults (8KB buffers, pool of 4) that work for most tests.
977    /// Prefer this over manually constructing `BufferPool` in tests.
978    ///
979    /// # Panics
980    /// Panics only if the built-in test buffer size stops satisfying
981    /// `BufferSize` validation.
982    #[cfg(test)]
983    #[must_use]
984    pub fn for_tests() -> Self {
985        Self::new(BufferSize::try_new(8192).expect("valid size"), 4).with_capture_pool(8192, 2)
986    }
987
988    /// Get the current number of available buffers in the pool
989    #[must_use]
990    pub fn available_buffers(&self) -> usize {
991        self.pool_size.load(Ordering::Relaxed)
992    }
993
994    /// Get the number of buffers currently in use
995    #[must_use]
996    pub fn buffers_in_use(&self) -> usize {
997        self.allocated_count
998            .load(Ordering::Relaxed)
999            .saturating_sub(self.available_buffers())
1000    }
1001
1002    /// Get buffer pool statistics (available, in-use, total)
1003    #[must_use]
1004    pub fn stats(&self) -> (usize, usize, usize) {
1005        let available = self.available_buffers();
1006        let in_use = self.buffers_in_use();
1007        (available, in_use, self.max_pool_size)
1008    }
1009
1010    /// Get a buffer from the pool or create a new one (lock-free)
1011    ///
1012    /// Returns a `PooledBuffer` that automatically returns to the pool when dropped.
1013    ///
1014    /// # Performance: Warm Zero-Allocation Hot Path
1015    ///
1016    /// This method allocates lazily until the pool reaches `max_pool_size`.
1017    /// Once buffers have been created and returned, reuse just pops from the
1018    /// lock-free queue. This keeps idle RSS low while preserving the warm
1019    /// zero-allocation forwarding path.
1020    ///
1021    /// Only if the pool is exhausted (all buffers in use) will a new buffer
1022    /// be allocated. Size the pool appropriately to avoid this fallback.
1023    ///
1024    /// # Safety Notes
1025    ///
1026    /// The buffer may contain old bytes outside its logical length, but immutable
1027    /// access only exposes `BytesMut::len()` initialized bytes.
1028    fn wrap_regular_buffer(
1029        &self,
1030        buffer: BytesMut,
1031        fallback: bool,
1032        counts_toward_pool: bool,
1033    ) -> PooledBuffer {
1034        PooledBuffer {
1035            expected_capacity: buffer.capacity(),
1036            buffer,
1037            pool: Arc::clone(&self.pool),
1038            pool_size: Arc::clone(&self.pool_size),
1039            allocated_count: Arc::clone(&self.allocated_count),
1040            max_pool_size: self.max_pool_size,
1041            writable_len: self.buffer_size.get(),
1042            acquired_at: Instant::now(),
1043            source: PooledBufferSource::regular(fallback),
1044            counts_toward_pool,
1045        }
1046    }
1047
1048    fn try_reserve_regular_slot(&self) -> bool {
1049        let mut allocated = self.allocated_count.load(Ordering::Relaxed);
1050        while allocated < self.max_pool_size {
1051            match self.allocated_count.compare_exchange_weak(
1052                allocated,
1053                allocated + 1,
1054                Ordering::Relaxed,
1055                Ordering::Relaxed,
1056            ) {
1057                Ok(_) => return true,
1058                Err(new_allocated) => allocated = new_allocated,
1059            }
1060        }
1061        false
1062    }
1063
1064    fn try_reserve_capture_slot(&self) -> bool {
1065        let mut allocated = self.capture_allocated_count.load(Ordering::Relaxed);
1066        while allocated < self.max_capture_pool_size {
1067            match self.capture_allocated_count.compare_exchange_weak(
1068                allocated,
1069                allocated + 1,
1070                Ordering::Relaxed,
1071                Ordering::Relaxed,
1072            ) {
1073                Ok(_) => return true,
1074                Err(new_allocated) => allocated = new_allocated,
1075            }
1076        }
1077        false
1078    }
1079
1080    fn acquire_now(&self) -> PooledBuffer {
1081        self.pool.pop().map_or_else(
1082            || {
1083                if self.try_reserve_regular_slot() {
1084                    return self.wrap_regular_buffer(
1085                        Self::create_aligned_buffer(self.buffer_size.get()),
1086                        false,
1087                        true,
1088                    );
1089                }
1090
1091                hot_path_allocation_metrics()
1092                    .regular_pool_fallback_allocations
1093                    .fetch_add(1, Ordering::Relaxed);
1094                warn!(
1095                    max_pool_size = self.max_pool_size,
1096                    buffer_size = self.buffer_size.get(),
1097                    "Regular buffer pool exhausted; allocating fallback buffer"
1098                );
1099                self.wrap_regular_buffer(
1100                    Self::create_aligned_buffer(self.buffer_size.get()),
1101                    true,
1102                    false,
1103                )
1104            },
1105            |buffer| {
1106                self.pool_size.fetch_sub(1, Ordering::Relaxed);
1107                // Buffer from pool is logically empty and has the expected allocation
1108                // capacity (enforced on return).
1109                debug_assert_eq!(buffer.len(), 0);
1110                self.wrap_regular_buffer(buffer, false, true)
1111            },
1112        )
1113    }
1114
1115    pub fn acquire(&self) -> PooledBuffer {
1116        self.acquire_now()
1117    }
1118
1119    /// Get a regular buffer only if one is already available in the pool.
1120    ///
1121    /// This is for strict proxy forwarding paths where falling back to a fresh
1122    /// allocation would hide hot-path pressure. Pool exhaustion remains visible
1123    /// through the separate exhaustion metric, but this method returns `None`
1124    /// instead of allocating.
1125    #[cfg(test)]
1126    pub(crate) fn try_acquire(&self) -> Option<PooledBuffer> {
1127        let buffer = self.pool.pop().map_or_else(
1128            || {
1129                hot_path_allocation_metrics()
1130                    .regular_pool_exhaustions
1131                    .fetch_add(1, Ordering::Relaxed);
1132                None
1133            },
1134            Some,
1135        )?;
1136        self.pool_size.fetch_sub(1, Ordering::Relaxed);
1137        debug_assert_eq!(buffer.len(), 0);
1138        Some(self.wrap_regular_buffer(buffer, false, true))
1139    }
1140
1141    /// Get a capture buffer from the capture pool
1142    ///
1143    /// Returns a `PooledBuffer` backed by a pooled capture buffer.
1144    /// Used for accumulating streaming data (e.g., caching articles).
1145    pub(crate) fn acquire_capture_now(&self) -> PooledBuffer {
1146        let mut fallback = false;
1147        let mut counts_toward_pool = true;
1148        let buffer = self.capture_pool.pop().map_or_else(
1149            || {
1150                let capacity = if self.capture_capacity > 0 {
1151                    self.capture_capacity
1152                } else {
1153                    self.buffer_size.get()
1154                };
1155
1156                if self.try_reserve_capture_slot() {
1157                    return BytesMut::with_capacity(capacity);
1158                }
1159
1160                fallback = true;
1161                counts_toward_pool = false;
1162                hot_path_allocation_metrics()
1163                    .capture_pool_fallback_allocations
1164                    .fetch_add(1, Ordering::Relaxed);
1165                warn!(
1166                    max_pool_size = self.max_capture_pool_size,
1167                    capacity, "Capture buffer pool exhausted; allocating fallback buffer"
1168                );
1169                // Pool exhausted - allocate a visible fallback buffer.
1170                BytesMut::with_capacity(capacity)
1171            },
1172            |mut buffer| {
1173                self.capture_pool_size.fetch_sub(1, Ordering::Relaxed);
1174                // Clear but keep capacity for reuse.
1175                buffer.clear();
1176                buffer
1177            },
1178        );
1179
1180        PooledBuffer {
1181            expected_capacity: buffer.capacity(),
1182            buffer,
1183            pool: Arc::clone(&self.capture_pool),
1184            pool_size: Arc::clone(&self.capture_pool_size),
1185            allocated_count: Arc::clone(&self.capture_allocated_count),
1186            max_pool_size: self.max_capture_pool_size,
1187            writable_len: 0,
1188            acquired_at: Instant::now(),
1189            source: PooledBufferSource::capture(fallback),
1190            counts_toward_pool,
1191        }
1192    }
1193
1194    pub fn acquire_capture(&self) -> PooledBuffer {
1195        self.acquire_capture_now()
1196    }
1197}
1198
1199#[cfg(test)]
1200mod tests {
1201    use super::*;
1202    use std::sync::{Mutex, MutexGuard};
1203
1204    fn hot_path_metrics_test_guard() -> MutexGuard<'static, ()> {
1205        static GUARD: Mutex<()> = Mutex::new(());
1206        GUARD.lock().expect("hot-path metrics test guard poisoned")
1207    }
1208    use tokio::io::AsyncWriteExt;
1209
1210    #[tokio::test]
1211    async fn test_buffer_pool_creation() {
1212        let pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 10);
1213
1214        assert_eq!(pool.stats(), (0, 0, 10));
1215
1216        // Pool should lazily allocate buffers on first use.
1217        let buffer1 = pool.acquire();
1218        assert_eq!(buffer1.capacity(), 8192);
1219        assert_eq!(buffer1.initialized(), 0); // No bytes initialized yet
1220        // Buffer automatically returned on drop
1221    }
1222
1223    #[tokio::test]
1224    async fn test_acquire_starts_empty_with_stable_capacity() {
1225        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1226
1227        let capacity = {
1228            let buffer = pool.acquire();
1229            assert_eq!(buffer.initialized(), 0);
1230            assert_eq!(buffer.len(), 0);
1231            buffer.capacity()
1232        };
1233
1234        let buffer = pool.acquire();
1235        assert_eq!(buffer.initialized(), 0);
1236        assert_eq!(buffer.len(), 0);
1237        assert_eq!(buffer.capacity(), capacity);
1238    }
1239
1240    #[tokio::test]
1241    async fn test_read_from_sets_initialized_without_reducing_capacity() {
1242        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1243        let mut buffer = pool.acquire();
1244        let capacity = buffer.capacity();
1245        let (mut writer, mut reader) = tokio::io::duplex(64);
1246
1247        writer.write_all(b"220 ready\r\n").await.unwrap();
1248        drop(writer);
1249
1250        let read = buffer.read_from(&mut reader).await.unwrap();
1251        assert_eq!(read, 11);
1252        assert_eq!(buffer.initialized(), 11);
1253        assert_eq!(&*buffer, b"220 ready\r\n");
1254        assert_eq!(buffer.capacity(), capacity);
1255    }
1256
1257    #[tokio::test]
1258    async fn test_read_from_is_limited_to_fixed_writable_region() {
1259        let pool = BufferPool::new(BufferSize::try_new(8).unwrap(), 1);
1260        let mut buffer = pool.acquire();
1261        assert_eq!(buffer.capacity(), 4096);
1262
1263        let (mut writer, mut reader) = tokio::io::duplex(64);
1264        writer
1265            .write_all(b"220 long response over eight bytes\r\n")
1266            .await
1267            .unwrap();
1268        drop(writer);
1269
1270        let read = buffer.read_from(&mut reader).await.unwrap();
1271        assert_eq!(read, 8);
1272        assert_eq!(buffer.initialized(), 8);
1273        assert_eq!(buffer[..read].len(), read);
1274    }
1275
1276    #[tokio::test]
1277    async fn test_read_more_appends_after_initialized_bytes() {
1278        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1279        let mut buffer = pool.acquire();
1280        buffer.copy_from_slice(b"22");
1281
1282        let (mut writer, mut reader) = tokio::io::duplex(64);
1283        writer.write_all(b"0 ready\r\n").await.unwrap();
1284        drop(writer);
1285
1286        let read = buffer.read_more(&mut reader).await.unwrap();
1287        assert_eq!(read, 9);
1288        assert_eq!(buffer.initialized(), 11);
1289        assert_eq!(&*buffer, b"220 ready\r\n");
1290    }
1291
1292    #[tokio::test]
1293    async fn test_read_more_is_limited_to_remaining_fixed_writable_region() {
1294        let pool = BufferPool::new(BufferSize::try_new(8).unwrap(), 1);
1295        let mut buffer = pool.acquire();
1296        buffer.copy_from_slice(b"22");
1297
1298        let (mut writer, mut reader) = tokio::io::duplex(64);
1299        writer.write_all(b"0 long response\r\n").await.unwrap();
1300        drop(writer);
1301
1302        let read = buffer.read_more(&mut reader).await.unwrap();
1303        assert_eq!(read, 6);
1304        assert_eq!(buffer.initialized(), 8);
1305        assert_eq!(&*buffer, b"220 long");
1306    }
1307
1308    #[tokio::test]
1309    async fn test_copy_from_slice_overwrites_and_allows_up_to_capacity() {
1310        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1311        let mut buffer = pool.acquire();
1312        let capacity = buffer.capacity();
1313
1314        buffer.copy_from_slice(b"previous bytes");
1315        assert_eq!(&*buffer, b"previous bytes");
1316
1317        let data = vec![b'X'; capacity];
1318        buffer.copy_from_slice(&data);
1319        assert_eq!(buffer.initialized(), capacity);
1320        assert_eq!(&buffer[..4], b"XXXX");
1321        assert_eq!(buffer.capacity(), capacity);
1322    }
1323
1324    #[tokio::test]
1325    #[should_panic(expected = "data exceeds buffer capacity")]
1326    async fn test_copy_from_slice_rejects_larger_than_capacity() {
1327        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1328        let mut buffer = pool.acquire();
1329        let too_large = vec![b'X'; buffer.capacity() + 1];
1330
1331        buffer.copy_from_slice(&too_large);
1332    }
1333
1334    #[tokio::test]
1335    async fn test_clear_only_changes_initialized_length() {
1336        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1337        let mut buffer = pool.acquire();
1338        let capacity = buffer.capacity();
1339
1340        buffer.copy_from_slice(b"abcdef");
1341        buffer.clear();
1342        assert_eq!(buffer.initialized(), 0);
1343        assert_eq!(buffer.len(), 0);
1344        assert_eq!(buffer.capacity(), capacity);
1345    }
1346
1347    #[tokio::test]
1348    async fn test_resized_buffers_are_discarded_instead_of_repooled() {
1349        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8, 1);
1350
1351        let normal_capacity = {
1352            let mut buffer = pool.acquire();
1353            let capacity = buffer.capacity();
1354            buffer.extend_from_slice(&vec![b'N'; capacity + 1]);
1355            assert!(buffer.capacity() > capacity);
1356            capacity
1357        };
1358
1359        let buffer = pool.acquire();
1360        assert_eq!(buffer.capacity(), normal_capacity);
1361        drop(buffer);
1362
1363        {
1364            let mut capture = pool.acquire_capture();
1365            capture.extend_from_slice(&[b'C'; 9]);
1366            assert!(capture.capacity() > 8);
1367        }
1368
1369        let capture = pool.acquire_capture();
1370        assert_eq!(capture.capacity(), 8);
1371    }
1372
1373    #[tokio::test]
1374    async fn test_buffer_pool_get_and_return() {
1375        let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 5);
1376
1377        // Get a buffer
1378        let buffer = pool.acquire();
1379        assert_eq!(buffer.capacity(), 4096);
1380        assert_eq!(buffer.initialized(), 0);
1381
1382        // Callers see an empty initialized slice until data is read or copied in.
1383
1384        // Drop it (automatically returns to pool)
1385        drop(buffer);
1386
1387        // Get it again - should be from pool
1388        let buffer2 = pool.acquire();
1389        assert_eq!(buffer2.capacity(), 4096);
1390    }
1391
1392    #[tokio::test]
1393    async fn test_buffer_pool_exhaustion() {
1394        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 2);
1395
1396        // Lazily allocate all pool-owned buffers.
1397        let buf1 = pool.acquire();
1398        let buf2 = pool.acquire();
1399
1400        // Pool is exhausted, should create new buffer
1401        let buf3 = pool.acquire();
1402        assert_eq!(buf3.capacity(), 4096);
1403
1404        // Drop buffers (automatically returned)
1405        drop(buf1);
1406        drop(buf2);
1407        drop(buf3);
1408    }
1409
1410    #[tokio::test]
1411    async fn test_try_acquire_reports_exhaustion_without_allocating() {
1412        let _guard = hot_path_metrics_test_guard();
1413        reset_hot_path_allocation_metrics();
1414        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1415
1416        assert!(pool.try_acquire().is_none());
1417        {
1418            let buffer = pool.acquire();
1419            assert_eq!(buffer.capacity(), 4096);
1420        }
1421
1422        let _held = pool.try_acquire().expect("warmed buffer available");
1423
1424        let before = hot_path_allocation_metrics_snapshot();
1425
1426        assert!(pool.try_acquire().is_none());
1427
1428        let after = hot_path_allocation_metrics_snapshot();
1429        assert_eq!(
1430            after.regular_pool_fallback_allocations,
1431            before.regular_pool_fallback_allocations
1432        );
1433        assert!(after.regular_pool_exhaustions > before.regular_pool_exhaustions);
1434        assert_eq!(pool.stats(), (0, 1, 1));
1435    }
1436
1437    #[tokio::test]
1438    async fn test_buffer_hold_metrics_record_regular_and_capture_drops() {
1439        let _guard = hot_path_metrics_test_guard();
1440        reset_hot_path_allocation_metrics();
1441        let pool =
1442            BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8192, 1);
1443        let before = hot_path_allocation_metrics_snapshot();
1444
1445        drop(pool.acquire());
1446        drop(pool.acquire_capture());
1447
1448        let after = hot_path_allocation_metrics_snapshot();
1449        assert!(after.regular_pool_buffer_holds > before.regular_pool_buffer_holds);
1450        assert!(after.capture_pool_buffer_holds > before.capture_pool_buffer_holds);
1451    }
1452
1453    #[tokio::test]
1454    async fn test_oversized_capture_buffer_is_not_reused() {
1455        let pool =
1456            BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8192, 1);
1457
1458        let mut capture = pool.acquire_capture();
1459        capture.extend_from_slice(&vec![b'X'; 8193]);
1460        assert!(capture.capacity() > 8192);
1461        drop(capture);
1462
1463        let capture2 = pool.acquire_capture();
1464        assert_eq!(capture2.capacity(), 8192);
1465    }
1466
1467    #[tokio::test]
1468    async fn test_capture_pool_lazily_allocates_configured_count() {
1469        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8, 4);
1470
1471        assert_eq!(pool.capture_pool_size.load(Ordering::Relaxed), 0);
1472        assert_eq!(pool.capture_allocated_count.load(Ordering::Relaxed), 0);
1473
1474        let buffers = [
1475            pool.acquire_capture(),
1476            pool.acquire_capture(),
1477            pool.acquire_capture(),
1478            pool.acquire_capture(),
1479        ];
1480
1481        assert_eq!(pool.capture_pool_size.load(Ordering::Relaxed), 0);
1482        assert_eq!(pool.capture_allocated_count.load(Ordering::Relaxed), 4);
1483        assert!(buffers.iter().all(|buffer| buffer.capacity() == 8));
1484        drop(buffers);
1485
1486        assert_eq!(pool.capture_pool_size.load(Ordering::Relaxed), 4);
1487    }
1488
1489    #[tokio::test]
1490    async fn test_chunked_response_helpers_without_flattening() {
1491        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8, 4);
1492        let mut response = ChunkedResponse::default();
1493        response.extend_from_slice(&pool, b"220 0 <id>\r\nBody\r\n.\r\n");
1494
1495        let chunks: Vec<&[u8]> = response.iter_chunks().collect();
1496        assert!(chunks.len() > 1, "test requires multi-chunk buffering");
1497        assert_eq!(chunks.concat(), b"220 0 <id>\r\nBody\r\n.\r\n");
1498
1499        let mut prefix = SmallVec::<[u8; 128]>::new();
1500        response.copy_prefix_into(12, &mut prefix);
1501        assert_eq!(&prefix[..], b"220 0 <id>\r\n");
1502    }
1503
1504    #[tokio::test]
1505    async fn test_chunked_response_uses_more_capture_buffers_instead_of_growing_one() {
1506        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8, 4);
1507        let mut response = ChunkedResponse::default();
1508
1509        response.extend_from_slice(&pool, b"123456789abcdefghi");
1510
1511        let chunks: Vec<&[u8]> = response.iter_chunks().collect();
1512        assert_eq!(
1513            chunks.iter().map(|chunk| chunk.len()).collect::<Vec<_>>(),
1514            vec![8, 8, 2]
1515        );
1516    }
1517
1518    #[tokio::test]
1519    async fn test_chunked_response_copy_prefix_clamps_to_length() {
1520        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1).with_capture_pool(8, 4);
1521        let mut response = ChunkedResponse::default();
1522        response.extend_from_slice(&pool, b"430\r\n");
1523
1524        let mut prefix = SmallVec::<[u8; 128]>::new();
1525        response.copy_prefix_into(64, &mut prefix);
1526        assert_eq!(&prefix[..], b"430\r\n");
1527    }
1528
1529    #[tokio::test]
1530    async fn test_chunked_response_can_own_range_without_copying() {
1531        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 2);
1532        let mut buffer = pool.acquire();
1533        buffer.copy_from_slice(b"xx220 0 <id>\r\nBody\r\n.\r\nyy");
1534        let expected_ptr = buffer.as_ref()[2..].as_ptr();
1535
1536        let mut response = ChunkedResponse::default();
1537        response.push_buffer_range(buffer, 2..23);
1538
1539        let first = response.first_chunk().expect("range chunk");
1540        assert_eq!(first.as_ptr(), expected_ptr);
1541        assert_eq!(first, b"220 0 <id>\r\nBody\r\n.\r\n");
1542    }
1543
1544    #[tokio::test]
1545    async fn test_chunked_response_can_share_pooled_range_without_copying() {
1546        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 2);
1547        let mut buffer = pool.acquire();
1548        buffer.copy_from_slice(b"xx220 0 <id>\r\nBody\r\n.\r\nyy");
1549        let shared = std::sync::Arc::new(buffer);
1550        let expected_ptr = shared.as_ref()[2..].as_ptr();
1551
1552        let mut response = ChunkedResponse::default();
1553        response.push_shared_pooled_range(std::sync::Arc::clone(&shared), 2..23);
1554
1555        let first = response.first_chunk().expect("shared pooled range chunk");
1556        assert_eq!(first.as_ptr(), expected_ptr);
1557        assert_eq!(first, b"220 0 <id>\r\nBody\r\n.\r\n");
1558        drop(response);
1559        drop(shared);
1560        assert_eq!(pool.available_buffers(), 1);
1561    }
1562
1563    #[tokio::test]
1564    async fn test_freeze_detaches_bytes_from_pool_ownership() {
1565        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1);
1566        let mut buffer = pool.acquire();
1567        buffer.copy_from_slice(b"223 0 <id>\r\n");
1568        assert_eq!(pool.available_buffers(), 0);
1569
1570        let bytes = buffer.freeze();
1571        assert_eq!(&bytes[..], b"223 0 <id>\r\n");
1572
1573        drop(bytes);
1574
1575        assert_eq!(
1576            pool.available_buffers(),
1577            0,
1578            "dropping frozen bytes must not return the detached allocation to the pool"
1579        );
1580    }
1581
1582    #[tokio::test]
1583    async fn test_buffer_pool_concurrent_access() {
1584        let pool = BufferPool::new(BufferSize::try_new(2048).unwrap(), 10);
1585
1586        // Spawn multiple tasks accessing the pool concurrently
1587        let mut handles = vec![];
1588
1589        for _ in 0..20 {
1590            let pool_clone = pool.clone();
1591            let handle = tokio::spawn(async move {
1592                let buffer = pool_clone.acquire();
1593                assert_eq!(buffer.capacity(), 4096);
1594                // Simulate some work
1595                tokio::time::sleep(tokio::time::Duration::from_millis(1)).await;
1596            });
1597            handles.push(handle);
1598        }
1599
1600        // Wait for all tasks to complete
1601        for handle in handles {
1602            handle.await.unwrap();
1603        }
1604    }
1605
1606    #[tokio::test]
1607    async fn test_buffer_alignment() {
1608        let pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 1);
1609        let buffer = pool.acquire();
1610
1611        // Buffer capacity should be aligned to page boundaries (4KB)
1612        assert!(buffer.capacity() >= 8192);
1613        // Should be page-aligned (multiple of 4096)
1614        assert_eq!(buffer.capacity() % 4096, 0);
1615    }
1616
1617    #[tokio::test]
1618    async fn test_buffer_clear_and_resize() {
1619        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 2);
1620
1621        let mut buffer = pool.acquire();
1622
1623        // Write data using copy_from_slice
1624        let data = vec![42u8; 101];
1625        buffer.copy_from_slice(&data);
1626        assert_eq!(buffer.initialized(), 101);
1627
1628        // Drop returns it to pool
1629        drop(buffer);
1630
1631        // Get it again - may contain old data (performance optimization)
1632        let buffer2 = pool.acquire();
1633        assert_eq!(buffer2.capacity(), 4096);
1634        // Note: buffer may contain previous bytes outside the initialized range.
1635    }
1636
1637    #[tokio::test]
1638    async fn test_buffer_pool_max_size_enforcement() {
1639        let pool = BufferPool::new(BufferSize::try_new(512).unwrap(), 3);
1640
1641        // Get all buffers
1642        let buf1 = pool.acquire();
1643        let buf2 = pool.acquire();
1644        let buf3 = pool.acquire();
1645
1646        // Get one more (should create new)
1647        let buf4 = pool.acquire();
1648
1649        // Drop all buffers (automatically returned)
1650        drop(buf1);
1651        drop(buf2);
1652        drop(buf3);
1653        drop(buf4);
1654
1655        // Pool should not exceed max size
1656        // (We can't directly test pool size, but the implementation handles it)
1657    }
1658
1659    #[tokio::test]
1660    async fn test_buffer_wrong_size_not_returned() {
1661        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 2);
1662
1663        let buffer = pool.acquire();
1664        assert_eq!(buffer.capacity(), 4096);
1665
1666        // PooledBuffer auto-returns on drop with correct size enforcement in Drop impl
1667        drop(buffer);
1668    }
1669
1670    #[tokio::test]
1671    async fn test_buffer_pool_multiple_get_return_cycles() {
1672        let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 5);
1673
1674        // Do multiple get/return cycles
1675        for i in 0u8..20 {
1676            let mut buffer = pool.acquire();
1677            assert_eq!(buffer.capacity(), 4096);
1678
1679            // Write some data using copy_from_slice
1680            let data = vec![i; 1];
1681            buffer.copy_from_slice(&data);
1682            assert_eq!(buffer.initialized(), 1);
1683        }
1684    }
1685
1686    #[test]
1687    fn test_buffer_pool_clone() {
1688        let pool1 = BufferPool::new(BufferSize::try_new(1024).unwrap(), 5);
1689        let _pool2 = pool1;
1690
1691        // Both should share the same underlying pool
1692        // (Arc ensures shared ownership)
1693    }
1694
1695    #[tokio::test]
1696    async fn test_different_buffer_sizes() {
1697        let small_pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 5);
1698        let medium_pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 5);
1699        let large_pool = BufferPool::new(BufferSize::try_new(65536).unwrap(), 5);
1700
1701        let small_buf = small_pool.acquire();
1702        let medium_buf = medium_pool.acquire();
1703        let large_buf = large_pool.acquire();
1704
1705        assert_eq!(small_buf.capacity(), 4096);
1706        assert_eq!(medium_buf.capacity(), 8192);
1707        assert_eq!(large_buf.capacity(), 65536);
1708
1709        // Buffers auto-return on drop
1710    }
1711
1712    #[tokio::test]
1713    async fn test_buffer_pool_stress() {
1714        let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 10);
1715
1716        // Stress test with many concurrent operations
1717        let mut handles = vec![];
1718
1719        for _ in 0..100 {
1720            let pool_clone = pool.clone();
1721            let handle = tokio::spawn(async move {
1722                for _ in 0..10 {
1723                    let buffer = pool_clone.acquire();
1724                    assert_eq!(buffer.capacity(), 4096);
1725                }
1726            });
1727            handles.push(handle);
1728        }
1729
1730        for handle in handles {
1731            handle.await.unwrap();
1732        }
1733    }
1734
1735    #[tokio::test]
1736    async fn test_pooled_buffer_deref() {
1737        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 5);
1738        let mut buffer = pool.acquire();
1739
1740        // Initially no initialized bytes
1741        assert_eq!(buffer.len(), 0);
1742
1743        // Write data
1744        buffer.copy_from_slice(b"Hello");
1745
1746        // Deref should return only initialized portion
1747        assert_eq!(buffer.len(), 5);
1748        assert_eq!(&*buffer, b"Hello");
1749    }
1750
1751    #[tokio::test]
1752    async fn test_pooled_buffer_as_ref() {
1753        let pool = BufferPool::new(BufferSize::try_new(512).unwrap(), 5);
1754        let mut buffer = pool.acquire();
1755
1756        buffer.copy_from_slice(b"Test data");
1757
1758        // AsRef should return initialized portion
1759        let slice: &[u8] = buffer.as_ref();
1760        assert_eq!(slice, b"Test data");
1761        assert_eq!(slice.len(), 9);
1762    }
1763
1764    #[tokio::test]
1765    async fn test_copy_from_slice_updates_initialized() {
1766        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 5);
1767        let mut buffer = pool.acquire();
1768
1769        assert_eq!(buffer.initialized(), 0);
1770
1771        buffer.copy_from_slice(b"abc");
1772        assert_eq!(buffer.initialized(), 3);
1773
1774        buffer.copy_from_slice(b"longer text");
1775        assert_eq!(buffer.initialized(), 11);
1776    }
1777
1778    #[tokio::test]
1779    #[should_panic(expected = "data exceeds buffer capacity")]
1780    async fn test_copy_from_slice_panic_on_overflow() {
1781        let pool = BufferPool::new(BufferSize::try_new(10).unwrap(), 5);
1782        let mut buffer = pool.acquire();
1783
1784        let too_large = vec![0u8; buffer.capacity() + 1];
1785        buffer.copy_from_slice(&too_large); // Should panic
1786    }
1787
1788    #[tokio::test]
1789    async fn test_buffer_pool_debug() {
1790        let pool = BufferPool::new(BufferSize::try_new(2048).unwrap(), 5);
1791        let debug_str = format!("{pool:?}");
1792        assert!(debug_str.contains("BufferPool"));
1793    }
1794
1795    #[tokio::test]
1796    async fn test_buffer_initialized_tracking() {
1797        let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 5);
1798        let mut buffer = pool.acquire();
1799
1800        // Test multiple writes update initialized correctly
1801        buffer.copy_from_slice(b"First");
1802        assert_eq!(buffer.initialized(), 5);
1803        assert_eq!(&*buffer, b"First");
1804
1805        buffer.copy_from_slice(b"Second write");
1806        assert_eq!(buffer.initialized(), 12);
1807        assert_eq!(&*buffer, b"Second write");
1808    }
1809
1810    #[tokio::test]
1811    async fn test_buffer_capacity_vs_initialized() {
1812        let pool = BufferPool::new(BufferSize::try_new(8192).unwrap(), 5);
1813        let mut buffer = pool.acquire();
1814
1815        // Capacity is full buffer size
1816        assert_eq!(buffer.capacity(), 8192);
1817
1818        // Initialized is what we've written
1819        assert_eq!(buffer.initialized(), 0);
1820
1821        buffer.copy_from_slice(b"Small");
1822        assert_eq!(buffer.capacity(), 8192);
1823        assert_eq!(buffer.initialized(), 5);
1824    }
1825
1826    #[tokio::test]
1827    async fn test_empty_slice_copy() {
1828        let pool = BufferPool::new(BufferSize::try_new(512).unwrap(), 5);
1829        let mut buffer = pool.acquire();
1830
1831        // Copying empty slice should work
1832        buffer.copy_from_slice(&[]);
1833        assert_eq!(buffer.initialized(), 0);
1834        assert_eq!(&*buffer, b"");
1835    }
1836
1837    #[tokio::test]
1838    async fn test_buffer_reuse_preserves_capacity() {
1839        let pool = BufferPool::new(BufferSize::try_new(2048).unwrap(), 5);
1840
1841        {
1842            let mut buffer = pool.acquire();
1843            buffer.copy_from_slice(b"test");
1844            assert_eq!(buffer.capacity(), 4096);
1845        } // Drop returns to pool
1846
1847        let buffer2 = pool.acquire();
1848        // Should have same capacity when reused
1849        assert_eq!(buffer2.capacity(), 4096);
1850    }
1851
1852    #[test]
1853    fn test_buffer_size_alignment() {
1854        // Test that create_aligned_buffer aligns to page boundaries
1855        let buffer = BufferPool::create_aligned_buffer(1000);
1856        // Should be aligned to 4096
1857        assert_eq!(buffer.len(), 0);
1858        assert_eq!(buffer.capacity() % 4096, 0);
1859
1860        let buffer2 = BufferPool::create_aligned_buffer(8192);
1861        assert_eq!(buffer2.len(), 0);
1862        assert_eq!(buffer2.capacity() % 4096, 0);
1863    }
1864}