From ce3c2b729d5d191bb3794e7be6fc1a9d0cedc40f Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Mon, 14 Sep 2026 17:49:20 -0600 Subject: [PATCH 01/23] perf: defer pooled buffer compaction --- src/pool/buffer.rs | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index 62ff03d5..3ab83676 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -145,6 +145,12 @@ impl BufferStorage { self.bytes = prefix; } + fn compact_visible_if_full(&mut self) { + if self.hidden_prefix.is_some() && self.bytes.capacity() == self.bytes.len() { + self.compact_visible(); + } + } + fn restore_hidden_prefix(&mut self) { let Some(mut prefix) = self.hidden_prefix.take() else { return; @@ -173,7 +179,6 @@ impl BufferStorage { where R: AsyncRead + Unpin, { - debug_assert!(self.hidden_prefix.is_none()); if read_len == 0 { return Ok(0); } @@ -281,7 +286,7 @@ impl PooledBuffer { self.read_limit() } ReadMode::Append => { - self.buffer.compact_visible(); + self.buffer.compact_visible_if_full(); self.read_limit() .saturating_sub(self.buffer.initialized_len()) } @@ -1487,7 +1492,7 @@ mod tests { } #[tokio::test] - async fn read_more_compacts_a_retained_prefix_before_appending() { + async fn read_more_appends_after_a_retained_prefix_without_compacting() { let pool = BufferPool::new(BufferSize::try_new(1024).unwrap(), 1); let mut buffer = pool.acquire(); buffer.copy_from_slice(b"discard22unused"); @@ -1504,6 +1509,24 @@ mod tests { assert_eq!(buffer.allocation_ptr(), allocation); } + #[tokio::test] + async fn read_more_compacts_a_retained_prefix_when_tail_is_full() { + let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(&vec![b'x'; 4096]); + let allocation = buffer.allocation_ptr(); + buffer.expose_initialized_range_without_copying(4094..4096); + let (mut writer, mut reader) = tokio::io::duplex(64); + writer.write_all(b"y").await.unwrap(); + drop(writer); + + let read = buffer.read_more(&mut reader).await.unwrap(); + + assert_eq!(read, 1); + assert_eq!(buffer.as_ref(), b"xxy"); + assert_eq!(buffer.allocation_ptr(), allocation); + } + #[tokio::test] async fn fresh_read_after_a_range_view_restores_the_full_writable_region() { let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); From 274ae916414e17420b5319c30f8809cfe85c1d9f Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Mon, 14 Sep 2026 17:54:05 -0600 Subject: [PATCH 02/23] refactor: model pooled buffer storage states --- src/pool/buffer.rs | 54 ++++++++++++++++++++++++++++++---------------- 1 file changed, 35 insertions(+), 19 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index 3ab83676..2fdf40ae 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -66,16 +66,25 @@ enum ReadMode { #[derive(Default)] struct BufferStorage { bytes: BytesMut, - hidden_prefix: Option, + state: BufferState, allocation_capacity: usize, } +#[derive(Default)] +enum BufferState { + #[default] + Contiguous, + Retained { + prefix: BytesMut, + }, +} + impl BufferStorage { fn new(bytes: BytesMut) -> Self { let allocation_capacity = bytes.capacity(); Self { bytes, - hidden_prefix: None, + state: BufferState::Contiguous, allocation_capacity, } } @@ -92,13 +101,13 @@ impl BufferStorage { #[inline] fn clear(&mut self) { - self.restore_hidden_prefix(); + self.restore_contiguous(); self.bytes.clear(); } #[inline] fn clear_for_fresh_read(&mut self) { - self.restore_hidden_prefix(); + self.restore_contiguous(); self.bytes.clear(); } @@ -126,15 +135,20 @@ impl BufferStorage { let mut visible = self.bytes.split_off(range.start); visible.truncate(range.len()); let newly_hidden_prefix = std::mem::replace(&mut self.bytes, visible); - if let Some(prefix) = &mut self.hidden_prefix { - prefix.unsplit(newly_hidden_prefix); - } else { - self.hidden_prefix = Some(newly_hidden_prefix); + match &mut self.state { + BufferState::Contiguous => { + self.state = BufferState::Retained { + prefix: newly_hidden_prefix, + }; + } + BufferState::Retained { prefix } => { + prefix.unsplit(newly_hidden_prefix); + } } } fn compact_visible(&mut self) { - let Some(mut prefix) = self.hidden_prefix.take() else { + let BufferState::Retained { mut prefix } = std::mem::take(&mut self.state) else { return; }; let prefix_len = prefix.len(); @@ -146,13 +160,15 @@ impl BufferStorage { } fn compact_visible_if_full(&mut self) { - if self.hidden_prefix.is_some() && self.bytes.capacity() == self.bytes.len() { + if matches!(self.state, BufferState::Retained { .. }) + && self.bytes.capacity() == self.bytes.len() + { self.compact_visible(); } } - fn restore_hidden_prefix(&mut self) { - let Some(mut prefix) = self.hidden_prefix.take() else { + fn restore_contiguous(&mut self) { + let BufferState::Retained { mut prefix } = std::mem::take(&mut self.state) else { return; }; prefix.unsplit(std::mem::take(&mut self.bytes)); @@ -164,7 +180,7 @@ impl BufferStorage { } fn take(&mut self) -> BytesMut { - drop(self.hidden_prefix.take()); + self.state = BufferState::Contiguous; self.allocation_capacity = 0; std::mem::take(&mut self.bytes) } @@ -172,7 +188,7 @@ impl BufferStorage { fn restore(&mut self, bytes: BytesMut) { self.allocation_capacity = bytes.capacity(); self.bytes = bytes; - self.hidden_prefix = None; + self.state = BufferState::Contiguous; } async fn read_from(&mut self, reader: &mut R, read_len: usize) -> std::io::Result @@ -301,7 +317,7 @@ impl PooledBuffer { #[must_use] pub(crate) fn is_exposed_range_view(&self) -> bool { - self.buffer.hidden_prefix.is_some() + matches!(self.buffer.state, BufferState::Retained { .. }) } #[must_use] @@ -311,10 +327,10 @@ impl PooledBuffer { #[cfg(test)] pub(crate) fn allocation_ptr(&self) -> *const u8 { - self.buffer - .hidden_prefix - .as_ref() - .map_or_else(|| self.buffer.bytes.as_ptr(), |prefix| prefix.as_ptr()) + match &self.buffer.state { + BufferState::Contiguous => self.buffer.bytes.as_ptr(), + BufferState::Retained { prefix } => prefix.as_ptr(), + } } /// Copy data into buffer and mark as initialized From df24ffc37db17b51a8be7af2a3281cd4e4d23fb6 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Mon, 14 Sep 2026 18:20:26 -0600 Subject: [PATCH 03/23] perf: retain pooled buffer views by offset --- src/pool/buffer.rs | 149 +++++++++++++------------------ src/session/multiline_framing.rs | 3 +- 2 files changed, 65 insertions(+), 87 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index 2fdf40ae..966474ea 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -57,63 +57,50 @@ struct BufferAcquisition { counts_toward_pool: bool, } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum ReadMode { - Reset, - Append, -} - #[derive(Default)] struct BufferStorage { bytes: BytesMut, - state: BufferState, - allocation_capacity: usize, + visible_start: VisibleStart, } -#[derive(Default)] -enum BufferState { - #[default] - Contiguous, - Retained { - prefix: BytesMut, - }, +#[derive(Clone, Copy, Default)] +struct VisibleStart(usize); + +impl VisibleStart { + fn get(self) -> usize { + self.0 + } } impl BufferStorage { fn new(bytes: BytesMut) -> Self { - let allocation_capacity = bytes.capacity(); Self { bytes, - state: BufferState::Contiguous, - allocation_capacity, + visible_start: VisibleStart::default(), } } #[inline] fn capacity(&self) -> usize { - self.allocation_capacity + self.bytes.capacity() } #[inline] fn initialized_len(&self) -> usize { - self.bytes.len() + self.as_slice().len() } #[inline] fn clear(&mut self) { - self.restore_contiguous(); - self.bytes.clear(); - } - - #[inline] - fn clear_for_fresh_read(&mut self) { - self.restore_contiguous(); self.bytes.clear(); + self.visible_start = VisibleStart::default(); } #[inline] fn as_slice(&self) -> &[u8] { - &self.bytes + self.bytes + .get(self.visible_start.get()..) + .expect("visible start must remain within initialized bytes") } fn copy_from_slice(&mut self, data: &[u8]) { @@ -124,7 +111,6 @@ impl BufferStorage { fn extend_from_slice(&mut self, data: &[u8]) { self.compact_visible(); self.bytes.extend_from_slice(data); - self.allocation_capacity = self.bytes.capacity(); } fn retain_range(&mut self, range: Range) { @@ -132,47 +118,39 @@ impl BufferStorage { range.start < range.end && range.end <= self.initialized_len(), "exposed range must be non-empty and inside initialized bytes" ); - let mut visible = self.bytes.split_off(range.start); - visible.truncate(range.len()); - let newly_hidden_prefix = std::mem::replace(&mut self.bytes, visible); - match &mut self.state { - BufferState::Contiguous => { - self.state = BufferState::Retained { - prefix: newly_hidden_prefix, - }; - } - BufferState::Retained { prefix } => { - prefix.unsplit(newly_hidden_prefix); - } - } + let start = self + .visible_start + .get() + .checked_add(range.start) + .expect("visible range start overflowed"); + let end = self + .visible_start + .get() + .checked_add(range.end) + .expect("visible range end overflowed"); + self.bytes.truncate(end); + self.visible_start = VisibleStart(start); } fn compact_visible(&mut self) { - let BufferState::Retained { mut prefix } = std::mem::take(&mut self.state) else { + let visible_start = self.visible_start.get(); + if visible_start == 0 { return; - }; - let prefix_len = prefix.len(); - let visible_len = self.bytes.len(); - prefix.unsplit(std::mem::take(&mut self.bytes)); - prefix.copy_within(prefix_len..prefix_len + visible_len, 0); - prefix.truncate(visible_len); - self.bytes = prefix; + } + + self.bytes.copy_within(visible_start.., 0); + self.bytes.truncate(self.bytes.len() - visible_start); + self.visible_start = VisibleStart::default(); } - fn compact_visible_if_full(&mut self) { - if matches!(self.state, BufferState::Retained { .. }) - && self.bytes.capacity() == self.bytes.len() - { + fn compact_visible_if_tail_full(&mut self) { + if self.visible_start.get() != 0 && self.bytes.capacity() == self.bytes.len() { self.compact_visible(); } } - fn restore_contiguous(&mut self) { - let BufferState::Retained { mut prefix } = std::mem::take(&mut self.state) else { - return; - }; - prefix.unsplit(std::mem::take(&mut self.bytes)); - self.bytes = prefix; + fn has_retained_prefix(&self) -> bool { + self.visible_start.get() != 0 } fn freeze(&mut self) -> Bytes { @@ -180,15 +158,19 @@ impl BufferStorage { } fn take(&mut self) -> BytesMut { - self.state = BufferState::Contiguous; - self.allocation_capacity = 0; - std::mem::take(&mut self.bytes) + let visible_start = self.visible_start.get(); + self.visible_start = VisibleStart::default(); + let mut bytes = std::mem::take(&mut self.bytes); + if visible_start == 0 { + bytes + } else { + bytes.split_off(visible_start) + } } fn restore(&mut self, bytes: BytesMut) { - self.allocation_capacity = bytes.capacity(); self.bytes = bytes; - self.state = BufferState::Contiguous; + self.visible_start = VisibleStart::default(); } async fn read_from(&mut self, reader: &mut R, read_len: usize) -> std::io::Result @@ -273,7 +255,8 @@ impl PooledBuffer { where R: AsyncRead + Unpin, { - self.read_into_spare(reader, ReadMode::Reset).await + self.buffer.clear(); + self.read_into_spare(reader, self.read_limit()).await } /// Read more data at the current initialized offset, accumulating bytes. @@ -289,24 +272,21 @@ impl PooledBuffer { where R: AsyncRead + Unpin, { - self.read_into_spare(reader, ReadMode::Append).await + self.buffer.compact_visible_if_tail_full(); + let read_len = self + .read_limit() + .saturating_sub(self.buffer.initialized_len()); + self.read_into_spare(reader, read_len).await } - async fn read_into_spare(&mut self, reader: &mut R, mode: ReadMode) -> std::io::Result + async fn read_into_spare( + &mut self, + reader: &mut R, + read_len: usize, + ) -> std::io::Result where R: AsyncRead + Unpin, { - let read_len = match mode { - ReadMode::Reset => { - self.buffer.clear_for_fresh_read(); - self.read_limit() - } - ReadMode::Append => { - self.buffer.compact_visible_if_full(); - self.read_limit() - .saturating_sub(self.buffer.initialized_len()) - } - }; self.buffer.read_from(reader, read_len).await } @@ -316,8 +296,8 @@ impl PooledBuffer { } #[must_use] - pub(crate) fn is_exposed_range_view(&self) -> bool { - matches!(self.buffer.state, BufferState::Retained { .. }) + pub(crate) fn has_retained_prefix(&self) -> bool { + self.buffer.has_retained_prefix() } #[must_use] @@ -327,10 +307,7 @@ impl PooledBuffer { #[cfg(test)] pub(crate) fn allocation_ptr(&self) -> *const u8 { - match &self.buffer.state { - BufferState::Contiguous => self.buffer.bytes.as_ptr(), - BufferState::Retained { prefix } => prefix.as_ptr(), - } + self.buffer.bytes.as_ptr() } /// Copy data into buffer and mark as initialized @@ -1500,11 +1477,13 @@ mod tests { let mut buffer = pool.acquire(); buffer.copy_from_slice(b"discardkeepdiscard"); let allocation = buffer.allocation_ptr(); + let capacity = buffer.capacity(); buffer.expose_initialized_range_without_copying(7..11); assert_eq!(buffer.as_ref(), b"keep"); assert_eq!(buffer.allocation_ptr(), allocation); + assert_eq!(buffer.capacity(), capacity); } #[tokio::test] diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 8db4f953..30261970 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -1351,8 +1351,7 @@ where .await? } FramedMultilineChunk::Incomplete(incomplete) => { - if io_buffer.is_exposed_range_view() && io_buffer.has_remaining_fixed_writable_region() - { + if io_buffer.has_retained_prefix() && io_buffer.has_remaining_fixed_writable_region() { incomplete .compact_packed_prefix_and_write_with_next_backend_chunk( writer, From 2bd12e5184c372d9b038a705eb0856833d1caa5d Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Mon, 14 Sep 2026 18:28:29 -0600 Subject: [PATCH 04/23] refactor: simplify retained buffer offset --- src/pool/buffer.rs | 101 ++++++++++++++++--------------- src/pool/mod.rs | 1 + src/session/multiline_framing.rs | 9 +-- 3 files changed, 58 insertions(+), 53 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index 966474ea..c50dc41f 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -39,6 +39,11 @@ pub struct PooledBuffer { counts_toward_pool: bool, } +/// Proof that a retained suffix has logical room for another backend read. +pub(crate) struct AppendableRetainedBuffer<'a> { + buffer: &'a mut PooledBuffer, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum PooledBufferKind { Regular, @@ -60,23 +65,14 @@ struct BufferAcquisition { #[derive(Default)] struct BufferStorage { bytes: BytesMut, - visible_start: VisibleStart, -} - -#[derive(Clone, Copy, Default)] -struct VisibleStart(usize); - -impl VisibleStart { - fn get(self) -> usize { - self.0 - } + visible_from: usize, } impl BufferStorage { fn new(bytes: BytesMut) -> Self { Self { bytes, - visible_start: VisibleStart::default(), + visible_from: 0, } } @@ -87,20 +83,18 @@ impl BufferStorage { #[inline] fn initialized_len(&self) -> usize { - self.as_slice().len() + self.bytes.len() - self.visible_from } #[inline] fn clear(&mut self) { self.bytes.clear(); - self.visible_start = VisibleStart::default(); + self.visible_from = 0; } #[inline] fn as_slice(&self) -> &[u8] { - self.bytes - .get(self.visible_start.get()..) - .expect("visible start must remain within initialized bytes") + &self.bytes[self.visible_from..] } fn copy_from_slice(&mut self, data: &[u8]) { @@ -118,39 +112,31 @@ impl BufferStorage { range.start < range.end && range.end <= self.initialized_len(), "exposed range must be non-empty and inside initialized bytes" ); - let start = self - .visible_start - .get() - .checked_add(range.start) - .expect("visible range start overflowed"); - let end = self - .visible_start - .get() - .checked_add(range.end) - .expect("visible range end overflowed"); + let start = self.visible_from + range.start; + let end = self.visible_from + range.end; self.bytes.truncate(end); - self.visible_start = VisibleStart(start); + self.visible_from = start; } fn compact_visible(&mut self) { - let visible_start = self.visible_start.get(); - if visible_start == 0 { + if self.visible_from == 0 { return; } - self.bytes.copy_within(visible_start.., 0); - self.bytes.truncate(self.bytes.len() - visible_start); - self.visible_start = VisibleStart::default(); + let len = self.initialized_len(); + self.bytes.copy_within(self.visible_from.., 0); + self.bytes.truncate(len); + self.visible_from = 0; } fn compact_visible_if_tail_full(&mut self) { - if self.visible_start.get() != 0 && self.bytes.capacity() == self.bytes.len() { + if self.visible_from != 0 && self.bytes.capacity() == self.bytes.len() { self.compact_visible(); } } fn has_retained_prefix(&self) -> bool { - self.visible_start.get() != 0 + self.visible_from != 0 } fn freeze(&mut self) -> Bytes { @@ -158,19 +144,18 @@ impl BufferStorage { } fn take(&mut self) -> BytesMut { - let visible_start = self.visible_start.get(); - self.visible_start = VisibleStart::default(); + let visible_from = std::mem::take(&mut self.visible_from); let mut bytes = std::mem::take(&mut self.bytes); - if visible_start == 0 { + if visible_from == 0 { bytes } else { - bytes.split_off(visible_start) + bytes.split_off(visible_from) } } fn restore(&mut self, bytes: BytesMut) { self.bytes = bytes; - self.visible_start = VisibleStart::default(); + self.visible_from = 0; } async fn read_from(&mut self, reader: &mut R, read_len: usize) -> std::io::Result @@ -295,14 +280,12 @@ impl PooledBuffer { self.buffer.retain_range(range); } - #[must_use] - pub(crate) fn has_retained_prefix(&self) -> bool { - self.buffer.has_retained_prefix() - } - - #[must_use] - pub(crate) fn has_remaining_fixed_writable_region(&self) -> bool { - self.buffer.initialized_len() < self.read_limit() + pub(crate) fn appendable_retained(&mut self) -> Option> { + if !self.buffer.has_retained_prefix() || self.buffer.initialized_len() >= self.read_limit() + { + return None; + } + Some(AppendableRetainedBuffer { buffer: self }) } #[cfg(test)] @@ -376,6 +359,19 @@ impl PooledBuffer { } } +impl<'a> AppendableRetainedBuffer<'a> { + pub(crate) async fn read_more(&mut self, reader: &mut R) -> std::io::Result + where + R: AsyncRead + Unpin, + { + self.buffer.read_more(reader).await + } + + pub(crate) fn into_inner(self) -> &'a mut PooledBuffer { + self.buffer + } +} + impl Deref for PooledBuffer { type Target = [u8]; @@ -1497,11 +1493,15 @@ mod tests { writer.write_all(b"0 ready\r\n").await.unwrap(); drop(writer); - let read = buffer.read_more(&mut reader).await.unwrap(); + let Some(mut appendable) = buffer.appendable_retained() else { + panic!("retained prefix should have spare tail capacity"); + }; + let read = appendable.read_more(&mut reader).await.unwrap(); assert_eq!(read, 9); assert_eq!(buffer.as_ref(), b"220 ready\r\n"); assert_eq!(buffer.allocation_ptr(), allocation); + assert!(buffer.buffer.has_retained_prefix()); } #[tokio::test] @@ -1515,7 +1515,10 @@ mod tests { writer.write_all(b"y").await.unwrap(); drop(writer); - let read = buffer.read_more(&mut reader).await.unwrap(); + let Some(mut appendable) = buffer.appendable_retained() else { + panic!("retained prefix should have logical room to append"); + }; + let read = appendable.read_more(&mut reader).await.unwrap(); assert_eq!(read, 1); assert_eq!(buffer.as_ref(), b"xxy"); diff --git a/src/pool/mod.rs b/src/pool/mod.rs index 733e05bb..2d6ce1c9 100644 --- a/src/pool/mod.rs +++ b/src/pool/mod.rs @@ -10,6 +10,7 @@ pub mod health_check; pub mod prewarming; pub mod provider; +pub(crate) use buffer::AppendableRetainedBuffer; pub use buffer::{ BufferPool, ChunkedResponse, HotPathAllocationMetricsSnapshot, PooledBuffer, hot_path_allocation_metrics_snapshot, reset_hot_path_allocation_metrics, diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 30261970..4b04366c 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -273,12 +273,12 @@ struct IncompleteMultilineWireChunk { impl IncompleteMultilineWireChunk { #[allow(clippy::too_many_arguments)] - async fn compact_packed_prefix_and_write_with_next_backend_chunk( + async fn append_to_packed_prefix_and_write_with_next_backend_chunk( self, writer: &mut W, current_len: usize, framer: &mut MultilineFramer, - io_buffer: &mut crate::pool::PooledBuffer, + mut io_buffer: crate::pool::AppendableRetainedBuffer<'_>, conn: &mut crate::stream::ConnectionStream, pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, @@ -300,6 +300,7 @@ impl IncompleteMultilineWireChunk { }, ); } + let io_buffer = io_buffer.into_inner(); let total_len = current_len + n; match framer.frame_next_multiline_chunk(self, &io_buffer[current_len..total_len]) { @@ -1351,9 +1352,9 @@ where .await? } FramedMultilineChunk::Incomplete(incomplete) => { - if io_buffer.has_retained_prefix() && io_buffer.has_remaining_fixed_writable_region() { + if let Some(io_buffer) = io_buffer.appendable_retained() { incomplete - .compact_packed_prefix_and_write_with_next_backend_chunk( + .append_to_packed_prefix_and_write_with_next_backend_chunk( writer, initial_len, &mut framer, From a695aded00efb82507d071dac141895145ce7b29 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Mon, 14 Sep 2026 19:24:25 -0600 Subject: [PATCH 05/23] Consume retained append permissions --- src/pool/buffer.rs | 107 +++++++++++++++++++++++++++---- src/pool/mod.rs | 2 +- src/session/multiline_framing.rs | 45 +++++++------ 3 files changed, 122 insertions(+), 32 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index c50dc41f..5b153d1f 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -39,11 +39,21 @@ pub struct PooledBuffer { counts_toward_pool: bool, } -/// Proof that a retained suffix has logical room for another backend read. -pub(crate) struct AppendableRetainedBuffer<'a> { +/// One permission to append a backend read to a retained buffer. +pub(crate) struct RetainedAppendPermit<'a> { buffer: &'a mut PooledBuffer, } +pub(crate) struct AppendedRead<'a> { + buffer: &'a mut PooledBuffer, + previous_len: usize, +} + +pub(crate) enum AppendOutcome<'a> { + Data(AppendedRead<'a>), + Eof(&'a mut PooledBuffer), +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum PooledBufferKind { Regular, @@ -280,12 +290,12 @@ impl PooledBuffer { self.buffer.retain_range(range); } - pub(crate) fn appendable_retained(&mut self) -> Option> { + pub(crate) fn retained_append_permit(&mut self) -> Option> { if !self.buffer.has_retained_prefix() || self.buffer.initialized_len() >= self.read_limit() { return None; } - Some(AppendableRetainedBuffer { buffer: self }) + Some(RetainedAppendPermit { buffer: self }) } #[cfg(test)] @@ -359,12 +369,38 @@ impl PooledBuffer { } } -impl<'a> AppendableRetainedBuffer<'a> { - pub(crate) async fn read_more(&mut self, reader: &mut R) -> std::io::Result +impl<'a> RetainedAppendPermit<'a> { + pub(crate) async fn read(self, reader: &mut R) -> std::io::Result> where R: AsyncRead + Unpin, { - self.buffer.read_more(reader).await + let Self { buffer } = self; + let previous_len = buffer.initialized(); + let read = buffer.read_more(reader).await?; + Ok(match read { + 0 => AppendOutcome::Eof(buffer), + _ => AppendOutcome::Data(AppendedRead { + buffer, + previous_len, + }), + }) + } +} + +impl<'a> AppendedRead<'a> { + #[must_use] + pub(crate) fn previous_len(&self) -> usize { + self.previous_len + } + + #[must_use] + pub(crate) fn as_new_bytes(&self) -> &[u8] { + &self.buffer.as_ref()[self.previous_len..] + } + + #[must_use] + pub(crate) fn total_len(&self) -> usize { + self.buffer.initialized() } pub(crate) fn into_inner(self) -> &'a mut PooledBuffer { @@ -1493,10 +1529,14 @@ mod tests { writer.write_all(b"0 ready\r\n").await.unwrap(); drop(writer); - let Some(mut appendable) = buffer.appendable_retained() else { + let Some(appendable) = buffer.retained_append_permit() else { panic!("retained prefix should have spare tail capacity"); }; - let read = appendable.read_more(&mut reader).await.unwrap(); + let AppendOutcome::Data(appended) = appendable.read(&mut reader).await.unwrap() else { + panic!("the reader has data available"); + }; + let read = appended.as_new_bytes().len(); + let buffer = appended.into_inner(); assert_eq!(read, 9); assert_eq!(buffer.as_ref(), b"220 ready\r\n"); @@ -1504,6 +1544,46 @@ mod tests { assert!(buffer.buffer.has_retained_prefix()); } + #[tokio::test] + async fn retained_append_result_owns_the_new_logical_window() { + let pool = BufferPool::new(BufferSize::try_new(8).unwrap(), 1); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(b"discard22unused"); + buffer.expose_initialized_range_without_copying(7..9); + + let mut reader = std::io::Cursor::new(b"abcdefghi"); + let Some(appendable) = buffer.retained_append_permit() else { + panic!("retained prefix should have logical room to append"); + }; + let AppendOutcome::Data(appended) = appendable.read(&mut reader).await.unwrap() else { + panic!("the reader has data available"); + }; + + assert_eq!(appended.previous_len(), 2); + assert_eq!(appended.as_new_bytes(), b"abcdef"); + let buffer = appended.into_inner(); + assert_eq!(buffer.as_ref(), b"22abcdef"); + assert!(buffer.retained_append_permit().is_none()); + } + + #[tokio::test] + async fn retained_append_result_distinguishes_eof_from_logical_exhaustion() { + let pool = BufferPool::new(BufferSize::try_new(8).unwrap(), 1); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(b"discard22unused"); + buffer.expose_initialized_range_without_copying(7..9); + + let mut reader = std::io::Cursor::new(Vec::::new()); + let Some(appendable) = buffer.retained_append_permit() else { + panic!("retained prefix should have logical room to append"); + }; + + let AppendOutcome::Eof(buffer) = appendable.read(&mut reader).await.unwrap() else { + panic!("an exhausted reader should produce the typed EOF outcome"); + }; + assert_eq!(buffer.initialized(), 2); + } + #[tokio::test] async fn read_more_compacts_a_retained_prefix_when_tail_is_full() { let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); @@ -1515,10 +1595,15 @@ mod tests { writer.write_all(b"y").await.unwrap(); drop(writer); - let Some(mut appendable) = buffer.appendable_retained() else { + let Some(appendable) = buffer.retained_append_permit() else { panic!("retained prefix should have logical room to append"); }; - let read = appendable.read_more(&mut reader).await.unwrap(); + let AppendOutcome::Data(appended) = appendable.read(&mut reader).await.unwrap() else { + panic!("the reader has data available"); + }; + let read = appended.as_new_bytes().len(); + assert_eq!(appended.as_new_bytes(), b"y"); + let buffer = appended.into_inner(); assert_eq!(read, 1); assert_eq!(buffer.as_ref(), b"xxy"); diff --git a/src/pool/mod.rs b/src/pool/mod.rs index 2d6ce1c9..0ecba585 100644 --- a/src/pool/mod.rs +++ b/src/pool/mod.rs @@ -10,7 +10,7 @@ pub mod health_check; pub mod prewarming; pub mod provider; -pub(crate) use buffer::AppendableRetainedBuffer; +pub(crate) use buffer::{AppendOutcome, RetainedAppendPermit}; pub use buffer::{ BufferPool, ChunkedResponse, HotPathAllocationMetricsSnapshot, PooledBuffer, hot_path_allocation_metrics_snapshot, reset_hot_path_allocation_metrics, diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 4b04366c..d22bd110 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -276,9 +276,8 @@ impl IncompleteMultilineWireChunk { async fn append_to_packed_prefix_and_write_with_next_backend_chunk( self, writer: &mut W, - current_len: usize, framer: &mut MultilineFramer, - mut io_buffer: crate::pool::AppendableRetainedBuffer<'_>, + io_buffer: crate::pool::RetainedAppendPermit<'_>, conn: &mut crate::stream::ConnectionStream, pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, @@ -287,27 +286,34 @@ impl IncompleteMultilineWireChunk { W: AsyncWrite + Unpin, { let initial_response_start = self.response.start; - let n = io_buffer.read_more(conn).await.map_err(|e| { + let appended = io_buffer.read(conn).await.map_err(|e| { crate::session::response_transfer::ResponseTransferError::Io( anyhow::Error::from(e).context("Failed to read remaining response body"), ) })?; - if n == 0 { - return Err( - crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received: self.response.len() as u64, - }, - ); - } - let io_buffer = io_buffer.into_inner(); - let total_len = current_len + n; + let appended = match appended { + crate::pool::AppendOutcome::Data(appended) => appended, + crate::pool::AppendOutcome::Eof(buffer) => { + let _ = buffer; + return Err( + crate::session::response_transfer::ResponseTransferError::BackendEof { + backend_id, + bytes_received: self.response.len() as u64, + }, + ); + } + }; + let previous_len = appended.previous_len(); + let total_len = appended.total_len(); + let frame = framer.frame_next_multiline_chunk(self, appended.as_new_bytes()); + let io_buffer = appended.into_inner(); - match framer.frame_next_multiline_chunk(self, &io_buffer[current_len..total_len]) { + match frame { FramedMultilineChunk::Complete(complete) => { - let combined_response = initial_response_start..current_len + complete.response.end; - let combined_next_response = current_len + complete.next_response_input.start - ..current_len + complete.next_response_input.end; + let combined_response = + initial_response_start..previous_len + complete.response.end; + let combined_next_response = previous_len + complete.next_response_input.start + ..previous_len + complete.next_response_input.end; write_response_chunk_preserving_suffix_on_error( writer, io_buffer, @@ -321,7 +327,7 @@ impl IncompleteMultilineWireChunk { } FramedMultilineChunk::Incomplete(incomplete) => { let combined_response = - initial_response_start..current_len + incomplete.response.end; + initial_response_start..previous_len + incomplete.response.end; let bytes_received = combined_response.len() as u64; if let Err(error) = writer .write_all(&io_buffer[combined_response.clone()]) @@ -1352,11 +1358,10 @@ where .await? } FramedMultilineChunk::Incomplete(incomplete) => { - if let Some(io_buffer) = io_buffer.appendable_retained() { + if let Some(io_buffer) = io_buffer.retained_append_permit() { incomplete .append_to_packed_prefix_and_write_with_next_backend_chunk( writer, - initial_len, &mut framer, io_buffer, conn, From b57a10cab6cc144bbff4f4aec18b85751d57658e Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Mon, 14 Sep 2026 21:33:22 -0600 Subject: [PATCH 06/23] Verify retained append edge cases --- Cargo.toml | 4 ++ benches/retained_buffer.rs | 34 +++++++++ src/pool/buffer.rs | 141 +++++++++++++++++++++++++++++++++++++ 3 files changed, 179 insertions(+) create mode 100644 benches/retained_buffer.rs diff --git a/Cargo.toml b/Cargo.toml index 95b4c30f..ec4d87c3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -144,6 +144,10 @@ harness = false name = "multiline_framing" harness = false +[[bench]] +name = "retained_buffer" +harness = false + [[bench]] name = "yenc_decoding" harness = false diff --git a/benches/retained_buffer.rs b/benches/retained_buffer.rs new file mode 100644 index 00000000..0382eaef --- /dev/null +++ b/benches/retained_buffer.rs @@ -0,0 +1,34 @@ +//! Retained-buffer append policy benchmarks. +//! +//! Run with: +//! `nix develop -c cargo bench --features framing-bench --bench retained_buffer` + +use divan::{Bencher, black_box}; +use nntp_proxy::pool::buffer::retained_append_benchmark; +use tokio::runtime::Builder; + +fn main() { + divan::main(); +} + +fn bench_tail(bencher: Bencher, physical_tail: usize) { + let runtime = Builder::new_current_thread().build().unwrap(); + bencher + .with_inputs(|| retained_append_benchmark(physical_tail)) + .bench_local_values(|case| black_box(runtime.block_on(case.drain()))); +} + +#[divan::bench(sample_count = 100, sample_size = 100)] +fn exact_full(bencher: Bencher) { + bench_tail(bencher, 0); +} + +#[divan::bench(sample_count = 100, sample_size = 100)] +fn one_byte_tail(bencher: Bencher) { + bench_tail(bencher, 1); +} + +#[divan::bench(sample_count = 100, sample_size = 100)] +fn roomy_tail(bencher: Bencher) { + bench_tail(bencher, 4096); +} diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index 5b153d1f..9e8cb4b3 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -54,6 +54,64 @@ pub(crate) enum AppendOutcome<'a> { Eof(&'a mut PooledBuffer), } +#[cfg(feature = "framing-bench")] +#[doc(hidden)] +pub struct RetainedAppendBenchmark { + buffer: PooledBuffer, + reader: std::io::Cursor>, +} + +#[cfg(feature = "framing-bench")] +#[doc(hidden)] +#[must_use] +pub fn retained_append_benchmark(physical_tail: usize) -> RetainedAppendBenchmark { + const CAPACITY: usize = 64 * 1024; + const RETAINED: usize = 32; + + assert!(physical_tail <= CAPACITY - RETAINED); + let initialized = CAPACITY - physical_tail; + let pool = BufferPool::new( + BufferSize::try_new(CAPACITY).expect("valid benchmark size"), + 1, + ); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(&vec![b'x'; initialized]); + buffer.expose_initialized_range_without_copying(initialized - RETAINED..initialized); + + RetainedAppendBenchmark { + buffer, + reader: std::io::Cursor::new(vec![b'y'; CAPACITY - RETAINED]), + } +} + +#[cfg(feature = "framing-bench")] +impl RetainedAppendBenchmark { + pub async fn drain(mut self) -> (usize, usize) { + let Some(permit) = self.buffer.retained_append_permit() else { + panic!("benchmark buffer should have logical room to append"); + }; + let AppendOutcome::Data(appended) = permit + .read(&mut self.reader) + .await + .expect("benchmark cursor read should succeed") + else { + panic!("benchmark cursor should contain data"); + }; + let first_read = appended.as_new_bytes().len(); + let buffer = appended.into_inner(); + let second_read = if self.reader.position() < self.reader.get_ref().len() as u64 { + buffer + .read_from(&mut self.reader) + .await + .expect("benchmark cursor read should succeed") + } else { + 0 + }; + assert_eq!(first_read + second_read, 64 * 1024 - 32); + (first_read, second_read) + } +} + #[derive(Clone, Copy, Debug, PartialEq, Eq)] enum PooledBufferKind { Regular, @@ -1610,6 +1668,89 @@ mod tests { assert_eq!(buffer.allocation_ptr(), allocation); } + #[tokio::test] + async fn retained_append_uses_a_tiny_physical_tail_without_compacting() { + let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(&vec![b'x'; 4095]); + let allocation = buffer.allocation_ptr(); + buffer.expose_initialized_range_without_copying(4063..4095); + let mut reader = std::io::Cursor::new(vec![b'y'; 4096]); + + let Some(permit) = buffer.retained_append_permit() else { + panic!("retained prefix should have logical room to append"); + }; + let AppendOutcome::Data(appended) = permit.read(&mut reader).await.unwrap() else { + panic!("the reader has data available"); + }; + + assert_eq!(appended.as_new_bytes(), b"y"); + let buffer = appended.into_inner(); + assert_eq!(buffer.initialized(), 33); + assert_eq!(buffer.allocation_ptr(), allocation); + assert_eq!(buffer.buffer.visible_from, 4063); + } + + #[tokio::test] + async fn cancelling_retained_append_after_compaction_preserves_visible_bytes() { + let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(&vec![b'x'; 4096]); + let allocation = buffer.allocation_ptr(); + buffer.expose_initialized_range_without_copying(4094..4096); + let (_writer, mut reader) = tokio::io::duplex(64); + + let Some(permit) = buffer.retained_append_permit() else { + panic!("retained prefix should have logical room to append"); + }; + let mut read = Box::pin(permit.read(&mut reader)); + match futures::poll!(read.as_mut()) { + std::task::Poll::Pending => {} + std::task::Poll::Ready(_) => panic!("empty open reader should remain pending"), + } + drop(read); + + assert_eq!(buffer.as_ref(), b"xx"); + assert_eq!(buffer.allocation_ptr(), allocation); + assert_eq!(buffer.buffer.visible_from, 0); + } + + #[tokio::test] + async fn retained_append_error_after_compaction_preserves_visible_bytes() { + struct ErrorReader; + + impl AsyncRead for ErrorReader { + fn poll_read( + self: Pin<&mut Self>, + _cx: &mut std::task::Context<'_>, + _buf: &mut ReadBuf<'_>, + ) -> std::task::Poll> { + std::task::Poll::Ready(Err(std::io::Error::from( + std::io::ErrorKind::ConnectionReset, + ))) + } + } + + let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(&vec![b'x'; 4096]); + let allocation = buffer.allocation_ptr(); + buffer.expose_initialized_range_without_copying(4094..4096); + + let Some(permit) = buffer.retained_append_permit() else { + panic!("retained prefix should have logical room to append"); + }; + let error = match permit.read(&mut ErrorReader).await { + Err(error) => error, + Ok(_) => panic!("error reader should fail the append"), + }; + + assert_eq!(error.kind(), std::io::ErrorKind::ConnectionReset); + assert_eq!(buffer.as_ref(), b"xx"); + assert_eq!(buffer.allocation_ptr(), allocation); + assert_eq!(buffer.buffer.visible_from, 0); + } + #[tokio::test] async fn fresh_read_after_a_range_view_restores_the_full_writable_region() { let pool = BufferPool::new(BufferSize::try_new(4096).unwrap(), 1); From 4764f6e5e49688d5ee117a3e14d4dc92d9bfb997 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 15 Sep 2026 08:45:45 -0600 Subject: [PATCH 07/23] Gate framing benchmarks by feature --- Cargo.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/Cargo.toml b/Cargo.toml index ec4d87c3..633ec463 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -143,10 +143,12 @@ harness = false [[bench]] name = "multiline_framing" harness = false +required-features = ["framing-bench"] [[bench]] name = "retained_buffer" harness = false +required-features = ["framing-bench"] [[bench]] name = "yenc_decoding" From ae3ff263d67f64c08c6f21584125f977ed5a7a14 Mon Sep 17 00:00:00 2001 From: Codex Date: Tue, 15 Sep 2026 10:26:08 -0600 Subject: [PATCH 08/23] Retire unfinished connections when exchanges are cancelled --- src/pool/connection_guard.rs | 33 +++++++++++++++++++++++---------- 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/src/pool/connection_guard.rs b/src/pool/connection_guard.rs index 464b1c0e..d4531d7c 100644 --- a/src/pool/connection_guard.rs +++ b/src/pool/connection_guard.rs @@ -211,10 +211,9 @@ impl ConnectionGuard { impl Drop for ConnectionGuard { fn drop(&mut self) { - debug_assert!( - self.conn.is_none(), - "ConnectionGuard dropped without explicit finalize" - ); + // Cancellation can drop an exchange between reads. Cleanup must run + // in debug builds too; panicking here would drop the pool object + // normally and return an unfinished connection to the pool. if let Some(conn) = self.conn.take() { tracing::debug!( connection_type = conn.connection_type(), @@ -393,7 +392,6 @@ mod tests { /// The guard shuts down the socket; pool recycle detects EOF /// and discards it; next `get()` creates a fresh TCP connection. /// Unknown/backend-error drop paths apply replacement cooldown. - #[cfg(not(debug_assertions))] #[tokio::test] async fn drop_without_release_forces_new_connection() { let (port, accept_count) = spawn_greeting_server().await; @@ -427,14 +425,29 @@ mod tests { ); } - #[cfg(debug_assertions)] #[tokio::test] - #[should_panic(expected = "ConnectionGuard dropped without explicit finalize")] - async fn drop_without_complete_success_panics_in_debug() { - let (port, _accept_count) = spawn_greeting_server().await; + async fn cancelling_an_unfinished_exchange_retires_its_connection() { + use std::future::Future; + use std::task::{Context, Poll, Waker}; + + let (port, accept_count) = spawn_greeting_server().await; let provider = make_provider(port); let conn = provider.get_pooled_connection().await.unwrap(); - let _guard = ConnectionGuard::new(conn, provider.clone()); + let mut guard = ConnectionGuard::new(conn, provider.clone()); + guard.stream_mut().write_all(b"DATE\r\n").await.unwrap(); + + // The future owns a used connection whose reply has not been consumed. + // Cancellation must retire it even if no input is locally queued. + let mut exchange = Box::pin(async move { + let _guard = guard; + std::future::pending::<()>().await; + }); + let mut context = Context::from_waker(Waker::noop()); + assert_eq!(exchange.as_mut().poll(&mut context), Poll::Pending); + drop(exchange); + + let _replacement = provider.get_pooled_connection().await.unwrap(); + assert_eq!(accept_count.load(Ordering::SeqCst), 2); } // ─── salvage_with_health_check notes ──────────────────────────────────── From 44f7d39a39abfaba73d70321657ecfb58cd7c987 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 15 Sep 2026 11:57:12 -0600 Subject: [PATCH 09/23] Bind response ownership to framing operations --- Cargo.toml | 2 +- docs/development.md | 37 + src/client/mod.rs | 104 +- src/pool/buffer.rs | 48 +- src/pool/connection_guard.rs | 6 +- src/pool/mod.rs | 2 +- src/protocol/article/mod.rs | 3 + src/protocol/mod.rs | 2 +- src/session/backend.rs | 203 +- src/session/handlers/command_execution.rs | 180 +- src/session/multiline_framing.rs | 2101 +++++++++------------ src/session/precheck.rs | 34 +- 12 files changed, 1101 insertions(+), 1621 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 633ec463..850f1332 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -111,7 +111,7 @@ framing-bench = [] tokio-console = ["dep:console-subscriber", "tokio/tracing"] # Runtime/task instrumentation for tokio-console; build with RUSTFLAGS="--cfg tokio_unstable" [lints.rust] -unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)'] } +unexpected_cfgs = { level = "warn", check-cfg = ['cfg(tokio_unstable)', 'cfg(response_contract, values(none(), "append_twice", "append_alias", "response_twice", "classified_buffer_reuse", "chunk_coordinate"))'] } [lints.clippy] branches_sharing_code = "deny" diff --git a/docs/development.md b/docs/development.md index 236e3f33..aa278a95 100644 --- a/docs/development.md +++ b/docs/development.md @@ -101,6 +101,43 @@ Current response responsibilities: ## Benchmarks +### Response ownership contracts + +`StreamingResponse` keeps the request's response shape, scanner, current window, +backend and pool together inside `multiline_framing.rs`. Write, observe and +capture consume this operation. `IsolatedMultilineResponse` owns the same +continuation relationship but rejects packed suffixes. Neither operation lends +an independently usable continuation to a handler. + +The scanner returns `ChunkConsumed`, a count relative to the latest push. +`WindowEnd` is an exclusive position in the buffer's visible window, not its +physical allocation. Translation occurs inside the framer, including after +compaction. Coordinates do not establish buffer identity: the operation's +exclusive borrow and the storage-owned `AppendedRead` establish that association. + +Storage does not classify responses. `RetainedAppendPermit::read` consumes one +permission and returns either EOF or an append result bound to the same buffer. +Logical exhaustion cannot create a permit; it is not EOF. Compaction policy and +pooled allocation reuse are unchanged. Ordinary forwarding borrows current +pooled bytes; only intentional capture/cache paths retain entire responses. + +nntpbench uses the same framing-versus-validation distinction and consuming +mutable operations, but freezes completed prefixes into immutable owners. +Its typed article accessor reuses validated layout and transformation metadata; +proxy pass-through does not manufacture that semantic guarantee. + +Cancellation may preserve bytes but does not prove a completed exchange. An +unreleased `ConnectionGuard` retires its connection on drop in every build mode. + +Explicit compiler-contract checks exercise real private types, with successful +controls and checked diagnostic codes: + +```bash +nix develop -c bash scripts/check-response-contracts.sh +``` + +### Measurement + Published benchmark numbers were intentionally removed from the docs until they are rerun. When you want fresh numbers: diff --git a/src/client/mod.rs b/src/client/mod.rs index 4174d4d7..1fbb7202 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -37,7 +37,7 @@ //! # fn process(_: &[u8]) {} //! ``` -use crate::pool::{BufferPool, ConnectionGuard, DeadpoolConnectionProvider, PooledBuffer}; +use crate::pool::{BufferPool, DeadpoolConnectionProvider, PooledBuffer}; use crate::protocol::{RequestContext, article_request, body_request, head_request, stat_request}; use crate::session::backend::execute_request_classified; use anyhow::{Context, Result}; @@ -46,7 +46,8 @@ use anyhow::{Context, Result}; /// /// Zero-allocation design using caller-provided buffer pool. /// Share one pool across multiple clients for minimal allocations. -/// Returns `PooledBuffer` - caller parses with `Article::parse()`. +/// Returns a framer-owned, terminator-stripped `PooledBuffer` - caller parses +/// with `Article::parse()`. #[derive(Clone)] pub struct NntpClient { conn_pool: DeadpoolConnectionProvider, @@ -67,7 +68,8 @@ impl NntpClient { /// Fetch article body (BODY command) /// - /// Returns `PooledBuffer` with the backend response bytes. + /// Returns `PooledBuffer` with the status line and payload bytes. The + /// multiline terminator has already been consumed by the framer. /// Parse with `Article::parse(&buffer, validate_yenc)`. /// /// # Arguments @@ -86,7 +88,8 @@ impl NntpClient { /// Fetch article headers (HEAD command) /// - /// Returns `PooledBuffer` with the backend response bytes. + /// Returns `PooledBuffer` with the status line and payload bytes. The + /// multiline terminator has already been consumed by the framer. /// Parse with `Article::parse(&buffer, false)`. /// /// # Arguments @@ -105,7 +108,8 @@ impl NntpClient { /// Fetch full article (ARTICLE command) /// - /// Returns `PooledBuffer` with the backend response bytes. + /// Returns `PooledBuffer` with the status line and payload bytes. The + /// multiline terminator has already been consumed by the framer. /// Parse with `Article::parse(&buffer, validate_yenc)`. /// /// # Arguments @@ -140,16 +144,16 @@ impl NntpClient { .checkout_connection_guard() .await .context("Failed to get connection from pool")?; - let mut buffer = self.buffer_pool.acquire(); + let buffer = self.buffer_pool.acquire(); - let response = execute_request_classified(conn.stream_mut(), &request, &mut buffer).await?; + let response = execute_request_classified(conn.stream_mut(), &request, buffer).await?; let Some(status_code) = response.status_code() else { anyhow::bail!("Invalid STAT response"); }; let result = Self::parse_stat_response(status_code); if result.is_ok() { - let completion = response.completion_proof(&request)?; + let completion = response.completion_proof()?; let _reusable = conn.complete_success(completion); } result @@ -176,51 +180,27 @@ impl NntpClient { .checkout_connection_guard() .await .context("Failed to get connection from pool")?; - let mut io_buffer = self.buffer_pool.acquire(); + let io_buffer = self.buffer_pool.acquire(); - let response = - execute_request_classified(conn.stream_mut(), &request, &mut io_buffer).await?; + let response = execute_request_classified(conn.stream_mut(), &request, io_buffer).await?; let Some(status_code) = response.status_code() else { anyhow::bail!("Invalid response from server"); }; Self::validate_response(status_code)?; - if request.has_response_body(status_code) { - return self - .fetch_captured_multiline_response(conn, io_buffer) - .await; - } - - let completion = response.completion_proof(&request)?; - let _reusable = conn.complete_success(completion); - Ok(io_buffer) - } - - async fn fetch_captured_multiline_response( - &self, - mut conn: ConnectionGuard, - mut io_buffer: PooledBuffer, - ) -> Result { - // This client helper is intentionally only an owner of the destination - // capture buffer. It delegates all multiline response completion and - // trailing-byte rejection to the backend/framer facade. - let mut capture = self.buffer_pool.acquire_capture(); - let completion = match crate::session::backend::capture_complete_multiline_response( - conn.stream_mut(), - &mut io_buffer, - &mut capture, - ) - .await + let (captured, completion) = match response + .capture_isolated(conn.stream_mut(), &self.buffer_pool) + .await { - Ok(completion) => completion, - Err(err) => { + Ok(result) => result, + Err(error) => { conn.fail_backend(); - return Err(err); + return Err(error); } }; let _reusable = conn.complete_success(completion); - Ok(capture) + Ok(captured) } /// Validate NNTP response status code @@ -429,10 +409,10 @@ mod tests { io_buffer: &mut PooledBuffer, capture: &mut PooledBuffer, ) -> Result<()> { - let _completion = - crate::session::backend::capture_complete_multiline_response(conn, io_buffer, capture) - .await?; - Ok(()) + crate::session::multiline_framing::capture_isolated_multiline_response( + conn, io_buffer, capture, + ) + .await } /// Verify the session response reader captures the complete response when it all @@ -559,7 +539,7 @@ mod tests { let buffer = client.fetch_head(&msg_id).await.unwrap(); - assert_eq!(&buffer[..], response); + assert_eq!(&buffer[..], &response[..response.len() - 5]); } #[tokio::test] @@ -571,7 +551,35 @@ mod tests { let buffer = client.fetch_body(&msg_id).await.unwrap(); - assert_eq!(&buffer[..], response); + assert_eq!(&buffer[..], &response[..response.len() - 5]); + let article = crate::protocol::Article::parse(&buffer, false).unwrap(); + assert_eq!(article.body, Some(&b"hello world"[..])); + } + + #[tokio::test] + async fn fetch_head_can_be_parsed_by_the_documented_article_api() { + let response = b"221 0 \r\nSubject: test\r\nFrom: tester\r\n.\r\n"; + let addr = spawn_fetch_test_server("HEAD ", response).await; + let client = make_test_client(addr); + let msg_id = crate::types::MessageId::new("".to_string()).unwrap(); + + let buffer = client.fetch_head(&msg_id).await.unwrap(); + let article = crate::protocol::Article::parse(&buffer, false).unwrap(); + assert_eq!(article.body, None); + assert_eq!(article.headers.unwrap().get("Subject"), Some(&b"test"[..])); + } + + #[tokio::test] + async fn fetch_body_preserves_wire_dot_stuffing_for_the_article_decoder() { + let response = b"222 0 \r\n..wire-dot\r\n.\r\n"; + let addr = spawn_fetch_test_server("BODY ", response).await; + let client = make_test_client(addr); + let msg_id = crate::types::MessageId::new("".to_string()).unwrap(); + + let buffer = client.fetch_body(&msg_id).await.unwrap(); + assert_eq!(&buffer[..], &response[..response.len() - 5]); + let article = crate::protocol::Article::parse(&buffer, false).unwrap(); + assert_eq!(article.body, Some(&b"..wire-dot"[..])); } #[tokio::test] @@ -587,6 +595,6 @@ mod tests { let buffer = client.fetch_body(&msg_id).await.unwrap(); - assert_eq!(&buffer[..], response); + assert_eq!(&buffer[..], &response[..response.len() - 5]); } } diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index 9e8cb4b3..ae2d070d 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -186,6 +186,11 @@ impl BufferStorage { self.visible_from = start; } + fn truncate_visible(&mut self, len: usize) { + assert!(len <= self.initialized_len()); + self.bytes.truncate(self.visible_from + len); + } + fn compact_visible(&mut self) { if self.visible_from == 0 { return; @@ -291,6 +296,10 @@ impl PooledBuffer { self.buffer.initialized_len() } + pub(crate) fn available_read_capacity(&self) -> usize { + self.read_limit().saturating_sub(self.initialized()) + } + #[inline] fn read_limit(&self) -> usize { if self.writable_len == 0 { @@ -384,6 +393,15 @@ impl PooledBuffer { self.buffer.clear(); } + /// Truncate initialized bytes without moving the backing allocation. + pub(crate) fn truncate_initialized(&mut self, len: usize) { + assert!( + len <= self.initialized(), + "truncate length exceeds initialized bytes" + ); + self.buffer.truncate_visible(len); + } + /// Append data to the buffer (accumulator mode) /// /// Used when `PooledBuffer` is acquired from capture pool for accumulating @@ -456,11 +474,7 @@ impl<'a> AppendedRead<'a> { &self.buffer.as_ref()[self.previous_len..] } - #[must_use] - pub(crate) fn total_len(&self) -> usize { - self.buffer.initialized() - } - + #[cfg(any(test, feature = "framing-bench"))] pub(crate) fn into_inner(self) -> &'a mut PooledBuffer { self.buffer } @@ -1800,6 +1814,7 @@ mod tests { let pool = BufferPool::new(BufferSize::try_new(8).unwrap(), 1); let mut buffer = pool.acquire(); buffer.copy_from_slice(b"22"); + assert_eq!(buffer.available_read_capacity(), 6); let (mut writer, mut reader) = tokio::io::duplex(64); writer.write_all(b"0 long response\r\n").await.unwrap(); @@ -1808,6 +1823,7 @@ mod tests { let read = buffer.read_more(&mut reader).await.unwrap(); assert_eq!(read, 6); assert_eq!(buffer.initialized(), 8); + assert_eq!(buffer.available_read_capacity(), 0); assert_eq!(&*buffer, b"220 long"); } @@ -2379,3 +2395,25 @@ mod tests { assert_eq!(buffer2.capacity() % 4096, 0); } } + +// These snippets compile against the actual private capabilities, not public +// stand-ins. The script checks a successful control and precise error codes. +#[cfg(response_contract)] +#[allow(dead_code)] +mod contracts { + use super::*; + + async fn append_permission(buffer: &mut PooledBuffer, reader: &mut (impl AsyncRead + Unpin)) { + let permit = buffer + .retained_append_permit() + .expect("retained writable input"); + let outcome = permit.read(reader).await.expect("read"); + #[cfg(response_contract = "append_twice")] + let _ = permit.read(reader).await; + if let AppendOutcome::Data(appended) = outcome { + #[cfg(response_contract = "append_alias")] + buffer.copy_from_slice(b"changed"); + std::hint::black_box(appended.as_new_bytes()); + } + } +} diff --git a/src/pool/connection_guard.rs b/src/pool/connection_guard.rs index d4531d7c..9e90cc85 100644 --- a/src/pool/connection_guard.rs +++ b/src/pool/connection_guard.rs @@ -78,6 +78,7 @@ const _SALVAGE_NO_LOOP: () = { pub struct ConnectionGuard { conn: Option, provider: DeadpoolConnectionProvider, + active: bool, } /// A connection whose last backend exchange was proven complete. @@ -92,6 +93,7 @@ impl ConnectionGuard { Self { conn: Some(conn), provider, + active: false, } } /// Return a checked-out connection that has not been used yet. @@ -99,6 +101,7 @@ impl ConnectionGuard { /// This is the only successful release available without response-completion /// evidence; callers cannot obtain protocol-stream access from this API. pub fn release_idle(mut self) { + assert!(!self.active, "cannot release a connection after I/O began"); drop(self.conn.take().expect("ConnectionGuard already consumed")); } @@ -160,7 +163,8 @@ impl ConnectionGuard { /// # Panics /// /// Panics if the guard has already been consumed. - pub(crate) const fn get_mut(&mut self) -> &mut PooledConnection { + pub(crate) fn get_mut(&mut self) -> &mut PooledConnection { + self.active = true; self.conn .as_mut() .expect("ConnectionGuard already consumed") diff --git a/src/pool/mod.rs b/src/pool/mod.rs index 0ecba585..f521d6fc 100644 --- a/src/pool/mod.rs +++ b/src/pool/mod.rs @@ -10,7 +10,7 @@ pub mod health_check; pub mod prewarming; pub mod provider; -pub(crate) use buffer::{AppendOutcome, RetainedAppendPermit}; +pub(crate) use buffer::AppendOutcome; pub use buffer::{ BufferPool, ChunkedResponse, HotPathAllocationMetricsSnapshot, PooledBuffer, hot_path_allocation_metrics_snapshot, reset_hot_path_allocation_metrics, diff --git a/src/protocol/article/mod.rs b/src/protocol/article/mod.rs index e7e7697f..dc81c859 100644 --- a/src/protocol/article/mod.rs +++ b/src/protocol/article/mod.rs @@ -28,6 +28,9 @@ pub struct Article<'a> { pub body: Option<&'a [u8]>, } +/// Consumer-facing view of an article whose framing has already been handled. +pub type ArticleView<'a> = Article<'a>; + impl<'a> TryFrom<&'a [u8]> for Article<'a> { type Error = ParseError; diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index a90eacf6..1a564bd6 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -15,7 +15,7 @@ mod response; mod responses; // Re-export article parsing types -pub use article::{Article, HeaderIter, Headers, ParseError, yenc}; +pub use article::{Article, ArticleView, HeaderIter, Headers, ParseError, yenc}; // Re-export response types and utilities pub(crate) use request::{ diff --git a/src/session/backend.rs b/src/session/backend.rs index 580ad48b..3d9e3431 100644 --- a/src/session/backend.rs +++ b/src/session/backend.rs @@ -12,21 +12,14 @@ //! should not rebuild command strings after request validation. use anyhow::Result; -use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; use crate::pool::PooledBuffer; use crate::protocol::RequestContext; pub(crate) use crate::session::multiline_framing::BackendResponseOrder; -/// Opaque result of reading enough backend bytes to classify one response. -/// -/// Callers can inspect status and single-line bytes through methods, but they do -/// not receive framing internals or boundary offsets. Any caller that needs to -/// transfer or capture a body must pass the same buffer back through this module. -pub(crate) struct BackendReadResult { - inner: BackendReadResultInner, -} +pub(crate) use crate::session::multiline_framing::ClassifiedResponse; #[must_use] pub(crate) struct BackendResponseComplete(()); @@ -69,11 +62,6 @@ pub(crate) enum SingleLineReplyReadError { Invalid { bytes_read: usize }, } -enum BackendReadResultInner { - Response(crate::session::multiline_framing::BackendResponseRead), - Invalid(crate::session::multiline_framing::ResponseReadError), -} - pub(crate) async fn read_single_line_reply( conn: &mut C, request: &RequestContext, @@ -148,85 +136,10 @@ pub(crate) const fn capabilities_without_authinfo_response() -> &'static [u8] { crate::session::multiline_framing::CAPABILITIES_WITHOUT_AUTHINFO_RESPONSE } -impl BackendReadResult { - #[must_use] - pub(crate) fn status_code(&self) -> Option { - match &self.inner { - BackendReadResultInner::Response(response) => Some(response.status_code()), - BackendReadResultInner::Invalid(_) => None, - } - } - - pub(crate) fn completion_proof( - &self, - request: &RequestContext, - ) -> Result { - let Some(status_code) = self.status_code() else { - anyhow::bail!("cannot prove completion for an invalid backend response"); - }; - if request.has_response_body(status_code) { - anyhow::bail!("multiline response requires framer-owned completion"); - } - Ok(BackendResponseComplete(())) - } - - #[must_use] - pub(crate) fn single_line_bytes<'a>(&self, buffer: &'a PooledBuffer) -> Option<&'a [u8]> { - match &self.inner { - BackendReadResultInner::Response(response) => response.single_line_bytes(buffer), - BackendReadResultInner::Invalid(_) => None, - } - } - - pub(crate) fn log_warnings( - &self, - buffer: &[u8], - client_addr: impl std::fmt::Display, - backend_id: crate::types::BackendId, - ) { - if let BackendReadResultInner::Invalid(err) = &self.inner { - err.log_warnings(buffer, client_addr, backend_id); - } - } -} - fn duration_micros_u64(duration: std::time::Duration) -> u64 { u64::try_from(duration.as_micros()).unwrap_or(u64::MAX) } -async fn read_until_backend_reply( - conn: &mut C, - request: &RequestContext, - buffer: &mut PooledBuffer, -) -> Result -where - C: AsyncReadExt + Unpin, -{ - loop { - match crate::session::multiline_framing::backend_response_read(request, buffer) { - Ok(response) => { - return Ok(BackendReadResult { - inner: BackendReadResultInner::Response(response), - }); - } - Err(err @ crate::session::multiline_framing::ResponseReadError::Invalid(_)) => { - return Ok(BackendReadResult { - inner: BackendReadResultInner::Invalid(err), - }); - } - Err(crate::session::multiline_framing::ResponseReadError::Incomplete) => { - let more = buffer.read_more(conn).await?; - if more == 0 { - anyhow::bail!( - "Backend EOF before complete backend response ({} bytes)", - buffer.initialized() - ); - } - } - } - } -} - /// Format a hex preview of response bytes for debugging Invalid responses /// /// # Arguments @@ -258,8 +171,8 @@ pub fn format_hex_preview(data: &[u8], max_bytes: usize) -> String { pub(crate) async fn execute_request_classified( conn: &mut C, request: &RequestContext, - buffer: &mut PooledBuffer, -) -> Result + mut buffer: PooledBuffer, +) -> Result where C: AsyncReadExt + AsyncWriteExt + Unpin, { @@ -270,7 +183,7 @@ where anyhow::bail!("Backend connection closed unexpectedly"); } - read_until_backend_reply(conn, request, buffer).await + ClassifiedResponse::read(conn, request, buffer).await } /// Read a response for a request that was already written as part of an @@ -279,7 +192,7 @@ pub(crate) async fn read_classified_response_for_already_sent_request( conn: &mut crate::stream::ConnectionStream, request: &RequestContext, pool: &crate::pool::BufferPool, -) -> Result<(BackendReadResult, PooledBuffer)> { +) -> Result { let mut buffer = crate::session::multiline_framing::take_queued_input_or_acquire_empty(conn, pool); if buffer.initialized() == 0 { @@ -289,15 +202,14 @@ pub(crate) async fn read_classified_response_for_already_sent_request( } } - let read = read_until_backend_reply(conn, request, &mut buffer).await?; - Ok((read, buffer)) + ClassifiedResponse::read(conn, request, buffer).await } pub(crate) async fn execute_request_classified_timed( conn: &mut C, request: &RequestContext, - buffer: &mut PooledBuffer, -) -> Result<(BackendReadResult, u64, u64, u64)> + mut buffer: PooledBuffer, +) -> Result<(ClassifiedResponse, u64, u64, u64)> where C: AsyncReadExt + AsyncWriteExt + Unpin, { @@ -312,7 +224,7 @@ where anyhow::bail!("Backend connection closed unexpectedly"); } - let response = read_until_backend_reply(conn, request, buffer).await?; + let response = ClassifiedResponse::read(conn, request, buffer).await?; let after_recv = Instant::now(); let send_elapsed = after_send.duration_since(start); let recv_elapsed = after_recv.duration_since(after_send); @@ -326,87 +238,6 @@ where )) } -/// Capture one backend response into owned pooled chunks. -/// -/// This is used for payload cache ingestion and other paths that explicitly need -/// ownership of the complete response. It still delegates all response boundary -/// detection to the framer. -pub(crate) async fn write_response_with_optional_capture( - request: &RequestContext, - buffer: &mut PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - writer: &mut W, - captured: &mut crate::pool::ChunkedResponse, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, -) -> Result<(u64, bool), crate::session::response_transfer::ResponseTransferError> -where - W: AsyncWrite + Unpin, -{ - crate::session::multiline_framing::write_response_with_optional_capture( - request, buffer, conn, writer, captured, pool, backend_id, - ) - .await -} - -/// Observe and drain one backend response without retaining its bytes. -/// -/// This is for retry/control paths that must consume the backend response to -/// keep the connection reusable but do not need ownership of the response body. -pub(crate) async fn observe_response( - request: &RequestContext, - buffer: &mut PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, -) -> Result { - crate::session::multiline_framing::observe_response(request, buffer, conn, pool, backend_id) - .await?; - Ok(BackendResponseComplete(())) -} - -/// Capture an isolated multiline response into a single pooled capture buffer. -/// -/// Used by client and precheck paths that issue commands outside the normal -/// per-command transfer loop. -pub(crate) async fn capture_complete_multiline_response( - conn: &mut crate::stream::ConnectionStream, - buffer: &mut PooledBuffer, - capture: &mut PooledBuffer, -) -> anyhow::Result { - crate::session::multiline_framing::capture_isolated_multiline_response(conn, buffer, capture) - .await?; - Ok(BackendResponseComplete(())) -} - -/// Capture an isolated multiline response into chunked pooled storage while it -/// remains within the framer-owned retention limit. -pub(crate) async fn capture_complete_multiline_response_chunked_optional( - conn: &mut crate::stream::ConnectionStream, - buffer: &mut PooledBuffer, - pool: &crate::pool::BufferPool, - response: &mut crate::pool::ChunkedResponse, -) -> anyhow::Result<(bool, BackendResponseComplete)> { - let retained = - crate::session::multiline_framing::capture_isolated_multiline_response_chunked_optional( - conn, buffer, pool, response, - ) - .await - .map_err(|err| anyhow::anyhow!("backend multiline response capture failed: {err:?}"))?; - Ok((retained, BackendResponseComplete(()))) -} - -/// Drain an isolated multiline response without retaining the bytes. -pub(crate) async fn observe_complete_multiline_response( - conn: &mut crate::stream::ConnectionStream, - buffer: &mut PooledBuffer, -) -> anyhow::Result { - crate::session::multiline_framing::observe_isolated_multiline_response(conn, buffer) - .await - .map_err(|err| anyhow::anyhow!("backend multiline response drain failed: {err:?}"))?; - Ok(BackendResponseComplete(())) -} - #[cfg(test)] mod tests { use super::*; @@ -474,10 +305,10 @@ mod tests { let mut stream = ChunkedStream::new(vec![b"20".to_vec(), b"0 OK\r\n".to_vec()]); let pool = crate::pool::BufferPool::for_tests(); - let mut buffer = pool.acquire(); + let buffer = pool.acquire(); let request = RequestContext::from_verb_args(b"DATE", b""); - let response = execute_request_classified(&mut stream, &request, &mut buffer) + let response = execute_request_classified(&mut stream, &request, buffer) .await .expect("send_request should handle partial reads"); let status_code = response @@ -495,10 +326,10 @@ mod tests { let mut stream = ChunkedStream::new(vec![b"111".to_vec(), b" 20260501173336\r\n".to_vec()]); let pool = crate::pool::BufferPool::for_tests(); - let mut buffer = pool.acquire(); + let buffer = pool.acquire(); let request = RequestContext::from_verb_args(b"DATE", b""); - let response = execute_request_classified(&mut stream, &request, &mut buffer) + let response = execute_request_classified(&mut stream, &request, buffer) .await .expect("send_request should read through complete backend response"); let status_code = response @@ -506,7 +337,7 @@ mod tests { .expect("DATE response should be valid"); assert_eq!(status_code, StatusCode::new(111)); assert!(!request.has_response_body(status_code)); - assert_eq!(&buffer[..buffer.initialized()], b"111 20260501173336\r\n"); + assert_eq!(response.received_bytes(), b"111 20260501173336\r\n"); } #[tokio::test] @@ -517,10 +348,10 @@ mod tests { let mut stream = ChunkedStream::new(chunks); let pool = crate::pool::BufferPool::for_tests(); - let mut buffer = pool.acquire(); + let buffer = pool.acquire(); let request = RequestContext::from_verb_args(b"GROUP", b"alt.test"); - let response = execute_request_classified(&mut stream, &request, &mut buffer) + let response = execute_request_classified(&mut stream, &request, buffer) .await .expect("send_request should handle single-byte reads"); let status_code = response diff --git a/src/session/handlers/command_execution.rs b/src/session/handlers/command_execution.rs index 24d1a2a8..23b1683a 100644 --- a/src/session/handlers/command_execution.rs +++ b/src/session/handlers/command_execution.rs @@ -116,26 +116,24 @@ pub(super) struct ResponseWriteParams<'a> { struct InvalidBackendResponseContext<'a> { provider: &'a crate::pool::DeadpoolConnectionProvider, request: &'a RequestContext, - buffer: &'a crate::pool::PooledBuffer, + buffer: &'a [u8], conn: crate::pool::ConnectionGuard, } pub(super) type PreparedBackendAttempt = Option<( crate::pool::ConnectionGuard, StatusCode, - crate::pool::PooledBuffer, + crate::session::backend::ClassifiedResponse, )>; type ExecutedBackendAttempt = ( crate::pool::ConnectionGuard, - crate::session::backend::BackendReadResult, - crate::pool::PooledBuffer, + crate::session::backend::ClassifiedResponse, Option, ); type BackendReadAttempt = ( - crate::session::backend::BackendReadResult, - crate::pool::PooledBuffer, + crate::session::backend::ClassifiedResponse, Option, ); @@ -290,7 +288,7 @@ impl ClientSession { return Ok(BackendAttemptResult::BackendUnavailable); }; - let Some((mut conn, status_code, mut buffer)) = self + let Some((mut conn, status_code, buffer)) = self .prepare_backend_attempt(provider, &backend, request, state, is_retry_attempt) .await? else { @@ -317,15 +315,10 @@ impl ClientSession { ) .await; } - let completion = crate::session::backend::observe_response( - request, - &mut buffer, - conn.stream_mut(), - &self.buffer_pool, - backend_id, - ) - .await - .map_err(SessionError::from)?; + let completion = buffer + .observe(conn.stream_mut(), &self.buffer_pool, backend_id) + .await + .map_err(SessionError::from)?; self.release_or_reuse_connection( conn, backend_id, @@ -422,29 +415,22 @@ impl ClientSession { Ok(conn) => conn, Err(_) => return RetryStatProbeOutcome::Unavailable(backend_id), }; - let mut buffer = self.buffer_pool.acquire(); - let read = backend::execute_request_classified( - conn.stream_mut(), - &stat_request, - &mut buffer, - ) - .await; - let status_code = match read { - Ok(read) => read.status_code(), + let buffer = self.buffer_pool.acquire(); + let read = + backend::execute_request_classified(conn.stream_mut(), &stat_request, buffer) + .await; + let response = match read { + Ok(read) => read, Err(_) => { conn.fail_backend(); return RetryStatProbeOutcome::Unavailable(backend_id); } }; - let completion = match crate::session::backend::observe_response( - &stat_request, - &mut buffer, - conn.stream_mut(), - &self.buffer_pool, - backend_id, - ) - .await + let status_code = response.status_code(); + let completion = match response + .observe(conn.stream_mut(), &self.buffer_pool, backend_id) + .await { Ok(completion) => completion, Err(_) => { @@ -556,28 +542,24 @@ impl ClientSession { Ok(conn) => conn, Err(_) => return, }; - let mut buffer = buffer_pool.acquire(); + let buffer = buffer_pool.acquire(); let read = backend::execute_request_classified( conn.stream_mut(), &stat_request, - &mut buffer, + buffer, ) .await; - let status_code = match read { - Ok(read) => read.status_code(), + let response = match read { + Ok(read) => read, Err(_) => { conn.fail_backend(); return; } }; - let completion = match crate::session::backend::observe_response( - &stat_request, - &mut buffer, - conn.stream_mut(), - &buffer_pool, - backend_id, - ) - .await + let status_code = response.status_code(); + let completion = match response + .observe(conn.stream_mut(), &buffer_pool, backend_id) + .await { Ok(completion) => completion, Err(_) => { @@ -605,7 +587,7 @@ impl ClientSession { conn: crate::pool::ConnectionGuard, client_writer: &mut crate::session::ClientWriter, backend: &ArticleBackend, - buffer: crate::pool::PooledBuffer, + buffer: crate::session::backend::ClassifiedResponse, params: ResponseWriteParams<'_>, backend_connection: &mut Option, ) -> Result { @@ -639,7 +621,7 @@ impl ClientSession { request_wire_len, "Preparing direct backend attempt" ); - let (conn, read_status, buffer, timings) = match retry_once!( + let (conn, buffer, timings) = match retry_once!( self.execute_backend_attempt( provider, backend, @@ -672,10 +654,10 @@ impl ClientSession { } *state.client_to_backend_bytes = state.client_to_backend_bytes.add(request_wire_len); - let status_code = match read_status.status_code() { + let status_code = match buffer.status_code() { Some(status_code) => status_code, None => { - read_status.log_warnings(&buffer, self.client_addr, backend_id); + buffer.log_warnings(self.client_addr, backend_id); trace!( client = %self.client_addr, backend = backend_id.as_index(), @@ -688,7 +670,7 @@ impl ClientSession { InvalidBackendResponseContext { provider, request, - buffer: &buffer, + buffer: buffer.received_bytes(), conn, }, ); @@ -703,7 +685,7 @@ impl ClientSession { command_verb = ?request.verb(), msg_id = ?request.message_id_value(), status_code = status_code.as_u16(), - backend_read_bytes = buffer.initialized(), + backend_read_bytes = buffer.received_bytes().len(), availability_missing_bits = format_args!("{:08b}", state.availability.missing_bits()), "Backend attempt received classifiable response bytes" ); @@ -751,7 +733,7 @@ impl ClientSession { mut conn: crate::pool::ConnectionGuard, client_write: &mut W, backend: &ArticleBackend, - backend_bytes: crate::pool::PooledBuffer, + backend_bytes: crate::session::backend::ClassifiedResponse, params: ResponseWriteParams<'_>, backend_connection: Option<&mut Option>, ) -> Result @@ -1163,7 +1145,7 @@ impl ClientSession { "Running STAT miss probe before backend article fetch" ); - let (probe_response, mut probe_buffer, probe_timings) = self + let (probe_response, probe_timings) = self .execute_and_read_response(guard.stream_mut(), backend, &stat_request) .await .map_err(|BackendReadAttemptError::Backend(e)| e)?; @@ -1178,17 +1160,12 @@ impl ClientSession { msg_id = ?request.message_id_value(), "STAT miss probe returned 430; skipping backend article fetch" ); - Ok((probe_response, probe_buffer, probe_timings)) + Ok((probe_response, probe_timings)) } else { - let _completion = crate::session::backend::observe_response( - &stat_request, - &mut probe_buffer, - guard.stream_mut(), - &self.buffer_pool, - backend.backend_id(), - ) - .await - .map_err(SessionError::from)?; + let _completion = probe_response + .observe(guard.stream_mut(), &self.buffer_pool, backend.backend_id()) + .await + .map_err(SessionError::from)?; let _ = probe_timings; self.execute_and_read_response(guard.stream_mut(), backend, request) .await @@ -1199,7 +1176,7 @@ impl ClientSession { }; match result { - Ok((read_status, buffer, timings)) => Ok((guard, read_status, buffer, timings)), + Ok((response, timings)) => Ok((guard, response, timings)), Err(BackendReadAttemptError::Backend(e)) => { debug!( client = %self.client_addr, @@ -1228,15 +1205,15 @@ impl ClientSession { self.metrics.record_command(backend_id); self.metrics.user_command(self.username()); - let mut buffer = self.buffer_pool.acquire(); + let buffer = self.buffer_pool.acquire(); let (response, timings) = if should_sample_backend_timing() { let (response, ttfb, send, recv) = - backend::execute_request_classified_timed(conn, request, &mut buffer) + backend::execute_request_classified_timed(conn, request, buffer) .await .map_err(BackendReadAttemptError::Backend)?; (response, Some((ttfb, send, recv))) } else { - let response = backend::execute_request_classified(conn, request, &mut buffer) + let response = backend::execute_request_classified(conn, request, buffer) .await .map_err(BackendReadAttemptError::Backend)?; (response, None) @@ -1249,7 +1226,7 @@ impl ClientSession { "Backend response classification completed for direct attempt" ); - Ok((response, buffer, timings)) + Ok((response, timings)) } pub(super) async fn forward_response_for_already_sent_request( @@ -1265,15 +1242,15 @@ impl ClientSession { W: AsyncWrite + Unpin, { let backend_id = backend.backend_id(); - let (read, mut buffer) = backend::read_classified_response_for_already_sent_request( + let buffer = backend::read_classified_response_for_already_sent_request( conn.stream_mut(), request, &self.buffer_pool, ) .await .map_err(AlreadySentResponseError::Read)?; - let Some(status_code) = read.status_code() else { - read.log_warnings(&buffer, self.client_addr, backend_id); + let Some(status_code) = buffer.status_code() else { + buffer.log_warnings(self.client_addr, backend_id); return Err(AlreadySentResponseError::Read(anyhow::anyhow!( "backend returned an invalid response to an already-sent request" ))); @@ -1291,15 +1268,10 @@ impl ClientSession { ) .await; } - let completion = backend::observe_response( - request, - &mut buffer, - conn.stream_mut(), - &self.buffer_pool, - backend_id, - ) - .await - .map_err(AlreadySentResponseError::Transfer)?; + let completion = buffer + .observe(conn.stream_mut(), &self.buffer_pool, backend_id) + .await + .map_err(AlreadySentResponseError::Transfer)?; self.send_430_to_client(client_write, backend_to_client_bytes) .await .map_err(|error| { @@ -1371,7 +1343,7 @@ impl ClientSession { pooled_conn: &mut crate::stream::ConnectionStream, client_write: &mut W, backend: &ArticleBackend, - backend_bytes: crate::pool::PooledBuffer, + backend_bytes: crate::session::backend::ClassifiedResponse, params: ResponseWriteParams<'_>, ) -> Result<(u64, crate::session::backend::BackendResponseComplete), ResponseTransferError> where @@ -1408,7 +1380,6 @@ impl ClientSession { client_write, backend_id, backend_bytes, - params, ) .await?; (bytes, None) @@ -1491,21 +1462,14 @@ impl ClientSession { pooled_conn: &mut crate::stream::ConnectionStream, client_write: &mut W, backend_id: BackendId, - mut backend_bytes: crate::pool::PooledBuffer, - params: ResponseWriteParams<'_>, + backend_bytes: crate::session::backend::ClassifiedResponse, ) -> Result where W: AsyncWrite + Unpin, { - crate::session::multiline_framing::write_response( - params.request, - &mut backend_bytes, - pooled_conn, - client_write, - &self.buffer_pool, - backend_id, - ) - .await + backend_bytes + .write(pooled_conn, client_write, &self.buffer_pool, backend_id) + .await } #[allow(clippy::too_many_arguments)] @@ -1514,7 +1478,7 @@ impl ClientSession { pooled_conn: &mut crate::stream::ConnectionStream, client_write: &mut W, backend_id: BackendId, - mut backend_bytes: crate::pool::PooledBuffer, + backend_bytes: crate::session::backend::ClassifiedResponse, params: ResponseWriteParams<'_>, ) -> Result<(u64, Option), ResponseTransferError> where @@ -1525,10 +1489,8 @@ impl ClientSession { // framer-owned retention limit. Larger responses are streamed through // without cache insertion so normal delivery and backend reuse can // continue. - let (bytes_written, retained) = - crate::session::backend::write_response_with_optional_capture( - params.request, - &mut backend_bytes, + let (bytes_written, retained) = backend_bytes + .capture_and_write( pooled_conn, client_write, &mut captured, @@ -1536,7 +1498,6 @@ impl ClientSession { backend_id, ) .await?; - drop(backend_bytes); self.log_body_response_written(backend_id, params, bytes_written as usize); if retained && let Some(article_request) = params.article_request { debug!( @@ -2991,6 +2952,13 @@ mod tests { let mut backend_bytes = BufferPool::new(BufferSize::try_new(8192).unwrap(), 1).acquire(); backend_bytes.copy_from_slice(b"223 0 status\r\n"); + let backend_bytes = crate::session::backend::ClassifiedResponse::read( + &mut tokio::io::empty(), + &request, + backend_bytes, + ) + .await + .unwrap(); let err = session .write_successful_backend_response( guard, @@ -3039,6 +3007,13 @@ mod tests { let mut backend_bytes = BufferPool::new(BufferSize::try_new(8192).unwrap(), 1).acquire(); backend_bytes.copy_from_slice(b"220 Article follows\r\nbody\r\n.\r\n"); + let backend_bytes = crate::session::backend::ClassifiedResponse::read( + &mut tokio::io::empty(), + &request, + backend_bytes, + ) + .await + .unwrap(); let err = session .write_successful_backend_response( guard, @@ -3088,6 +3063,13 @@ mod tests { b"220 Article follows\r\nbody\r\n.\r\n223 0 status\r\n", ); + let backend_bytes = crate::session::backend::ClassifiedResponse::read( + &mut tokio::io::empty(), + &request, + backend_bytes, + ) + .await + .unwrap(); let err = session .write_successful_backend_response( guard, diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index d22bd110..74d802b8 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -120,17 +120,40 @@ impl CompleteMultilinePayloadSplit { } } -/// Complete multiline frame result for the current backend buffer. +/// Complete response window in the current backend buffer. /// /// The ranges remain private to this module so callers cannot make their own /// response-boundary decisions after framing. #[derive(Debug, PartialEq, Eq)] -struct CompleteMultilineWireChunk { +struct CompleteResponseWindow { response: Range, next_response_input: Range, } -impl CompleteMultilineWireChunk { +impl CompleteResponseWindow { + async fn write_from( + &self, + writer: &mut W, + io_buffer: &mut crate::pool::PooledBuffer, + conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, + total_len: usize, + ) -> Result + where + W: AsyncWrite + Unpin, + { + write_response_chunk_preserving_suffix_on_error( + writer, + io_buffer, + conn, + pool, + total_len, + &self.response, + &self.next_response_input, + ) + .await + } + fn extend_capture_from(&self, source: &[u8], capture: &mut crate::pool::PooledBuffer) { capture.extend_from_slice(&source[self.response.clone()]); } @@ -157,29 +180,6 @@ impl CompleteMultilineWireChunk { Ok(()) } - async fn write_from( - &self, - writer: &mut W, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - total_len: usize, - ) -> Result - where - W: AsyncWrite + Unpin, - { - write_response_chunk_preserving_suffix_on_error( - writer, - io_buffer, - conn, - pool, - total_len, - &self.response, - &self.next_response_input, - ) - .await - } - fn push_from_buffer( &self, io_buffer: &mut crate::pool::PooledBuffer, @@ -267,311 +267,11 @@ where } #[derive(Debug, PartialEq, Eq)] -struct IncompleteMultilineWireChunk { +struct IncompleteResponseWindow { response: Range, } -impl IncompleteMultilineWireChunk { - #[allow(clippy::too_many_arguments)] - async fn append_to_packed_prefix_and_write_with_next_backend_chunk( - self, - writer: &mut W, - framer: &mut MultilineFramer, - io_buffer: crate::pool::RetainedAppendPermit<'_>, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, - ) -> Result - where - W: AsyncWrite + Unpin, - { - let initial_response_start = self.response.start; - let appended = io_buffer.read(conn).await.map_err(|e| { - crate::session::response_transfer::ResponseTransferError::Io( - anyhow::Error::from(e).context("Failed to read remaining response body"), - ) - })?; - let appended = match appended { - crate::pool::AppendOutcome::Data(appended) => appended, - crate::pool::AppendOutcome::Eof(buffer) => { - let _ = buffer; - return Err( - crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received: self.response.len() as u64, - }, - ); - } - }; - let previous_len = appended.previous_len(); - let total_len = appended.total_len(); - let frame = framer.frame_next_multiline_chunk(self, appended.as_new_bytes()); - let io_buffer = appended.into_inner(); - - match frame { - FramedMultilineChunk::Complete(complete) => { - let combined_response = - initial_response_start..previous_len + complete.response.end; - let combined_next_response = previous_len + complete.next_response_input.start - ..previous_len + complete.next_response_input.end; - write_response_chunk_preserving_suffix_on_error( - writer, - io_buffer, - conn, - pool, - total_len, - &combined_response, - &combined_next_response, - ) - .await - } - FramedMultilineChunk::Incomplete(incomplete) => { - let combined_response = - initial_response_start..previous_len + incomplete.response.end; - let bytes_received = combined_response.len() as u64; - if let Err(error) = writer - .write_all(&io_buffer[combined_response.clone()]) - .await - { - consume_remaining_multiline_response( - framer, - incomplete, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await?; - return Err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect( - error, - ), - ); - } - let mut stats = ResponseWriteStats::default(); - stats.add_chunk(combined_response.len()); - incomplete - .write_after_current_chunk( - writer, - framer, - io_buffer, - conn, - pool, - backend_id, - stats, - bytes_received, - ) - .await - } - } - } - - #[allow(clippy::too_many_arguments)] - async fn write_and_consume_response( - self, - writer: &mut W, - current_len: usize, - framer: &mut MultilineFramer, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, - ) -> Result - where - W: AsyncWrite + Unpin, - { - let mut bytes_received = 0; - let mut stats = { - let response = &io_buffer[..current_len][self.response.clone()]; - bytes_received += response.len() as u64; - if let Err(e) = writer.write_all(response).await { - consume_remaining_multiline_response( - framer, - self, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await?; - return Err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect(e), - ); - } - let mut stats = ResponseWriteStats::default(); - stats.add_chunk(response.len()); - stats - }; - - let mut continuation = Some(self); - loop { - let n = io_buffer.read_from(conn).await.map_err(|e| { - crate::session::response_transfer::ResponseTransferError::Io( - anyhow::Error::from(e).context("Failed to read remaining response body"), - ) - })?; - if n == 0 { - return Err( - crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received, - }, - ); - } - let prior = continuation - .take() - .expect("incomplete response continuation must be consumed by the framer"); - match framer.frame_next_multiline_chunk(prior, &io_buffer[..n]) { - FramedMultilineChunk::Complete(complete) => { - return complete - .write_from(writer, io_buffer, conn, pool, n) - .await - .map(|next| { - stats += next; - stats - }); - } - FramedMultilineChunk::Incomplete(incomplete) => { - let response = &io_buffer[incomplete.response.clone()]; - bytes_received += response.len() as u64; - if let Err(e) = writer.write_all(response).await { - consume_remaining_multiline_response( - framer, - incomplete, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await?; - return Err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect( - e, - ), - ); - } - stats.add_chunk(response.len()); - continuation = Some(incomplete); - } - } - } - } - - #[allow(clippy::too_many_arguments)] - async fn write_current_and_after( - self, - writer: &mut W, - current_len: usize, - framer: &mut MultilineFramer, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, - mut stats: ResponseWriteStats, - mut bytes_received: u64, - ) -> Result - where - W: AsyncWrite + Unpin, - { - let response = &io_buffer[..current_len][self.response.clone()]; - bytes_received += response.len() as u64; - if let Err(e) = writer.write_all(response).await { - consume_remaining_multiline_response( - framer, - self, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await?; - return Err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect(e), - ); - } - stats.add_chunk(response.len()); - self.write_after_current_chunk( - writer, - framer, - io_buffer, - conn, - pool, - backend_id, - stats, - bytes_received, - ) - .await - } - - #[allow(clippy::too_many_arguments)] - async fn write_after_current_chunk( - self, - writer: &mut W, - framer: &mut MultilineFramer, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, - mut stats: ResponseWriteStats, - mut bytes_received: u64, - ) -> Result - where - W: AsyncWrite + Unpin, - { - let mut continuation = Some(self); - loop { - let n = io_buffer.read_from(conn).await.map_err(|e| { - crate::session::response_transfer::ResponseTransferError::Io( - anyhow::Error::from(e).context("Failed to read remaining response body"), - ) - })?; - if n == 0 { - return Err( - crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received, - }, - ); - } - let prior = continuation - .take() - .expect("incomplete response continuation must be consumed by the framer"); - match framer.frame_next_multiline_chunk(prior, &io_buffer[..n]) { - FramedMultilineChunk::Complete(complete) => { - return complete - .write_from(writer, io_buffer, conn, pool, n) - .await - .map(|next| { - stats += next; - stats - }); - } - FramedMultilineChunk::Incomplete(incomplete) => { - let response = &io_buffer[incomplete.response.clone()]; - bytes_received += response.len() as u64; - if let Err(e) = writer.write_all(response).await { - consume_remaining_multiline_response( - framer, - incomplete, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await?; - return Err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect( - e, - ), - ); - } - stats.add_chunk(response.len()); - continuation = Some(incomplete); - } - } - } - } - +impl IncompleteResponseWindow { fn extend_capture_from(&self, source: &[u8], capture: &mut crate::pool::PooledBuffer) { capture.extend_from_slice(&source[self.response.clone()]); } @@ -587,10 +287,50 @@ impl IncompleteMultilineWireChunk { } } +/// Count consumed from one scanner push, never from an accumulated window. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ChunkConsumed(usize); + +/// Exclusive position in the current visible buffer, independent of physical +/// allocation placement. This coordinate alone does not identify a buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct WindowEnd(usize); + +impl WindowEnd { + fn after_chunk(self, consumed: ChunkConsumed) -> Self { + Self(self.0 + consumed.0) + } +} + +#[derive(Debug, PartialEq, Eq)] +enum ChunkProgress { + Complete(ChunkConsumed), + Incomplete, +} + +impl ChunkProgress { + // Translation lives here; operation contexts supply the origin from their + // own buffer-bound append result, never from an application caller. + fn in_window(self, origin: WindowEnd, window_len: usize) -> ResponseWindow { + match self { + Self::Complete(consumed) => { + let end = origin.after_chunk(consumed); + ResponseWindow::Complete(CompleteResponseWindow { + response: 0..end.0, + next_response_input: end.0..window_len, + }) + } + Self::Incomplete => ResponseWindow::Incomplete(IncompleteResponseWindow { + response: 0..window_len, + }), + } + } +} + #[derive(Debug, PartialEq, Eq)] -enum FramedMultilineChunk { - Complete(CompleteMultilineWireChunk), - Incomplete(IncompleteMultilineWireChunk), +enum ResponseWindow { + Complete(CompleteResponseWindow), + Incomplete(IncompleteResponseWindow), } #[derive(Debug, Clone, PartialEq, Eq)] @@ -600,58 +340,11 @@ struct FramedSingleLineChunk { } impl FramedSingleLineChunk { - async fn write_from( - &self, - writer: &mut W, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - total_len: usize, - ) -> Result - where - W: AsyncWrite + Unpin, - { - write_response_chunk_preserving_suffix_on_error( - writer, - io_buffer, - conn, - pool, - total_len, - &self.response, - &self.next_response_input, - ) - .await - } - - fn push_from_buffer( - &self, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - response: &mut crate::pool::ChunkedResponse, - pool: &crate::pool::BufferPool, - ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - let total_len = io_buffer.initialized(); - let old = std::mem::replace(io_buffer, pool.acquire()); - if self.next_response_input.start < total_len { - conn.queue_pending_bytes_first(&old[self.next_response_input.clone()]) - .map_err(crate::session::response_transfer::ResponseTransferError::Io)?; - } - response.push_buffer_range(old, self.response.clone()); - Ok(()) - } - - fn observe_from_buffer( - &self, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - let total_len = io_buffer.initialized(); - if self.next_response_input.start < total_len { - let old = std::mem::replace(io_buffer, pool.acquire()); - conn.queue_pooled_pending_bytes_first(old, self.next_response_input.clone()) - .map_err(crate::session::response_transfer::ResponseTransferError::Io)?; - } + fn require_isolated(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.next_response_input.is_empty(), + "unexpected bytes after isolated single-line response" + ); Ok(()) } } @@ -904,48 +597,224 @@ impl ResponseFrame { } } -pub(crate) enum BackendResponseRead { - SingleLine { - status: crate::protocol::StatusCode, - response: Range, - }, - Multiline { - status: crate::protocol::StatusCode, - }, +/// Classification retains its input and request-scoped shape. The buffer cannot +/// be replaced or mutated before the consuming transfer operation uses it. +pub(crate) struct ClassifiedResponse { + buffer: crate::pool::PooledBuffer, + frame: Result, } -impl BackendResponseRead { - #[inline] - #[must_use] - pub(crate) const fn status_code(&self) -> crate::protocol::StatusCode { - match self { - Self::SingleLine { status, .. } | Self::Multiline { status } => *status, +impl ClassifiedResponse { + pub(crate) async fn read( + conn: &mut C, + request: &crate::protocol::RequestContext, + mut buffer: crate::pool::PooledBuffer, + ) -> anyhow::Result { + loop { + let frame = ResponseFrame::parse(request, &buffer); + match frame { + Err(ResponseReadError::Incomplete) => { + let more = buffer.read_more(conn).await?; + if more == 0 { + if buffer.available_read_capacity() == 0 { + anyhow::bail!( + "Backend response exceeded the read buffer capacity ({} bytes)", + buffer.initialized() + ); + } + anyhow::bail!( + "Backend EOF before complete backend response ({} bytes)", + buffer.initialized() + ); + } + } + frame => return Ok(Self { buffer, frame }), + } } } - #[must_use] - pub(crate) fn single_line_bytes<'a>( + pub(crate) fn status_code(&self) -> Option { + match &self.frame { + Ok(ResponseFrame::SingleLine { status, .. } | ResponseFrame::Multiline { status }) => { + Some(*status) + } + Err(_) => None, + } + } + + pub(crate) fn single_line_bytes(&self) -> Option<&[u8]> { + match &self.frame { + Ok(ResponseFrame::SingleLine { framed, .. }) => { + Some(&self.buffer[framed.response.clone()]) + } + Ok(ResponseFrame::Multiline { .. }) | Err(_) => None, + } + } + + pub(crate) fn received_bytes(&self) -> &[u8] { + &self.buffer + } + + pub(crate) async fn capture_isolated( + mut self, + conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, + ) -> anyhow::Result<( + crate::pool::PooledBuffer, + crate::session::backend::BackendResponseComplete, + )> { + match self.frame { + Ok(ResponseFrame::SingleLine { framed, .. }) => { + framed.require_isolated()?; + Ok(( + self.buffer, + crate::session::backend::BackendResponseComplete::response(), + )) + } + Ok(ResponseFrame::Multiline { .. }) => { + let mut capture = pool.acquire_capture(); + IsolatedMultilineResponse::begin(conn, &mut self.buffer) + .map_err(isolated_multiline_error)? + .capture_into(&mut capture) + .await?; + // Framing owns the already-proven terminator. The public + // article parser receives the status line and payload only. + strip_captured_terminator(&mut capture)?; + Ok(( + capture, + crate::session::backend::BackendResponseComplete::response(), + )) + } + Err(error) => anyhow::bail!("cannot capture invalid backend response: {error:?}"), + } + } + + pub(crate) async fn capture_isolated_chunked_optional( + mut self, + conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, + captured: &mut crate::pool::ChunkedResponse, + ) -> anyhow::Result<(bool, crate::session::backend::BackendResponseComplete)> { + match self.frame { + Ok(ResponseFrame::Multiline { .. }) => {} + _ => anyhow::bail!("isolated multiline capture requires a multiline response"), + } + let retained = capture_isolated_multiline_response_chunked_optional( + conn, + &mut self.buffer, + pool, + captured, + ) + .await + .map_err(|error| anyhow::anyhow!("backend multiline response capture failed: {error:?}"))?; + Ok(( + retained, + crate::session::backend::BackendResponseComplete::response(), + )) + } + + pub(crate) async fn observe_isolated( + mut self, + conn: &mut crate::stream::ConnectionStream, + ) -> anyhow::Result { + match self.frame { + Ok(ResponseFrame::SingleLine { framed, .. }) => framed.require_isolated()?, + Ok(ResponseFrame::Multiline { .. }) => { + observe_isolated_multiline_response(conn, &mut self.buffer) + .await + .map_err(|error| { + anyhow::anyhow!("backend multiline response drain failed: {error:?}") + })? + } + Err(error) => anyhow::bail!("cannot observe invalid backend response: {error:?}"), + } + Ok(crate::session::backend::BackendResponseComplete::response()) + } + + pub(crate) fn log_warnings( &self, - buffer: &'a crate::pool::PooledBuffer, - ) -> Option<&'a [u8]> { - match self { - Self::SingleLine { response, .. } => Some(&buffer[response.clone()]), - Self::Multiline { .. } => None, + client_addr: impl std::fmt::Display, + backend_id: crate::types::BackendId, + ) { + if let Err(error) = &self.frame { + error.log_warnings(&self.buffer, client_addr, backend_id); + } + } + + pub(crate) fn completion_proof( + &self, + ) -> anyhow::Result { + match &self.frame { + Ok(ResponseFrame::SingleLine { framed, .. }) => { + framed.require_isolated()?; + Ok(crate::session::backend::BackendResponseComplete::response()) + } + Ok(ResponseFrame::Multiline { .. }) => { + anyhow::bail!("multiline response requires framer-owned completion") + } + Err(_) => anyhow::bail!("cannot prove completion for an invalid backend response"), } } + + fn stream<'a>( + &'a mut self, + conn: &'a mut crate::stream::ConnectionStream, + pool: &'a crate::pool::BufferPool, + backend_id: crate::types::BackendId, + ) -> Result, crate::session::response_transfer::ResponseTransferError> + { + let frame = std::mem::replace(&mut self.frame, Err(ResponseReadError::Incomplete)); + StreamingResponse::from_classification(frame, &mut self.buffer, conn, pool, backend_id) + } + + pub(crate) async fn write( + mut self, + conn: &mut crate::stream::ConnectionStream, + writer: &mut W, + pool: &crate::pool::BufferPool, + backend_id: crate::types::BackendId, + ) -> Result { + let stats = self.stream(conn, pool, backend_id)?.write(writer).await?; + stats.record(); + Ok(stats.bytes_written_u64()) + } + + pub(crate) async fn observe( + mut self, + conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, + backend_id: crate::types::BackendId, + ) -> Result< + crate::session::backend::BackendResponseComplete, + crate::session::response_transfer::ResponseTransferError, + > { + self.stream(conn, pool, backend_id)?.observe().await?; + Ok(crate::session::backend::BackendResponseComplete::response()) + } + + pub(crate) async fn capture_and_write( + mut self, + conn: &mut crate::stream::ConnectionStream, + writer: &mut W, + captured: &mut crate::pool::ChunkedResponse, + pool: &crate::pool::BufferPool, + backend_id: crate::types::BackendId, + ) -> Result<(u64, bool), crate::session::response_transfer::ResponseTransferError> { + self.stream(conn, pool, backend_id)? + .capture_and_write(writer, captured, MAX_CAPTURED_MULTILINE_RESPONSE_BYTES) + .await + } } -pub(crate) fn backend_response_read( - request: &crate::protocol::RequestContext, - buffer: &crate::pool::PooledBuffer, -) -> Result { - ResponseFrame::parse(request, buffer).map(|frame| match frame { - ResponseFrame::SingleLine { status, framed } => BackendResponseRead::SingleLine { - status, - response: framed.response, - }, - ResponseFrame::Multiline { status } => BackendResponseRead::Multiline { status }, - }) +fn strip_captured_terminator(capture: &mut crate::pool::PooledBuffer) -> anyhow::Result<()> { + let len = capture.initialized(); + anyhow::ensure!(len >= TERMINATOR.len(), "captured response is too short"); + anyhow::ensure!( + capture[len - TERMINATOR.len()..] == TERMINATOR[..], + "captured response does not end at the framed terminator" + ); + capture.truncate_initialized(len - TERMINATOR.len()); + Ok(()) } pub(crate) fn unpacked_single_line_response<'a>( @@ -1020,12 +889,14 @@ fn log_response_warnings( } } +#[cfg(test)] pub(crate) async fn capture_isolated_multiline_response( conn: &mut crate::stream::ConnectionStream, io_buffer: &mut crate::pool::PooledBuffer, capture: &mut crate::pool::PooledBuffer, ) -> anyhow::Result<()> { - IsolatedMultilineResponse { conn, io_buffer } + IsolatedMultilineResponse::begin(conn, io_buffer) + .map_err(isolated_multiline_error)? .capture_into(capture) .await } @@ -1034,7 +905,7 @@ pub(crate) async fn observe_isolated_multiline_response( conn: &mut crate::stream::ConnectionStream, io_buffer: &mut crate::pool::PooledBuffer, ) -> Result<(), FramingError> { - IsolatedMultilineResponse { conn, io_buffer } + IsolatedMultilineResponse::begin(conn, io_buffer)? .observe() .await } @@ -1045,185 +916,108 @@ pub(crate) async fn capture_isolated_multiline_response_chunked_optional( pool: &crate::pool::BufferPool, response: &mut crate::pool::ChunkedResponse, ) -> Result { - IsolatedMultilineResponse { conn, io_buffer } + IsolatedMultilineResponse::begin(conn, io_buffer)? .capture_chunked_optional(pool, response, MAX_CAPTURED_MULTILINE_RESPONSE_BYTES) .await } +/// The isolated operation rejects packed suffixes and owns its continuation. +/// A read can only advance the scanner belonging to this borrowed input window. struct IsolatedMultilineResponse<'a> { conn: &'a mut crate::stream::ConnectionStream, io_buffer: &'a mut crate::pool::PooledBuffer, + framer: MultilineFramer, + frame: ResponseWindow, } impl<'a> IsolatedMultilineResponse<'a> { - async fn capture_into(self, capture: &mut crate::pool::PooledBuffer) -> anyhow::Result<()> { + fn begin( + conn: &'a mut crate::stream::ConnectionStream, + io_buffer: &'a mut crate::pool::PooledBuffer, + ) -> Result { let mut framer = MultilineFramer::default(); - let initial_len = self.io_buffer.initialized(); - let mut continuation = match framer - .frame_initial_isolated_multiline_chunk(&self.io_buffer[..initial_len]) - .map_err(isolated_multiline_error)? - { - FramedMultilineChunk::Complete(complete) => { - complete.extend_capture_from(&self.io_buffer[..initial_len], capture); - None - } - FramedMultilineChunk::Incomplete(incomplete) => { - incomplete.extend_capture_from(&self.io_buffer[..initial_len], capture); - Some(incomplete) - } - }; + let frame = framer.frame_isolated_multiline_chunk(io_buffer)?; + Ok(Self { + conn, + io_buffer, + framer, + frame, + }) + } - while let Some(prior) = continuation.take() { - let n = self + async fn capture_into(mut self, capture: &mut crate::pool::PooledBuffer) -> anyhow::Result<()> { + loop { + match &self.frame { + ResponseWindow::Complete(chunk) => { + chunk.extend_capture_from(self.io_buffer, capture); + return Ok(()); + } + ResponseWindow::Incomplete(chunk) => { + chunk.extend_capture_from(self.io_buffer, capture); + } + } + // Preserve the underlying I/O error in the fallible capture API. + let read = self .io_buffer .read_from(self.conn) .await .context("Failed to read multiline response from backend")?; - if n == 0 { - anyhow::bail!("Backend closed connection before complete multiline response"); + if read == 0 { + return Err(isolated_multiline_error(FramingError::BackendEof)); } - continuation = match framer - .frame_next_isolated_multiline_chunk(prior, &self.io_buffer[..n]) - .map_err(isolated_multiline_error)? - { - FramedMultilineChunk::Complete(complete) => { - complete.extend_capture_from(&self.io_buffer[..n], capture); - None - } - FramedMultilineChunk::Incomplete(incomplete) => { - incomplete.extend_capture_from(&self.io_buffer[..n], capture); - Some(incomplete) - } - }; + self.frame = self + .framer + .frame_isolated_multiline_chunk(self.io_buffer) + .map_err(isolated_multiline_error)?; } - Ok(()) } - async fn observe(self) -> Result<(), FramingError> { - let mut framer = MultilineFramer::default(); - let initial_len = self.io_buffer.initialized(); - let mut continuation = - match framer.frame_initial_isolated_multiline_chunk(&self.io_buffer[..initial_len])? { - FramedMultilineChunk::Complete(_) => None, - FramedMultilineChunk::Incomplete(incomplete) => Some(incomplete), - }; + async fn read_next_chunk(&mut self) -> Result<(), FramingError> { + let read = self + .io_buffer + .read_from(self.conn) + .await + .map_err(|_| FramingError::Io)?; + if read == 0 { + return Err(FramingError::BackendEof); + } + self.frame = self.framer.frame_isolated_multiline_chunk(self.io_buffer)?; + Ok(()) + } - while let Some(prior) = continuation.take() { - let n = self - .io_buffer - .read_from(self.conn) - .await - .map_err(|_| FramingError::Io)?; - if n == 0 { - return Err(FramingError::BackendEof); + async fn observe(mut self) -> Result<(), FramingError> { + loop { + match &self.frame { + ResponseWindow::Complete(_) => return Ok(()), + ResponseWindow::Incomplete(_) => self.read_next_chunk().await?, } - continuation = - match framer.frame_next_isolated_multiline_chunk(prior, &self.io_buffer[..n])? { - FramedMultilineChunk::Complete(_) => None, - FramedMultilineChunk::Incomplete(incomplete) => Some(incomplete), - }; } - Ok(()) } async fn capture_chunked_optional( - self, + mut self, pool: &crate::pool::BufferPool, response: &mut crate::pool::ChunkedResponse, retention_limit: usize, - ) -> Result { - let mut framer = MultilineFramer::default(); - let initial_len = self.io_buffer.initialized(); - let mut continuation = - match framer.frame_initial_isolated_multiline_chunk(&self.io_buffer[..initial_len])? { - FramedMultilineChunk::Complete(complete) => { - if capture_would_exceed_limit( - response.len(), - complete.response.len(), - retention_limit, - ) { - response.clear(); - return Ok(false); - } - complete.push_isolated_buffer_to(response, pool, self.io_buffer); - None - } - FramedMultilineChunk::Incomplete(incomplete) => { - if capture_would_exceed_limit( - response.len(), - incomplete.response.len(), - retention_limit, - ) { - response.clear(); - return self - .drain_after_optional_capture_limit(&mut framer, incomplete) - .await; - } - incomplete.push_buffer_to(response, pool, self.io_buffer); - Some(incomplete) - } - }; - - while let Some(prior) = continuation.take() { - let n = self - .io_buffer - .read_from(self.conn) - .await - .map_err(|_| FramingError::Io)?; - if n == 0 { - return Err(FramingError::BackendEof); - } - continuation = - match framer.frame_next_isolated_multiline_chunk(prior, &self.io_buffer[..n])? { - FramedMultilineChunk::Complete(complete) => { - if capture_would_exceed_limit( - response.len(), - complete.response.len(), - retention_limit, - ) { - response.clear(); - return Ok(false); - } - complete.push_isolated_buffer_to(response, pool, self.io_buffer); - None - } - FramedMultilineChunk::Incomplete(incomplete) => { - if capture_would_exceed_limit( - response.len(), - incomplete.response.len(), - retention_limit, - ) { - response.clear(); - return self - .drain_after_optional_capture_limit(&mut framer, incomplete) - .await; - } - incomplete.push_buffer_to(response, pool, self.io_buffer); - Some(incomplete) - } - }; - } - Ok(true) - } - - async fn drain_after_optional_capture_limit( - self, - framer: &mut MultilineFramer, - mut continuation: IncompleteMultilineWireChunk, ) -> Result { loop { - let n = self - .io_buffer - .read_from(self.conn) - .await - .map_err(|_| FramingError::Io)?; - if n == 0 { - return Err(FramingError::BackendEof); + let chunk_len = match &self.frame { + ResponseWindow::Complete(chunk) => chunk.response.len(), + ResponseWindow::Incomplete(chunk) => chunk.response.len(), + }; + if capture_would_exceed_limit(response.len(), chunk_len, retention_limit) { + response.clear(); + self.observe().await?; + return Ok(false); } - match framer.frame_next_isolated_multiline_chunk(continuation, &self.io_buffer[..n])? { - FramedMultilineChunk::Complete(_) => return Ok(false), - FramedMultilineChunk::Incomplete(incomplete) => { - continuation = incomplete; + match &self.frame { + ResponseWindow::Complete(chunk) => { + chunk.push_isolated_buffer_to(response, pool, self.io_buffer); + return Ok(true); + } + ResponseWindow::Incomplete(chunk) => { + chunk.push_buffer_to(response, pool, self.io_buffer); + self.read_next_chunk().await?; } } } @@ -1239,484 +1033,243 @@ async fn capture_response( pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - ResponseCapture { - request, - io_buffer, - conn, - response, - pool, - backend_id, - } - .capture() - .await -} - -pub(crate) async fn write_response_with_optional_capture( - request: &crate::protocol::RequestContext, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - writer: &mut W, - response: &mut crate::pool::ChunkedResponse, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, -) -> Result<(u64, bool), crate::session::response_transfer::ResponseTransferError> -where - W: AsyncWrite + Unpin, -{ - ResponseCaptureAndWrite { - request, - io_buffer, - conn, - writer, - response, - pool, - backend_id, - retention_limit: MAX_CAPTURED_MULTILINE_RESPONSE_BYTES, - } - .capture_and_write() - .await -} - -pub(crate) async fn observe_response( - request: &crate::protocol::RequestContext, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, -) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - ResponseObserver { - request, - io_buffer, - conn, - pool, - backend_id, - } - .observe() - .await + StreamingResponse::begin(request, io_buffer, conn, pool, backend_id)? + .capture(response) + .await } -#[cfg(test)] -struct ResponseCapture<'a> { - request: &'a crate::protocol::RequestContext, +/// One response operation owns the matching scanner and exclusive I/O window. +/// No continuation or response range escapes independently of these resources. +struct StreamingResponse<'a> { io_buffer: &'a mut crate::pool::PooledBuffer, conn: &'a mut crate::stream::ConnectionStream, - response: &'a mut crate::pool::ChunkedResponse, pool: &'a crate::pool::BufferPool, backend_id: crate::types::BackendId, + framer: MultilineFramer, + frame: ResponseWindow, + shape: ResponseShape, + bytes_received: u64, } -struct ResponseObserver<'a> { - request: &'a crate::protocol::RequestContext, - io_buffer: &'a mut crate::pool::PooledBuffer, - conn: &'a mut crate::stream::ConnectionStream, - pool: &'a crate::pool::BufferPool, - backend_id: crate::types::BackendId, +enum ResponseShape { + SingleLine, + Multiline, } -struct ResponseCaptureAndWrite<'a, W> { - request: &'a crate::protocol::RequestContext, - io_buffer: &'a mut crate::pool::PooledBuffer, - conn: &'a mut crate::stream::ConnectionStream, - writer: &'a mut W, - response: &'a mut crate::pool::ChunkedResponse, - pool: &'a crate::pool::BufferPool, - backend_id: crate::types::BackendId, - retention_limit: usize, -} +impl<'a> StreamingResponse<'a> { + #[cfg(test)] + fn begin( + request: &crate::protocol::RequestContext, + io_buffer: &'a mut crate::pool::PooledBuffer, + conn: &'a mut crate::stream::ConnectionStream, + pool: &'a crate::pool::BufferPool, + backend_id: crate::types::BackendId, + ) -> Result { + let frame = ResponseFrame::parse(request, io_buffer); + Self::from_classification(frame, io_buffer, conn, pool, backend_id) + } -pub(crate) async fn write_response( - request: &crate::protocol::RequestContext, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - writer: &mut W, - pool: &crate::pool::BufferPool, - backend_id: crate::types::BackendId, -) -> Result -where - W: AsyncWrite + Unpin, -{ - let mut framer = MultilineFramer::default(); - let initial_len = io_buffer.initialized(); - let frame = ResponseFrame::parse(request, io_buffer).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io(anyhow::anyhow!( - "Failed to frame response: {err:?}" - )) - })?; + fn from_classification( + classification: Result, + io_buffer: &'a mut crate::pool::PooledBuffer, + conn: &'a mut crate::stream::ConnectionStream, + pool: &'a crate::pool::BufferPool, + backend_id: crate::types::BackendId, + ) -> Result { + let mut framer = MultilineFramer::default(); + let (shape, frame) = match classification.map_err(|err| { + crate::session::response_transfer::ResponseTransferError::Io(anyhow::anyhow!( + "Failed to frame response: {err:?}" + )) + })? { + ResponseFrame::SingleLine { framed, .. } => ( + ResponseShape::SingleLine, + ResponseWindow::Complete(CompleteResponseWindow { + response: framed.response, + next_response_input: framed.next_response_input, + }), + ), + ResponseFrame::Multiline { .. } => ( + ResponseShape::Multiline, + framer.frame_multiline_chunk(io_buffer), + ), + }; + let bytes_received = io_buffer.initialized() as u64; + Ok(Self { + io_buffer, + conn, + pool, + backend_id, + framer, + frame, + shape, + bytes_received, + }) + } - if let ResponseFrame::SingleLine { framed, .. } = &frame { - let stats = framed - .write_from(writer, io_buffer, conn, pool, initial_len) - .await?; - stats.record(); - return Ok(stats.bytes_written_u64()); - } - - let stats = match framer.frame_initial_multiline_chunk(&io_buffer[..initial_len]) { - FramedMultilineChunk::Complete(complete) => { - complete - .write_from(writer, io_buffer, conn, pool, initial_len) - .await? - } - FramedMultilineChunk::Incomplete(incomplete) => { - if let Some(io_buffer) = io_buffer.retained_append_permit() { - incomplete - .append_to_packed_prefix_and_write_with_next_backend_chunk( - writer, - &mut framer, - io_buffer, - conn, - pool, - backend_id, - ) - .await? - } else { - incomplete - .write_and_consume_response( - writer, - initial_len, - &mut framer, - io_buffer, - conn, - pool, - backend_id, - ) - .await? + /// Only newly appended bytes enter the scanner. Translation back into the + /// visible window happens while the append result still borrows that window. + async fn append_to_retained_prefix_if_writable( + &mut self, + ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { + match &self.frame { + ResponseWindow::Complete(_) => return Ok(()), + ResponseWindow::Incomplete(_) => {} + }; + let Some(permit) = self.io_buffer.retained_append_permit() else { + return Ok(()); + }; + let appended = match permit.read(self.conn).await.map_err(|error| { + crate::session::response_transfer::ResponseTransferError::Io( + anyhow::Error::from(error).context("Failed to read remaining response body"), + ) + })? { + crate::pool::AppendOutcome::Data(appended) => appended, + crate::pool::AppendOutcome::Eof(buffer) => { + let _ = buffer; + return Err( + crate::session::response_transfer::ResponseTransferError::BackendEof { + backend_id: self.backend_id, + bytes_received: self.bytes_received, + }, + ); } - } - }; - stats.record(); - Ok(stats.bytes_written_u64()) -} + }; + let origin = WindowEnd(appended.previous_len()); + let new_bytes = appended.as_new_bytes(); + self.bytes_received += new_bytes.len() as u64; + self.frame = self + .framer + .split_chunk(new_bytes, PackedPendingBytesPolicy::AllowIfStatusPrefix) + .expect("pending bytes policy cannot reject trailing bytes") + .in_window(origin, origin.0 + new_bytes.len()); + Ok(()) + } -async fn consume_remaining_multiline_response( - framer: &mut MultilineFramer, - mut prior: IncompleteMultilineWireChunk, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - backend_id: crate::types::BackendId, - mut bytes_received: u64, -) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - loop { - let n = io_buffer.read_from(conn).await.map_err(|e| { + async fn read_next_chunk( + &mut self, + ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { + let read = self.io_buffer.read_from(self.conn).await.map_err(|error| { crate::session::response_transfer::ResponseTransferError::Io( - anyhow::Error::from(e).context("Failed to read remaining response body"), + anyhow::Error::from(error).context("Failed to read remaining response body"), ) })?; - if n == 0 { + if read == 0 { return Err( crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received, + backend_id: self.backend_id, + bytes_received: self.bytes_received, }, ); } - match framer.frame_next_multiline_chunk(prior, &io_buffer[..n]) { - FramedMultilineChunk::Complete(complete) => { - complete.queue_next_response_input(&io_buffer[..n], conn)?; - return Ok(()); - } - FramedMultilineChunk::Incomplete(incomplete) => { - bytes_received += incomplete.response.len() as u64; - prior = incomplete; - } - } + self.bytes_received += read as u64; + self.frame = self.framer.frame_multiline_chunk(self.io_buffer); + Ok(()) } -} - -#[cfg(test)] -impl<'a> ResponseCapture<'a> { - async fn capture(self) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - let ResponseCapture { - request, - io_buffer, - conn, - response, - pool, - backend_id, - } = self; - let mut framer = MultilineFramer::default(); - let initial_len = io_buffer.initialized(); - let frame = ResponseFrame::parse(request, io_buffer).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io(anyhow::anyhow!( - "Failed to frame response: {err:?}" - )) - })?; - - response.clear(); - if let ResponseFrame::SingleLine { framed, .. } = &frame { - return framed.push_from_buffer(io_buffer, conn, response, pool); - } - let mut continuation = match framer.frame_initial_multiline_chunk(&io_buffer[..initial_len]) - { - FramedMultilineChunk::Complete(framed) => { - framed.push_from_buffer(io_buffer, conn, response, pool, initial_len)?; - ensure_capture_len(response.len()).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io( - isolated_multiline_error(err), - ) - })?; - None - } - FramedMultilineChunk::Incomplete(incomplete) => { - incomplete.push_buffer_to(response, pool, io_buffer); - ensure_capture_len(response.len()).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io( - isolated_multiline_error(err), - ) - })?; - Some(incomplete) + async fn drain( + mut self, + ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { + loop { + match &self.frame { + ResponseWindow::Complete(chunk) => { + return chunk.queue_next_response_input(self.io_buffer, self.conn); + } + ResponseWindow::Incomplete(_) => self.read_next_chunk().await?, } - }; + } + } - while let Some(prior) = continuation.take() { - let n = io_buffer.read_from(conn).await.map_err(|e| { - crate::session::response_transfer::ResponseTransferError::Io( - anyhow::Error::from(e).context("Failed to read remaining response body"), - ) - })?; - if n == 0 { - return Err( - crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received: response.len() as u64, - }, - ); + async fn observe(self) -> Result<(), crate::session::response_transfer::ResponseTransferError> { + match &self.frame { + ResponseWindow::Complete(chunk) => { + let total_len = self.io_buffer.initialized(); + chunk.observe_from_buffer(self.io_buffer, self.conn, self.pool, total_len) } - continuation = match framer.frame_next_multiline_chunk(prior, &io_buffer[..n]) { - FramedMultilineChunk::Complete(framed) => { - framed.push_from_buffer(io_buffer, conn, response, pool, n)?; - ensure_capture_len(response.len()).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io( - isolated_multiline_error(err), - ) - })?; - None - } - FramedMultilineChunk::Incomplete(incomplete) => { - incomplete.push_buffer_to(response, pool, io_buffer); - ensure_capture_len(response.len()).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io( - isolated_multiline_error(err), - ) - })?; - Some(incomplete) - } - }; + ResponseWindow::Incomplete(_) => self.drain().await, } - Ok(()) } -} -impl ResponseCaptureAndWrite<'_, W> -where - W: AsyncWrite + Unpin, -{ - async fn flush_unretained_prefix( + async fn write( + mut self, writer: &mut W, - response: &mut crate::pool::ChunkedResponse, - ) -> Result { - let mut stats = ResponseWriteStats::default(); - stats.add_buffered_response(response); - response - .write_all_to_recording(writer) - .await - .map_err(crate::session::response_transfer::ResponseTransferError::ClientDisconnect)?; - response.clear(); - Ok(stats) + ) -> Result { + self.append_to_retained_prefix_if_writable().await?; + self.write_current_and_remaining_chunks(writer).await } - async fn flush_unretained_prefix_before_complete( + async fn write_current_and_remaining_chunks( + mut self, writer: &mut W, - response: &mut crate::pool::ChunkedResponse, - framed: &CompleteMultilineWireChunk, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - total_len: usize, ) -> Result { - match Self::flush_unretained_prefix(writer, response).await { - Ok(stats) => Ok(stats), - Err(crate::session::response_transfer::ResponseTransferError::ClientDisconnect(e)) => { - framed.observe_from_buffer(io_buffer, conn, pool, total_len)?; - Err(crate::session::response_transfer::ResponseTransferError::ClientDisconnect(e)) + let mut stats = ResponseWriteStats::default(); + loop { + match &self.frame { + ResponseWindow::Complete(chunk) => { + let total_len = self.io_buffer.initialized(); + stats += chunk + .write_from(writer, self.io_buffer, self.conn, self.pool, total_len) + .await?; + return Ok(stats); + } + ResponseWindow::Incomplete(chunk) => { + let bytes = &self.io_buffer[chunk.response.clone()]; + if let Err(error) = writer.write_all(bytes).await { + self.drain().await?; + return Err(crate::session::response_transfer::ResponseTransferError::ClientDisconnect(error)); + } + stats.add_chunk(bytes.len()); + self.read_next_chunk().await?; + } } - Err(e) => Err(e), } } - #[allow(clippy::too_many_arguments)] - async fn flush_unretained_prefix_or_drain( + async fn capture_and_write( + mut self, writer: &mut W, response: &mut crate::pool::ChunkedResponse, - framer: &mut MultilineFramer, - incomplete: IncompleteMultilineWireChunk, - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - backend_id: crate::types::BackendId, - ) -> Result< - (ResponseWriteStats, IncompleteMultilineWireChunk), - crate::session::response_transfer::ResponseTransferError, - > { - let bytes_received = response.len() as u64; - match Self::flush_unretained_prefix(writer, response).await { - Ok(stats) => Ok((stats, incomplete)), - Err(crate::session::response_transfer::ResponseTransferError::ClientDisconnect(e)) => { - consume_remaining_multiline_response( - framer, - incomplete, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await?; - Err(crate::session::response_transfer::ResponseTransferError::ClientDisconnect(e)) - } - Err(e) => Err(e), - } - } - - #[inline] - fn would_exceed_retention( - response: &crate::pool::ChunkedResponse, - next_len: usize, retention_limit: usize, - ) -> bool { - ensure_capture_len_with_limit(response.len().saturating_add(next_len), retention_limit) - .is_err() - } - - async fn capture_and_write( - self, ) -> Result<(u64, bool), crate::session::response_transfer::ResponseTransferError> { - let ResponseCaptureAndWrite { - request, - io_buffer, - conn, - writer, - response, - pool, - backend_id, - retention_limit, - } = self; - let mut framer = MultilineFramer::default(); - let initial_len = io_buffer.initialized(); - let frame = ResponseFrame::parse(request, io_buffer).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io(anyhow::anyhow!( - "Failed to frame response: {err:?}" - )) - })?; - response.clear(); - if let ResponseFrame::SingleLine { framed, .. } = &frame { - framed.push_from_buffer(io_buffer, conn, response, pool)?; - let mut stats = ResponseWriteStats::default(); - stats.add_buffered_response(response); - response.write_all_to_recording(writer).await.map_err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect, - )?; - stats.record(); - return Ok((stats.bytes_written_u64(), true)); - } - - let mut continuation = match framer.frame_initial_multiline_chunk(&io_buffer[..initial_len]) - { - FramedMultilineChunk::Complete(framed) => { - if Self::would_exceed_retention(response, framed.response.len(), retention_limit) { - let mut stats = Self::flush_unretained_prefix_before_complete( - writer, - response, - &framed, - io_buffer, - conn, - pool, - initial_len, - ) - .await?; - stats += framed - .write_from(writer, io_buffer, conn, pool, initial_len) - .await?; - stats.record(); - return Ok((stats.bytes_written_u64(), false)); + loop { + let chunk_len = match &self.frame { + ResponseWindow::Complete(chunk) => chunk.response.len(), + ResponseWindow::Incomplete(chunk) => chunk.response.len(), + }; + let exceeds_limit = match self.shape { + ResponseShape::SingleLine => false, + ResponseShape::Multiline => { + capture_would_exceed_limit(response.len(), chunk_len, retention_limit) } - framed.push_from_buffer(io_buffer, conn, response, pool, initial_len)?; + }; + if exceeds_limit { let mut stats = ResponseWriteStats::default(); stats.add_buffered_response(response); - response.write_all_to_recording(writer).await.map_err( - crate::session::response_transfer::ResponseTransferError::ClientDisconnect, - )?; - stats.record(); - return Ok((stats.bytes_written_u64(), true)); - } - FramedMultilineChunk::Incomplete(incomplete) => { - if Self::would_exceed_retention( - response, - incomplete.response.len(), - retention_limit, - ) { - let (stats, incomplete) = Self::flush_unretained_prefix_or_drain( - writer, - response, - &mut framer, - incomplete, - io_buffer, - conn, - backend_id, - ) - .await?; - let bytes_received = stats.bytes_written_u64(); - let stats = incomplete - .write_current_and_after( - writer, - initial_len, - &mut framer, - io_buffer, - conn, - pool, - backend_id, - stats, - bytes_received, - ) - .await?; - stats.record(); - return Ok((stats.bytes_written_u64(), false)); + if let Err(error) = response.write_all_to_recording(writer).await { + self.observe().await?; + return Err( + crate::session::response_transfer::ResponseTransferError::ClientDisconnect( + error, + ), + ); } - incomplete.push_buffer_to(response, pool, io_buffer); - Some(incomplete) + response.clear(); + stats += self.write_current_and_remaining_chunks(writer).await?; + stats.record(); + return Ok((stats.bytes_written_u64(), false)); } - }; - while let Some(prior) = continuation.take() { - let n = io_buffer.read_from(conn).await.map_err(|e| { - crate::session::response_transfer::ResponseTransferError::Io( - anyhow::Error::from(e).context("Failed to read remaining response body"), - ) - })?; - if n == 0 { - return Err( - crate::session::response_transfer::ResponseTransferError::BackendEof { - backend_id, - bytes_received: response.len() as u64, - }, - ); - } - continuation = match framer.frame_next_multiline_chunk(prior, &io_buffer[..n]) { - FramedMultilineChunk::Complete(framed) => { - if Self::would_exceed_retention( + match &self.frame { + ResponseWindow::Complete(chunk) => { + let total_len = self.io_buffer.initialized(); + chunk.push_from_buffer( + self.io_buffer, + self.conn, response, - framed.response.len(), - retention_limit, - ) { - let mut stats = Self::flush_unretained_prefix_before_complete( - writer, response, &framed, io_buffer, conn, pool, n, - ) - .await?; - stats += framed.write_from(writer, io_buffer, conn, pool, n).await?; - stats.record(); - return Ok((stats.bytes_written_u64(), false)); - } - framed.push_from_buffer(io_buffer, conn, response, pool, n)?; + self.pool, + total_len, + )?; let mut stats = ResponseWriteStats::default(); stats.add_buffered_response(response); response.write_all_to_recording(writer).await.map_err( @@ -1725,121 +1278,85 @@ where stats.record(); return Ok((stats.bytes_written_u64(), true)); } - FramedMultilineChunk::Incomplete(incomplete) => { - if Self::would_exceed_retention( - response, - incomplete.response.len(), - retention_limit, - ) { - let (stats, incomplete) = Self::flush_unretained_prefix_or_drain( - writer, - response, - &mut framer, - incomplete, - io_buffer, - conn, - backend_id, - ) - .await?; - let bytes_received = stats.bytes_written_u64(); - let stats = incomplete - .write_current_and_after( - writer, - n, - &mut framer, - io_buffer, - conn, - pool, - backend_id, - stats, - bytes_received, - ) - .await?; - stats.record(); - return Ok((stats.bytes_written_u64(), false)); - } - incomplete.push_buffer_to(response, pool, io_buffer); - Some(incomplete) + ResponseWindow::Incomplete(chunk) => { + chunk.push_buffer_to(response, self.pool, self.io_buffer); + self.read_next_chunk().await?; } - }; + } } - - let mut stats = ResponseWriteStats::default(); - stats.add_buffered_response(response); - stats.record(); - Ok((stats.bytes_written_u64(), true)) } -} - -impl<'a> ResponseObserver<'a> { - async fn observe(self) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - let ResponseObserver { - request, - io_buffer, - conn, - pool, - backend_id, - } = self; - let mut framer = MultilineFramer::default(); - let initial_len = io_buffer.initialized(); - let frame = ResponseFrame::parse(request, io_buffer).map_err(|err| { - crate::session::response_transfer::ResponseTransferError::Io(anyhow::anyhow!( - "Failed to frame response: {err:?}" - )) - })?; - - if let ResponseFrame::SingleLine { framed, .. } = &frame { - return framed.observe_from_buffer(io_buffer, conn, pool); - } - match framer.frame_initial_multiline_chunk(&io_buffer[..initial_len]) { - FramedMultilineChunk::Complete(framed) => { - framed.observe_from_buffer(io_buffer, conn, pool, initial_len) + #[cfg(test)] + async fn capture( + mut self, + response: &mut crate::pool::ChunkedResponse, + ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { + response.clear(); + loop { + let complete = match &self.frame { + ResponseWindow::Complete(chunk) => { + let total_len = self.io_buffer.initialized(); + chunk.push_from_buffer( + self.io_buffer, + self.conn, + response, + self.pool, + total_len, + )?; + true + } + ResponseWindow::Incomplete(chunk) => { + chunk.push_buffer_to(response, self.pool, self.io_buffer); + false + } + }; + match self.shape { + ResponseShape::SingleLine => {} + ResponseShape::Multiline => { + ensure_capture_len(response.len()).map_err(|error| { + crate::session::response_transfer::ResponseTransferError::Io( + isolated_multiline_error(error), + ) + })? + } } - FramedMultilineChunk::Incomplete(incomplete) => { - let bytes_received = incomplete.response.len() as u64; - consume_remaining_multiline_response( - &mut framer, - incomplete, - io_buffer, - conn, - backend_id, - bytes_received, - ) - .await + if complete { + return Ok(()); } + self.read_next_chunk().await?; } } } -impl MultilineFramer { - fn frame_initial_multiline_chunk(&mut self, chunk: &[u8]) -> FramedMultilineChunk { - self.split_chunk(chunk, PackedPendingBytesPolicy::AllowIfStatusPrefix) - .expect("pending bytes policy cannot reject trailing bytes") - } +#[cfg(test)] +async fn write_response( + request: &crate::protocol::RequestContext, + buffer: crate::pool::PooledBuffer, + conn: &mut crate::stream::ConnectionStream, + writer: &mut W, + pool: &crate::pool::BufferPool, + backend_id: crate::types::BackendId, +) -> Result { + ClassifiedResponse::read(conn, request, buffer) + .await + .map_err(crate::session::response_transfer::ResponseTransferError::Io)? + .write(conn, writer, pool, backend_id) + .await +} - fn frame_next_multiline_chunk( - &mut self, - _continuation: IncompleteMultilineWireChunk, - chunk: &[u8], - ) -> FramedMultilineChunk { +impl MultilineFramer { + fn frame_multiline_chunk(&mut self, chunk: &[u8]) -> ResponseWindow { self.split_chunk(chunk, PackedPendingBytesPolicy::AllowIfStatusPrefix) .expect("pending bytes policy cannot reject trailing bytes") + .in_window(WindowEnd(0), chunk.len()) } - fn frame_initial_isolated_multiline_chunk( - &mut self, - chunk: &[u8], - ) -> Result { - self.split_chunk(chunk, PackedPendingBytesPolicy::Reject) - } - - fn frame_next_isolated_multiline_chunk( + fn frame_isolated_multiline_chunk( &mut self, - _continuation: IncompleteMultilineWireChunk, chunk: &[u8], - ) -> Result { + ) -> Result { self.split_chunk(chunk, PackedPendingBytesPolicy::Reject) + .map(|progress| progress.in_window(WindowEnd(0), chunk.len())) } /// Update tail with the last bytes from a chunk @@ -1875,13 +1392,10 @@ impl MultilineFramer { &mut self, chunk: &[u8], suffix_policy: PackedPendingBytesPolicy, - ) -> Result { + ) -> Result { for end in self.terminator_ends(chunk) { if end == chunk.len() { - return Ok(FramedMultilineChunk::Complete(CompleteMultilineWireChunk { - response: 0..end, - next_response_input: end..end, - })); + return Ok(ChunkProgress::Complete(ChunkConsumed(end))); } match suffix_policy { @@ -1891,21 +1405,14 @@ impl MultilineFramer { PackedPendingBytesPolicy::AllowIfStatusPrefix if plausible_status_prefix(&chunk[end..]) => { - return Ok(FramedMultilineChunk::Complete(CompleteMultilineWireChunk { - response: 0..end, - next_response_input: end..chunk.len(), - })); + return Ok(ChunkProgress::Complete(ChunkConsumed(end))); } PackedPendingBytesPolicy::AllowIfStatusPrefix => {} } } self.update(chunk); - Ok(FramedMultilineChunk::Incomplete( - IncompleteMultilineWireChunk { - response: 0..chunk.len(), - }, - )) + Ok(ChunkProgress::Incomplete) } /// Find spanning terminator offset in chunk @@ -2042,12 +1549,10 @@ impl PendingRequestFrame { match framer .split_chunk(&chunk[end..], PackedPendingBytesPolicy::AllowIfStatusPrefix) { - Ok(FramedMultilineChunk::Complete(complete)) => { - Some(FramedResponseForRequest { - response: offset..end + complete.response.end, - }) - } - Ok(FramedMultilineChunk::Incomplete(_)) => { + Ok(ChunkProgress::Complete(consumed)) => Some(FramedResponseForRequest { + response: offset..WindowEnd(end).after_chunk(consumed).0, + }), + Ok(ChunkProgress::Incomplete) => { self.state = PendingRequestFrameState::ReadingMultiline { framer }; None } @@ -2061,12 +1566,10 @@ impl PendingRequestFrame { &chunk[offset..], PackedPendingBytesPolicy::AllowIfStatusPrefix, ) { - Ok(FramedMultilineChunk::Complete(complete)) => { - Some(FramedResponseForRequest { - response: offset..offset + complete.response.end, - }) - } - Ok(FramedMultilineChunk::Incomplete(_)) => None, + Ok(ChunkProgress::Complete(consumed)) => Some(FramedResponseForRequest { + response: offset..WindowEnd(offset).after_chunk(consumed).0, + }), + Ok(ChunkProgress::Incomplete) => None, Err(_) => Some(FramedResponseForRequest { response: offset..chunk.len(), }), @@ -2131,8 +1634,8 @@ fn complete_multiline_payload_split(payload: &[u8]) -> Option { - let response_end = complete.response.end; + Ok(ChunkProgress::Complete(consumed)) => { + let response_end = consumed.0; let body_end = response_end.checked_sub(TERMINATOR.len())?; let terminator_start = response_end.checked_sub(3)?; Some(CompleteMultilinePayloadSplit::new( @@ -2140,9 +1643,7 @@ fn complete_multiline_payload_split(payload: &[u8]) -> Option None, + Ok(ChunkProgress::Incomplete) | Err(_) => None, } } @@ -2201,10 +1702,10 @@ pub fn benchmark_incremental_multiline_frame(response: &[u8], chunk_size: usize) let chunk_size = chunk_size.max(1); for chunk in response.chunks(chunk_size) { - if let Ok(FramedMultilineChunk::Complete(complete)) = + if let Ok(ChunkProgress::Complete(consumed)) = framer.split_chunk(chunk, PackedPendingBytesPolicy::AllowIfStatusPrefix) { - return complete.response.end; + return consumed.0; } } @@ -2221,11 +1722,11 @@ pub fn benchmark_stateless_multiline_frame(response: &[u8], chunk_size: usize) - for chunk in response.chunks(chunk_size) { received += chunk.len(); let mut framer = MultilineFramer::default(); - if let Ok(FramedMultilineChunk::Complete(complete)) = framer.split_chunk( + if let Ok(ChunkProgress::Complete(consumed)) = framer.split_chunk( &response[..received], PackedPendingBytesPolicy::AllowIfStatusPrefix, ) { - return complete.response.end; + return consumed.0; } } @@ -2397,6 +1898,82 @@ mod tests { crate::pool::BufferPool::new(BufferSize::try_new(65536).unwrap(), 2) } + #[tokio::test] + async fn isolated_single_line_reply_rejects_a_packed_next_response() { + let pool = make_pool(); + let request = crate::protocol::RequestContext::from_verb_args(b"STAT", b""); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(b"223 0 \r\n223 0 \r\n"); + let response = ClassifiedResponse::read(&mut tokio::io::empty(), &request, buffer) + .await + .unwrap(); + assert_eq!( + response.single_line_bytes(), + Some(b"223 0 \r\n".as_slice()) + ); + assert!( + response.completion_proof().is_err(), + "an isolated request must not authorize reuse with a following reply in its buffer" + ); + let mut conn = loopback_connection_stream().await; + assert!(response.capture_isolated(&mut conn, &pool).await.is_err()); + } + + #[tokio::test] + async fn classified_forwarding_preserves_shape_and_suffix_at_every_split() { + let cases: &[(&[u8], &[u8])] = &[ + (b"STAT \r\n", b"223 0 \r\n"), + (b"BODY \r\n", b"430 not found\r\n"), + (b"BODY \r\n", b"222 0 \r\n.\r\n"), + ( + b"BODY \r\n", + b"222 0 \r\n..dot\r\nbody\r\n.\r\n", + ), + ( + b"HEAD \r\n", + b"221 0 \r\nSubject: folded\r\n continuation\r\n.\r\n", + ), + (b"GROUP alt.test\r\n", b"211 1 1 1 alt.test\r\n"), + ( + b"LISTGROUP alt.test\r\n", + b"211 1 1 1 alt.test\r\n1\r\n.\r\n", + ), + ]; + let next = b"223 0 \r\n"; + for &(command, wire) in cases { + let request = crate::protocol::RequestContext::parse(command).unwrap(); + for split in 0..=wire.len() { + let pool = make_pool(); + let mut buffer = pool.acquire(); + buffer.copy_from_slice(&wire[..split]); + let tail = [wire[split..].as_ref(), next.as_slice()].concat(); + let mut conn = mock_backend_conn(vec![tail]).await; + let response = ClassifiedResponse::read(&mut conn, &request, buffer) + .await + .unwrap(); + let mut writer = RecordingWriter::default(); + let bytes = response + .write( + &mut conn, + &mut writer, + &pool, + crate::types::BackendId::from_index(0), + ) + .await + .unwrap(); + assert_eq!( + bytes, + wire.len() as u64, + "command={command:?}, split={split}" + ); + assert_eq!(writer.bytes, wire, "command={command:?}, split={split}"); + let mut following = [0; 14]; + conn.read_exact(&mut following).await.unwrap(); + assert_eq!(&following, next, "command={command:?}, split={split}"); + } + } + } + async fn loopback_connection_stream() -> crate::stream::ConnectionStream { let listener = TcpListener::bind("127.0.0.1:0") .await @@ -2449,19 +2026,15 @@ mod tests { let mut captured = crate::pool::ChunkedResponse::default(); let mut writer = RecordingWriter::default(); - let (bytes_written, retained) = ResponseCaptureAndWrite { - request: &request, - io_buffer: &mut io_buffer, - conn: &mut conn, - writer: &mut writer, - response: &mut captured, - pool: &pool, - backend_id: crate::types::BackendId::from_index(1), - retention_limit: first.len() - 1, - } - .capture_and_write() - .await - .expect("oversized retained response should stream through"); + let mut classified = ClassifiedResponse::read(&mut conn, &request, io_buffer) + .await + .unwrap(); + let (bytes_written, retained) = classified + .stream(&mut conn, &pool, crate::types::BackendId::from_index(1)) + .unwrap() + .capture_and_write(&mut writer, &mut captured, first.len() - 1) + .await + .expect("oversized retained response should stream through"); assert_eq!(bytes_written as usize, expected.len()); assert!(!retained); @@ -2481,19 +2054,15 @@ mod tests { let mut captured = crate::pool::ChunkedResponse::default(); let mut writer = RecordingWriter::default(); - let (bytes_written, retained) = ResponseCaptureAndWrite { - request: &request, - io_buffer: &mut io_buffer, - conn: &mut conn, - writer: &mut writer, - response: &mut captured, - pool: &pool, - backend_id: crate::types::BackendId::from_index(1), - retention_limit: response.len() - 1, - } - .capture_and_write() - .await - .expect("complete oversized retained response should still write"); + let mut classified = ClassifiedResponse::read(&mut conn, &request, io_buffer) + .await + .unwrap(); + let (bytes_written, retained) = classified + .stream(&mut conn, &pool, crate::types::BackendId::from_index(1)) + .unwrap() + .capture_and_write(&mut writer, &mut captured, response.len() - 1) + .await + .expect("complete oversized retained response should still write"); assert_eq!(bytes_written as usize, response.len()); assert!(!retained); @@ -2511,13 +2080,11 @@ mod tests { io_buffer.copy_from_slice(first); let mut captured = crate::pool::ChunkedResponse::default(); - let retained = IsolatedMultilineResponse { - conn: &mut conn, - io_buffer: &mut io_buffer, - } - .capture_chunked_optional(&pool, &mut captured, first.len() - 1) - .await - .expect("oversized isolated response should drain cleanly"); + let retained = IsolatedMultilineResponse::begin(&mut conn, &mut io_buffer) + .expect("valid initial response") + .capture_chunked_optional(&pool, &mut captured, first.len() - 1) + .await + .expect("oversized isolated response should drain cleanly"); assert!(!retained); assert!(captured.is_empty()); @@ -2533,13 +2100,11 @@ mod tests { io_buffer.copy_from_slice(response); let mut captured = crate::pool::ChunkedResponse::default(); - let retained = IsolatedMultilineResponse { - conn: &mut conn, - io_buffer: &mut io_buffer, - } - .capture_chunked_optional(&pool, &mut captured, response.len()) - .await - .expect("boundary-sized isolated response should capture cleanly"); + let retained = IsolatedMultilineResponse::begin(&mut conn, &mut io_buffer) + .expect("valid initial response") + .capture_chunked_optional(&pool, &mut captured, response.len()) + .await + .expect("boundary-sized isolated response should capture cleanly"); assert!(retained); assert_eq!(captured.to_vec(), response); @@ -2555,13 +2120,11 @@ mod tests { io_buffer.copy_from_slice(response); let mut captured = crate::pool::ChunkedResponse::default(); - let retained = IsolatedMultilineResponse { - conn: &mut conn, - io_buffer: &mut io_buffer, - } - .capture_chunked_optional(&pool, &mut captured, response.len() - 1) - .await - .expect("complete oversized isolated response should drain cleanly"); + let retained = IsolatedMultilineResponse::begin(&mut conn, &mut io_buffer) + .expect("valid initial response") + .capture_chunked_optional(&pool, &mut captured, response.len() - 1) + .await + .expect("complete oversized isolated response should drain cleanly"); assert!(!retained); assert!(captured.is_empty()); @@ -2873,7 +2436,7 @@ mod tests { crate::pool::buffer::reset_hot_path_allocation_metrics(); let chunk = b"220 article\r\nbody\r\n.\r\n223 0 \r\n"; let response_len = b"220 article\r\nbody\r\n.\r\n".len(); - let framed = CompleteMultilineWireChunk { + let framed = CompleteResponseWindow { response: 0..response_len, next_response_input: response_len..chunk.len(), }; @@ -2900,22 +2463,26 @@ mod tests { crate::pool::buffer::reset_hot_path_allocation_metrics(); let chunk = b"223 0 \r\n223 0 \r\n"; let response_len = b"223 0 \r\n".len(); - let framed = FramedSingleLineChunk { - response: 0..response_len, - next_response_input: response_len..chunk.len(), - }; + let request = + crate::protocol::RequestContext::parse(b"STAT \r\n").expect("valid request"); let pool = make_pool(); let mut io_buffer = pool.acquire(); io_buffer.copy_from_slice(chunk); let mut conn = loopback_connection_stream().await; let mut writer = Vec::new(); - let written = framed - .write_from(&mut writer, &mut io_buffer, &mut conn, &pool, chunk.len()) - .await - .expect("single-line response should write"); + let written = write_response( + &request, + io_buffer, + &mut conn, + &mut writer, + &pool, + crate::types::BackendId::from_index(1), + ) + .await + .expect("single-line response should write"); - assert_eq!(written.bytes_written_u64(), response_len as u64); + assert_eq!(written, response_len as u64); assert_eq!(writer, &chunk[..response_len]); assert_eq!(conn.pending_bytes_len(), b"223 0 \r\n".len()); let metrics = crate::pool::buffer::hot_path_allocation_metrics_snapshot(); @@ -2926,7 +2493,7 @@ mod tests { async fn complete_multiline_write_client_error_preserves_packed_suffix() { let chunk = b"220 article\r\nbody\r\n.\r\n223 0 \r\n"; let response_len = b"220 article\r\nbody\r\n.\r\n".len(); - let framed = CompleteMultilineWireChunk { + let framed = CompleteResponseWindow { response: 0..response_len, next_response_input: response_len..chunk.len(), }; @@ -2950,20 +2517,23 @@ mod tests { #[tokio::test] async fn single_line_write_client_error_preserves_packed_suffix() { let chunk = b"223 0 \r\n223 0 \r\n"; - let response_len = b"223 0 \r\n".len(); - let framed = FramedSingleLineChunk { - response: 0..response_len, - next_response_input: response_len..chunk.len(), - }; + let request = + crate::protocol::RequestContext::parse(b"STAT \r\n").expect("valid request"); let pool = make_pool(); let mut io_buffer = pool.acquire(); io_buffer.copy_from_slice(chunk); let mut conn = loopback_connection_stream().await; let mut writer = FailingWriter; - let err = framed - .write_from(&mut writer, &mut io_buffer, &mut conn, &pool, chunk.len()) - .await; + let err = write_response( + &request, + io_buffer, + &mut conn, + &mut writer, + &pool, + crate::types::BackendId::from_index(1), + ) + .await; assert!(matches!( err, @@ -2983,13 +2553,13 @@ mod tests { let mut continuation = Vec::from(&first_response[split..]); continuation.extend_from_slice(next_response); let mut framer = MultilineFramer::default(); - let prior = match framer.frame_initial_multiline_chunk(initial) { - FramedMultilineChunk::Incomplete(prior) => prior, - FramedMultilineChunk::Complete(_) => panic!("split={split} unexpectedly complete"), + match framer.frame_multiline_chunk(initial) { + ResponseWindow::Incomplete(_) => {} + ResponseWindow::Complete(_) => panic!("split={split} unexpectedly complete"), }; - let complete = match framer.frame_next_multiline_chunk(prior, &continuation) { - FramedMultilineChunk::Complete(complete) => complete, - FramedMultilineChunk::Incomplete(_) => { + let complete = match framer.frame_multiline_chunk(&continuation) { + ResponseWindow::Complete(complete) => complete, + ResponseWindow::Incomplete(_) => { panic!("split={split} did not complete on continuation") } }; @@ -3038,7 +2608,7 @@ mod tests { let written = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3052,8 +2622,8 @@ mod tests { assert_eq!(conn.pending_bytes_len(), next_response.len()); assert_eq!( pool.available_buffers(), - 0, - "packed suffix should hold the original pooled read buffer" + 1, + "the consumed response returns its replacement scratch buffer; the packed suffix retains the original" ); let mut pending = vec![0; next_response.len()]; @@ -3061,8 +2631,8 @@ mod tests { assert_eq!(pending, next_response); assert_eq!( pool.available_buffers(), - 1, - "draining pooled pending input should return the original buffer" + 2, + "draining pooled pending input also returns the original buffer" ); } @@ -3084,7 +2654,7 @@ mod tests { write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3092,7 +2662,6 @@ mod tests { ) .await .expect("complete multiline response should write"); - drop(io_buffer); let next_buffer = take_queued_input_or_acquire_empty(&mut conn, &pool); @@ -3121,7 +2690,7 @@ mod tests { write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3129,15 +2698,14 @@ mod tests { ) .await .expect("first response should write"); - drop(io_buffer); - let mut second_buffer = take_queued_input_or_acquire_empty(&mut conn, &pool); + let second_buffer = take_queued_input_or_acquire_empty(&mut conn, &pool); assert_eq!(second_buffer.allocation_ptr(), original_allocation); assert_eq!(second_buffer.as_ref(), &chunk[first_response.len()..]); writer.clear(); write_response( &request, - &mut second_buffer, + second_buffer, &mut conn, &mut writer, &pool, @@ -3146,7 +2714,6 @@ mod tests { .await .expect("second response should write"); assert_eq!(writer, second_response); - drop(second_buffer); let third_buffer = take_queued_input_or_acquire_empty(&mut conn, &pool); assert_eq!(third_buffer.allocation_ptr(), original_allocation); @@ -3176,7 +2743,7 @@ mod tests { let written = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3217,7 +2784,7 @@ mod tests { let error = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3249,7 +2816,7 @@ mod tests { write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3280,7 +2847,7 @@ mod tests { let written = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3294,8 +2861,8 @@ mod tests { assert_eq!(conn.pending_bytes_len(), next_response.len()); assert_eq!( pool.available_buffers(), - 0, - "packed suffix should hold the original pooled read buffer" + 1, + "the consumed response returns its replacement scratch buffer; the packed suffix retains the original" ); let mut pending = vec![0; next_response.len()]; @@ -3303,8 +2870,8 @@ mod tests { assert_eq!(pending, next_response); assert_eq!( pool.available_buffers(), - 1, - "draining pooled pending input should return the original buffer" + 2, + "draining pooled pending input also returns the original buffer" ); } @@ -3326,7 +2893,7 @@ mod tests { let written = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3369,7 +2936,7 @@ mod tests { let written = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3409,7 +2976,7 @@ mod tests { let err = write_response( &request, - &mut io_buffer, + io_buffer, &mut conn, &mut writer, &pool, @@ -3655,22 +3222,6 @@ mod tests { ); } - #[test] - fn framed_write_paths_use_shared_suffix_preservation_chokepoint() { - let source_path = std::path::Path::new(env!("CARGO_MANIFEST_DIR")) - .join("src") - .join("session") - .join("multiline_framing.rs"); - let source = std::fs::read_to_string(source_path).expect("read multiline_framing source"); - let delegations = source - .match_indices("write_response_chunk_preserving_suffix_on_error(") - .count(); - assert!( - delegations >= 3, - "single-line and complete-multiline write paths must delegate through the shared suffix-preservation helper" - ); - } - #[test] fn update_preserves_rolling_tail() { let mut framer = MultilineFramer::default(); @@ -3747,11 +3298,9 @@ mod tests { assert_eq!( split, - Ok(FramedMultilineChunk::Complete(CompleteMultilineWireChunk { - response: 0..b"220 article\r\nbody\r\n.\r\n".len(), - next_response_input: b"220 article\r\nbody\r\n.\r\n".len() - ..b"220 article\r\nbody\r\n.\r\n".len(), - })) + Ok(ChunkProgress::Complete(ChunkConsumed( + b"220 article\r\nbody\r\n.\r\n".len() + ))) ); } @@ -3764,10 +3313,9 @@ mod tests { assert_eq!( split, - Ok(FramedMultilineChunk::Complete(CompleteMultilineWireChunk { - response: 0..b"220 article\r\nbody\r\n.\r\n".len(), - next_response_input: b"220 article\r\nbody\r\n.\r\n".len()..chunk.len(), - })) + Ok(ChunkProgress::Complete(ChunkConsumed( + b"220 article\r\nbody\r\n.\r\n".len() + ))) ); } @@ -3790,10 +3338,7 @@ mod tests { assert_eq!( split, - Ok(FramedMultilineChunk::Complete(CompleteMultilineWireChunk { - response: 0..chunk.len(), - next_response_input: chunk.len()..chunk.len(), - })) + Ok(ChunkProgress::Complete(ChunkConsumed(chunk.len()))) ); } @@ -3806,10 +3351,7 @@ mod tests { assert_eq!( split, - Ok(FramedMultilineChunk::Complete(CompleteMultilineWireChunk { - response: 0..b".\r\n".len(), - next_response_input: b".\r\n".len()..b".\r\n".len(), - })) + Ok(ChunkProgress::Complete(ChunkConsumed(b".\r\n".len()))) ); } @@ -3827,18 +3369,16 @@ mod tests { ) .expect("first push should not reject a valid response"); let actual = match first { - FramedMultilineChunk::Complete(complete) => Some(complete.response.end), - FramedMultilineChunk::Incomplete(_) => framer + ChunkProgress::Complete(consumed) => Some(consumed.0), + ChunkProgress::Incomplete => framer .split_chunk( &response[split..], PackedPendingBytesPolicy::AllowIfStatusPrefix, ) .ok() .and_then(|result| match result { - FramedMultilineChunk::Complete(complete) => { - Some(split + complete.response.end) - } - FramedMultilineChunk::Incomplete(_) => None, + ChunkProgress::Complete(consumed) => Some(split + consumed.0), + ChunkProgress::Incomplete => None, }), }; @@ -3848,10 +3388,10 @@ mod tests { let rescanned_end = stateless .split_chunk(response, PackedPendingBytesPolicy::AllowIfStatusPrefix) .expect("rescan should accept a valid response"); - let FramedMultilineChunk::Complete(rescanned) = rescanned_end else { + let ChunkProgress::Complete(rescanned) = rescanned_end else { panic!("rescan did not complete for split={split}"); }; - assert_eq!(actual, Some(rescanned.response.end), "split={split}"); + assert_eq!(actual, Some(rescanned.0), "split={split}"); } } @@ -3964,3 +3504,46 @@ mod tests { }); } } + +#[cfg(response_contract)] +#[allow(dead_code)] +mod contracts { + use super::*; + + fn chunk_coordinate() { + let origin = WindowEnd(12); + let consumed = ChunkConsumed(3); + #[cfg(response_contract = "chunk_coordinate")] + let consumed = WindowEnd(consumed.0); + std::hint::black_box(origin.after_chunk(consumed)); + } + + async fn classified_buffer_reuse( + request: &crate::protocol::RequestContext, + buffer: crate::pool::PooledBuffer, + reader: &mut (impl tokio::io::AsyncRead + Unpin), + ) { + let response = ClassifiedResponse::read(reader, request, buffer) + .await + .expect("classified response"); + #[cfg(response_contract = "classified_buffer_reuse")] + let _replaced = buffer; + std::hint::black_box(response.single_line_bytes()); + } + + async fn consuming_response( + request: &crate::protocol::RequestContext, + buffer: crate::pool::PooledBuffer, + conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, + writer: &mut (impl AsyncWrite + Unpin), + ) { + let response = ClassifiedResponse::read(conn, request, buffer) + .await + .expect("classified response"); + let transfer = response.write(conn, writer, pool, crate::types::BackendId::from_index(0)); + #[cfg(response_contract = "response_twice")] + let _conflicting = response.status_code(); + std::hint::black_box(transfer); + } +} diff --git a/src/session/precheck.rs b/src/session/precheck.rs index 7ca01679..a2ffcf34 100644 --- a/src/session/precheck.rs +++ b/src/session/precheck.rs @@ -217,14 +217,14 @@ async fn execute_backend_query( return Ok(QueryAttemptResult::Error); }; - let mut buffer = deps.buffer_pool.acquire(); + let buffer = deps.buffer_pool.acquire(); let response = if should_sample_backend_timing() { - backend::execute_request_classified_timed(conn.stream_mut(), request, &mut buffer) + backend::execute_request_classified_timed(conn.stream_mut(), request, buffer) .await .map(|(response, ttfb, send, recv)| (response, Some((ttfb, send, recv)))) } else { - backend::execute_request_classified(conn.stream_mut(), request, &mut buffer) + backend::execute_request_classified(conn.stream_mut(), request, buffer) .await .map(|response| (response, None)) }; @@ -233,22 +233,21 @@ async fn execute_backend_query( match response { Ok((response, timings)) => { let Some(status_code) = response.status_code() else { - response.log_warnings(&buffer, "adaptive-precheck", backend_id); + response.log_warnings("adaptive-precheck", backend_id); conn.fail_backend(); return Err(()); }; let single_line_payload = response - .single_line_bytes(&buffer) + .single_line_bytes() .map(crate::cache::CacheIngestResponse::from); let (response, completion) = build_precheck_hit( deps, request, - &response, + response, status_code, single_line_payload, &mut conn, - &mut buffer, ) .await?; @@ -267,11 +266,10 @@ async fn execute_backend_query( async fn build_precheck_hit( deps: &OwnedDeps, request: &RequestContext, - response: &crate::session::backend::BackendReadResult, + response: crate::session::backend::ClassifiedResponse, status_code: StatusCode, single_line_payload: Option, conn: &mut crate::pool::ConnectionGuard, - buffer: &mut crate::pool::PooledBuffer, ) -> Result< ( PrecheckHit, @@ -280,10 +278,10 @@ async fn build_precheck_hit( (), > { if request.has_response_body(status_code) { - return read_complete_precheck_hit(deps, status_code, conn, buffer).await; + return read_complete_precheck_hit(deps, status_code, conn, response).await; } - let completion = response.completion_proof(request).map_err(|_| ())?; + let completion = response.completion_proof().map_err(|_| ())?; let hit = if let Some(payload) = single_line_payload { PrecheckHit::Payload(payload) } else { @@ -296,7 +294,7 @@ async fn read_complete_precheck_hit( deps: &OwnedDeps, status_code: StatusCode, conn: &mut crate::pool::ConnectionGuard, - buffer: &mut crate::pool::PooledBuffer, + classified: crate::session::backend::ClassifiedResponse, ) -> Result< ( PrecheckHit, @@ -310,13 +308,8 @@ async fn read_complete_precheck_hit( .then(crate::pool::ChunkedResponse::default); let completion = if let Some(response) = &mut response { - let (retained, completion) = - crate::session::backend::capture_complete_multiline_response_chunked_optional( - conn.stream_mut(), - buffer, - &deps.buffer_pool, - response, - ) + let (retained, completion) = classified + .capture_isolated_chunked_optional(conn.stream_mut(), &deps.buffer_pool, response) .await .map_err(|_| ())?; if !retained { @@ -325,7 +318,8 @@ async fn read_complete_precheck_hit( } completion } else { - crate::session::backend::observe_complete_multiline_response(conn.stream_mut(), buffer) + classified + .observe_isolated(conn.stream_mut()) .await .map_err(|_| ())? }; From 7dfea683a9325990f0cc3872d43e2db806662529 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 15 Sep 2026 12:14:05 -0600 Subject: [PATCH 10/23] Bind pool reuse to completed response state --- src/client/mod.rs | 10 ++- src/pool/connection_guard.rs | 62 ++++++++++++--- src/session/backend.rs | 27 ------- src/session/handlers/command_execution.rs | 93 +++++++++-------------- src/session/handlers/mod.rs | 11 +-- src/session/handlers/per_command.rs | 16 +--- src/session/multiline_framing.rs | 39 +++------- src/session/precheck.rs | 36 +++------ 8 files changed, 124 insertions(+), 170 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 1fbb7202..9762697d 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -153,8 +153,9 @@ impl NntpClient { let result = Self::parse_stat_response(status_code); if result.is_ok() { - let completion = response.completion_proof()?; - let _reusable = conn.complete_success(completion); + response.completion_proof()?; + conn.mark_response_complete(); + let _reusable = conn.complete_success(); } result } @@ -189,7 +190,7 @@ impl NntpClient { Self::validate_response(status_code)?; - let (captured, completion) = match response + let captured = match response .capture_isolated(conn.stream_mut(), &self.buffer_pool) .await { @@ -199,7 +200,8 @@ impl NntpClient { return Err(error); } }; - let _reusable = conn.complete_success(completion); + conn.mark_response_complete(); + let _reusable = conn.complete_success(); Ok(captured) } diff --git a/src/pool/connection_guard.rs b/src/pool/connection_guard.rs index 9e90cc85..91e84705 100644 --- a/src/pool/connection_guard.rs +++ b/src/pool/connection_guard.rs @@ -13,7 +13,6 @@ use crate::constants::pool::HEALTH_CHECK_TIMEOUT; use crate::pool::deadpool_connection::PooledConnection; use crate::pool::provider::DeadpoolConnectionProvider; -use crate::session::backend::BackendResponseComplete; // ═══════════════════════════════════════════════════════════════════════════ // COMPILE-TIME SAFEGUARDS: Connection Hold Time Limits @@ -78,7 +77,14 @@ const _SALVAGE_NO_LOOP: () = { pub struct ConnectionGuard { conn: Option, provider: DeadpoolConnectionProvider, - active: bool, + phase: ConnectionPhase, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum ConnectionPhase { + Idle, + InFlight, + Complete, } /// A connection whose last backend exchange was proven complete. @@ -93,7 +99,7 @@ impl ConnectionGuard { Self { conn: Some(conn), provider, - active: false, + phase: ConnectionPhase::Idle, } } /// Return a checked-out connection that has not been used yet. @@ -101,7 +107,11 @@ impl ConnectionGuard { /// This is the only successful release available without response-completion /// evidence; callers cannot obtain protocol-stream access from this API. pub fn release_idle(mut self) { - assert!(!self.active, "cannot release a connection after I/O began"); + assert_eq!( + self.phase, + ConnectionPhase::Idle, + "cannot release a connection after I/O began" + ); drop(self.conn.take().expect("ConnectionGuard already consumed")); } @@ -113,10 +123,12 @@ impl ConnectionGuard { /// # Panics /// /// Panics if the guard has already been consumed (double-release). - pub(crate) fn complete_success( - mut self, - _completion: BackendResponseComplete, - ) -> ReusableConnection { + pub(crate) fn complete_success(mut self) -> ReusableConnection { + assert_eq!( + self.phase, + ConnectionPhase::Complete, + "complete_success() requires a fully consumed response" + ); ReusableConnection { _connection: self .conn @@ -164,7 +176,7 @@ impl ConnectionGuard { /// /// Panics if the guard has already been consumed. pub(crate) fn get_mut(&mut self) -> &mut PooledConnection { - self.active = true; + self.phase = ConnectionPhase::InFlight; self.conn .as_mut() .expect("ConnectionGuard already consumed") @@ -174,6 +186,21 @@ impl ConnectionGuard { self.get_mut() } + /// Record that the framer consumed the complete response for this guard. + /// + /// The phase is kept on the guard that owns the socket. This removes the + /// old detached zero-sized completion token, which could be paired with a + /// different connection by a caller. + pub(crate) fn mark_response_complete(&mut self) { + match self.phase { + ConnectionPhase::InFlight => self.phase = ConnectionPhase::Complete, + ConnectionPhase::Complete => {} + ConnectionPhase::Idle => { + panic!("response completion requires an active connection") + } + } + } + /// Transfer ownership to the bounded health-check path without claiming success. pub(crate) fn into_connection_for_health_check(mut self) -> PooledConnection { self.conn @@ -377,9 +404,9 @@ mod tests { let conn = provider.get_pooled_connection().await.unwrap(); assert_eq!(accept_count.load(Ordering::SeqCst), 1); - // complete_success() returns conn to pool (no shutdown) + // An untouched checkout can be returned as idle. let guard = ConnectionGuard::new(conn, provider.clone()); - drop(guard.complete_success(crate::session::backend::BackendResponseComplete::for_test())); + guard.release_idle(); // Second get — pool recycles the existing connection (no new TCP handshake) let _conn2 = provider.get_pooled_connection().await.unwrap(); @@ -391,6 +418,19 @@ mod tests { ); } + /// An active guard cannot be returned as reusable before the framer records + /// a complete response. + #[tokio::test] + #[should_panic(expected = "complete_success() requires a fully consumed response")] + async fn complete_success_rejects_incomplete_response() { + let (port, _accept_count) = spawn_greeting_server().await; + let provider = make_provider(port); + let mut guard = + ConnectionGuard::new(provider.get_pooled_connection().await.unwrap(), provider); + let _ = guard.stream_mut(); + let _ = guard.complete_success(); + } + /// Invariant: drop without `complete_success()` removes the connection from the pool. /// /// The guard shuts down the socket; pool recycle detects EOF diff --git a/src/session/backend.rs b/src/session/backend.rs index 3d9e3431..75954cda 100644 --- a/src/session/backend.rs +++ b/src/session/backend.rs @@ -21,33 +21,6 @@ pub(crate) use crate::session::multiline_framing::BackendResponseOrder; pub(crate) use crate::session::multiline_framing::ClassifiedResponse; -#[must_use] -pub(crate) struct BackendResponseComplete(()); - -impl BackendResponseComplete { - pub(super) fn from_reusable_response( - reuse: &crate::session::response_transfer::ResponseConnectionReuse, - ) -> Option { - if matches!( - reuse, - crate::session::response_transfer::ResponseConnectionReuse::Reusable - ) { - Some(Self(())) - } else { - None - } - } - - pub(crate) const fn response() -> Self { - Self(()) - } - - #[cfg(test)] - pub(crate) const fn for_test() -> Self { - Self(()) - } -} - /// Failure while reading a complete single-line backend reply into caller-owned /// scratch storage. #[derive(Debug)] diff --git a/src/session/handlers/command_execution.rs b/src/session/handlers/command_execution.rs index 23b1683a..90713e07 100644 --- a/src/session/handlers/command_execution.rs +++ b/src/session/handlers/command_execution.rs @@ -315,15 +315,15 @@ impl ClientSession { ) .await; } - let completion = buffer + buffer .observe(conn.stream_mut(), &self.buffer_pool, backend_id) .await .map_err(SessionError::from)?; + conn.mark_response_complete(); self.release_or_reuse_connection( conn, backend_id, request, - completion, Some(state.backend_connection), ); return Ok(BackendAttemptResult::ArticleNotFound { missing }); @@ -428,17 +428,18 @@ impl ClientSession { }; let status_code = response.status_code(); - let completion = match response + match response .observe(conn.stream_mut(), &self.buffer_pool, backend_id) .await { - Ok(completion) => completion, + Ok(()) => {} Err(_) => { conn.fail_backend(); return RetryStatProbeOutcome::Unavailable(backend_id); } - }; - let _ = conn.complete_success(completion); + } + conn.mark_response_complete(); + let _ = conn.complete_success(); if status_code.is_some_and(|status| status.as_u16() == 430) { RetryStatProbeOutcome::Missing(backend_id) @@ -557,17 +558,18 @@ impl ClientSession { } }; let status_code = response.status_code(); - let completion = match response + match response .observe(conn.stream_mut(), &buffer_pool, backend_id) .await { - Ok(completion) => completion, + Ok(()) => {} Err(_) => { conn.fail_backend(); return; } - }; - let _ = conn.complete_success(completion); + } + conn.mark_response_complete(); + let _ = conn.complete_success(); if status_code.is_some_and(|status| status.as_u16() == 430) && let Ok(msg_id) = crate::types::MessageId::new(msg_id_text) @@ -751,7 +753,7 @@ impl ClientSession { has_response_body, "Writing backend response to client" ); - let (bytes_written, completion) = match self + let bytes_written = match self .write_response_to_client( conn.stream_mut(), client_write, @@ -771,6 +773,7 @@ impl ClientSession { )); } }; + conn.mark_response_complete(); if let Err(e) = client_write.flush().await { return Err(self.handle_response_transfer_error( conn, @@ -794,13 +797,7 @@ impl ClientSession { params.request.request_wire_len().as_u64(), bytes_written, ); - self.release_or_reuse_connection( - conn, - backend_id, - params.request, - completion, - backend_connection, - ); + self.release_or_reuse_connection(conn, backend_id, params.request, backend_connection); Ok(RequestResponseMetadata::new( status_code, @@ -810,14 +807,12 @@ impl ClientSession { pub(super) fn handle_response_transfer_error( &self, - conn: crate::pool::ConnectionGuard, + mut conn: crate::pool::ConnectionGuard, backend_id: BackendId, request: &RequestContext, error: ResponseTransferError, ) -> SessionError { let reuse = crate::session::response_transfer::connection_reuse_after_response(&conn); - let completion = - crate::session::backend::BackendResponseComplete::from_reusable_response(&reuse); match error.pool_fate(reuse) { BackendConnectionOutcome::BackendFailed => { warn!( @@ -851,9 +846,8 @@ impl ClientSession { command_verb = ?request.verb(), "Response write error left backend connection reusable; releasing it back to pool" ); - let completion = - completion.expect("healthy response fate requires a reusable connection"); - let _ = conn.complete_success(completion); + conn.mark_response_complete(); + let _ = conn.complete_success(); } } @@ -865,7 +859,6 @@ impl ClientSession { conn: crate::pool::ConnectionGuard, backend_id: BackendId, request: &RequestContext, - completion: crate::session::backend::BackendResponseComplete, backend_connection: Option<&mut Option>, ) { let reuse = crate::session::response_transfer::connection_reuse_after_response(&conn); @@ -907,7 +900,7 @@ impl ClientSession { pool_waiting = status.waiting, "Direct per-command response finished cleanly; keeping backend connection for this client batch" ); - *slot = Some(BackendLease::new(backend_id, conn, completion)); + *slot = Some(BackendLease::new(backend_id, conn)); } else { debug!( client = %self.client_addr, @@ -918,7 +911,7 @@ impl ClientSession { pending_bytes = conn.pending_bytes_len(), "Direct per-command response finished cleanly; releasing backend connection" ); - let _ = conn.complete_success(completion); + let _ = conn.complete_success(); } } BackendConnectionOutcome::BackendFailed => unreachable!( @@ -953,7 +946,6 @@ impl ClientSession { Some(BackendLease { backend_id: cached_backend_id, connection: mut guard, - completion, }) if cached_backend_id == backend_id => { if !self.cached_batch_connection_is_healthy(&mut guard, backend_id, request) { guard.fail_backend(); @@ -986,8 +978,7 @@ impl ClientSession { pending_bytes = guard.pending_bytes_len(), "Using idle pool capacity before reusing cached batch connection" ); - *backend_connection = - Some(BackendLease::new(cached_backend_id, guard, completion)); + *backend_connection = Some(BackendLease::new(cached_backend_id, guard)); match provider.checkout_connection_guard().await { Ok(conn) => conn, Err(err) => { @@ -1028,7 +1019,6 @@ impl ClientSession { let BackendLease { backend_id: cached_backend_id, connection: cached_conn, - completion, } = cached; trace!( client = %self.client_addr, @@ -1040,7 +1030,7 @@ impl ClientSession { pending_bytes = cached_conn.pending_bytes_len(), "Releasing cached backend connection before switching backend" ); - let _ = cached_conn.complete_success(completion); + let _ = cached_conn.complete_success(); let checkout_status = provider.status_counts(); trace!( client = %self.client_addr, @@ -1162,10 +1152,11 @@ impl ClientSession { ); Ok((probe_response, probe_timings)) } else { - let _completion = probe_response + probe_response .observe(guard.stream_mut(), &self.buffer_pool, backend.backend_id()) .await .map_err(SessionError::from)?; + guard.mark_response_complete(); let _ = probe_timings; self.execute_and_read_response(guard.stream_mut(), backend, request) .await @@ -1237,7 +1228,7 @@ impl ClientSession { request: &mut RequestContext, availability: &mut crate::cache::ArticleAvailability, backend_to_client_bytes: &mut BackendToClientBytes, - ) -> Result + ) -> Result<(), AlreadySentResponseError> where W: AsyncWrite + Unpin, { @@ -1268,10 +1259,11 @@ impl ClientSession { ) .await; } - let completion = buffer + buffer .observe(conn.stream_mut(), &self.buffer_pool, backend_id) .await .map_err(AlreadySentResponseError::Transfer)?; + conn.mark_response_complete(); self.send_430_to_client(client_write, backend_to_client_bytes) .await .map_err(|error| { @@ -1280,7 +1272,7 @@ impl ClientSession { client_write.flush().await.map_err(|error| { AlreadySentResponseError::Transfer(classify_response_write_err(error)) })?; - return Ok(completion); + return Ok(()); } let params = ResponseWriteParams { @@ -1288,10 +1280,11 @@ impl ClientSession { article_request: crate::command::CommandHandler::article_lookup_request(request), status_code, }; - let (bytes_written, completion) = self + let bytes_written = self .write_response_to_client(conn.stream_mut(), client_write, backend, buffer, params) .await .map_err(AlreadySentResponseError::Transfer)?; + conn.mark_response_complete(); client_write.flush().await.map_err(|error| { AlreadySentResponseError::Transfer(classify_response_write_err(error)) })?; @@ -1308,7 +1301,7 @@ impl ClientSession { ); request.record_backend_response(backend_id, response); *backend_to_client_bytes = backend_to_client_bytes.add(response.wire_len().get()); - Ok(completion) + Ok(()) } #[inline] @@ -1345,7 +1338,7 @@ impl ClientSession { backend: &ArticleBackend, backend_bytes: crate::session::backend::ClassifiedResponse, params: ResponseWriteParams<'_>, - ) -> Result<(u64, crate::session::backend::BackendResponseComplete), ResponseTransferError> + ) -> Result where W: AsyncWrite + Unpin, { @@ -1401,10 +1394,7 @@ impl ClientSession { if let Some(backend) = backend_after_write { self.apply_cache_action(cache_action, params, backend, captured); } - Ok(( - bytes_written, - crate::session::backend::BackendResponseComplete::response(), - )) + Ok(bytes_written) } fn apply_cache_action( @@ -2489,11 +2479,7 @@ mod tests { .checkout_connection_guard() .await .expect("initial pooled connection should be created"); - let mut backend_connection = Some(BackendLease::new( - backend_id, - conn, - crate::session::backend::BackendResponseComplete::for_test(), - )); + let mut backend_connection = Some(BackendLease::new(backend_id, conn)); if let Some(lease) = backend_connection.as_mut() { lease .connection @@ -2511,7 +2497,7 @@ mod tests { ) .await .expect("checkout should recover with a fresh pooled connection"); - drop(guard.complete_success(crate::session::backend::BackendResponseComplete::for_test())); + guard.release_idle(); assert_eq!( accept_count.load(Ordering::SeqCst), @@ -2982,7 +2968,7 @@ mod tests { assert_eq!(client_write.writes, b"223 0 status\r\n"); let reused = provider.checkout_connection_guard().await.unwrap(); - drop(reused.complete_success(crate::session::backend::BackendResponseComplete::for_test())); + reused.release_idle(); assert_eq!( accept_count.load(Ordering::SeqCst), 1, @@ -3036,7 +3022,7 @@ mod tests { )); let reused = provider.checkout_connection_guard().await.unwrap(); - drop(reused.complete_success(crate::session::backend::BackendResponseComplete::for_test())); + reused.release_idle(); assert_eq!( accept_count.load(Ordering::SeqCst), 1, @@ -3092,10 +3078,7 @@ mod tests { )); let replacement = provider.checkout_connection_guard().await.unwrap(); - drop( - replacement - .complete_success(crate::session::backend::BackendResponseComplete::for_test()), - ); + replacement.release_idle(); assert_eq!( accept_count.load(Ordering::SeqCst), 2, diff --git a/src/session/handlers/mod.rs b/src/session/handlers/mod.rs index fd357f2b..ae064eaf 100644 --- a/src/session/handlers/mod.rs +++ b/src/session/handlers/mod.rs @@ -10,30 +10,23 @@ //! All handler functions are implemented as methods on `ClientSession` in their //! respective modules. No need to re-export since they're all impl blocks. use crate::pool::ConnectionGuard; -use crate::session::backend::BackendResponseComplete; use crate::types::BackendId; pub(super) struct BackendLease { pub(super) backend_id: BackendId, pub(super) connection: ConnectionGuard, - pub(super) completion: BackendResponseComplete, } impl BackendLease { - pub(super) const fn new( - backend_id: BackendId, - connection: ConnectionGuard, - completion: BackendResponseComplete, - ) -> Self { + pub(super) const fn new(backend_id: BackendId, connection: ConnectionGuard) -> Self { Self { backend_id, connection, - completion, } } pub(super) fn complete_success(self) { - let _ = self.connection.complete_success(self.completion); + let _ = self.connection.complete_success(); } pub(super) fn fail_backend(self) { diff --git a/src/session/handlers/per_command.rs b/src/session/handlers/per_command.rs index b7ec364b..546940de 100644 --- a/src/session/handlers/per_command.rs +++ b/src/session/handlers/per_command.rs @@ -1074,7 +1074,6 @@ impl ClientSession { batch, backend_connection, } = params; - let mut last_completion = None; let mut unread_requests = requests.into_iter().enumerate(); while let Some((i, request_route)) = unread_requests.next() { match request_route { @@ -1111,8 +1110,7 @@ impl ClientSession { )) })?; match result { - Ok(completion) => { - last_completion = Some(completion); + Ok(()) => { guard.complete(); } Err( @@ -1161,12 +1159,10 @@ impl ClientSession { } client_writer.get_mut().flush().await?; - let completion = - last_completion.expect("sent article window contained an upstream request"); if conn.has_pending_bytes() { conn.fail_client(); } else { - *backend_connection.slot() = Some(BackendLease::new(backend_id, conn, completion)); + *backend_connection.slot() = Some(BackendLease::new(backend_id, conn)); } Ok(()) } @@ -1353,11 +1349,7 @@ mod tests { let handle = tokio::spawn(async move { let _batch = super::BatchBackendConnection { - conn: Some(super::BackendLease::new( - BackendId::from_index(0), - conn, - crate::session::backend::BackendResponseComplete::for_test(), - )), + conn: Some(super::BackendLease::new(BackendId::from_index(0), conn)), }; tokio::time::sleep(Duration::from_secs(1)).await; drop(_batch); @@ -1369,7 +1361,7 @@ mod tests { tokio::time::sleep(Duration::from_millis(50)).await; let next = provider.checkout_connection_guard().await.unwrap(); - drop(next.complete_success(crate::session::backend::BackendResponseComplete::for_test())); + next.release_idle(); assert_eq!( accept_count.load(Ordering::SeqCst), 2, diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 74d802b8..b4a197fd 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -659,17 +659,11 @@ impl ClassifiedResponse { mut self, conn: &mut crate::stream::ConnectionStream, pool: &crate::pool::BufferPool, - ) -> anyhow::Result<( - crate::pool::PooledBuffer, - crate::session::backend::BackendResponseComplete, - )> { + ) -> anyhow::Result { match self.frame { Ok(ResponseFrame::SingleLine { framed, .. }) => { framed.require_isolated()?; - Ok(( - self.buffer, - crate::session::backend::BackendResponseComplete::response(), - )) + Ok(self.buffer) } Ok(ResponseFrame::Multiline { .. }) => { let mut capture = pool.acquire_capture(); @@ -680,10 +674,7 @@ impl ClassifiedResponse { // Framing owns the already-proven terminator. The public // article parser receives the status line and payload only. strip_captured_terminator(&mut capture)?; - Ok(( - capture, - crate::session::backend::BackendResponseComplete::response(), - )) + Ok(capture) } Err(error) => anyhow::bail!("cannot capture invalid backend response: {error:?}"), } @@ -694,7 +685,7 @@ impl ClassifiedResponse { conn: &mut crate::stream::ConnectionStream, pool: &crate::pool::BufferPool, captured: &mut crate::pool::ChunkedResponse, - ) -> anyhow::Result<(bool, crate::session::backend::BackendResponseComplete)> { + ) -> anyhow::Result { match self.frame { Ok(ResponseFrame::Multiline { .. }) => {} _ => anyhow::bail!("isolated multiline capture requires a multiline response"), @@ -707,16 +698,13 @@ impl ClassifiedResponse { ) .await .map_err(|error| anyhow::anyhow!("backend multiline response capture failed: {error:?}"))?; - Ok(( - retained, - crate::session::backend::BackendResponseComplete::response(), - )) + Ok(retained) } pub(crate) async fn observe_isolated( mut self, conn: &mut crate::stream::ConnectionStream, - ) -> anyhow::Result { + ) -> anyhow::Result<()> { match self.frame { Ok(ResponseFrame::SingleLine { framed, .. }) => framed.require_isolated()?, Ok(ResponseFrame::Multiline { .. }) => { @@ -728,7 +716,7 @@ impl ClassifiedResponse { } Err(error) => anyhow::bail!("cannot observe invalid backend response: {error:?}"), } - Ok(crate::session::backend::BackendResponseComplete::response()) + Ok(()) } pub(crate) fn log_warnings( @@ -741,13 +729,11 @@ impl ClassifiedResponse { } } - pub(crate) fn completion_proof( - &self, - ) -> anyhow::Result { + pub(crate) fn completion_proof(&self) -> anyhow::Result<()> { match &self.frame { Ok(ResponseFrame::SingleLine { framed, .. }) => { framed.require_isolated()?; - Ok(crate::session::backend::BackendResponseComplete::response()) + Ok(()) } Ok(ResponseFrame::Multiline { .. }) => { anyhow::bail!("multiline response requires framer-owned completion") @@ -784,12 +770,9 @@ impl ClassifiedResponse { conn: &mut crate::stream::ConnectionStream, pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, - ) -> Result< - crate::session::backend::BackendResponseComplete, - crate::session::response_transfer::ResponseTransferError, - > { + ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { self.stream(conn, pool, backend_id)?.observe().await?; - Ok(crate::session::backend::BackendResponseComplete::response()) + Ok(()) } pub(crate) async fn capture_and_write( diff --git a/src/session/precheck.rs b/src/session/precheck.rs index a2ffcf34..82621590 100644 --- a/src/session/precheck.rs +++ b/src/session/precheck.rs @@ -241,7 +241,7 @@ async fn execute_backend_query( .single_line_bytes() .map(crate::cache::CacheIngestResponse::from); - let (response, completion) = build_precheck_hit( + let response = build_precheck_hit( deps, request, response, @@ -253,7 +253,8 @@ async fn execute_backend_query( let result = classify_precheck_result(deps, backend, status_code, timings, response); - let _ = conn.complete_success(completion); + conn.mark_response_complete(); + let _ = conn.complete_success(); Ok(result) } Err(_) => { @@ -270,24 +271,18 @@ async fn build_precheck_hit( status_code: StatusCode, single_line_payload: Option, conn: &mut crate::pool::ConnectionGuard, -) -> Result< - ( - PrecheckHit, - crate::session::backend::BackendResponseComplete, - ), - (), -> { +) -> Result { if request.has_response_body(status_code) { return read_complete_precheck_hit(deps, status_code, conn, response).await; } - let completion = response.completion_proof().map_err(|_| ())?; + response.completion_proof().map_err(|_| ())?; let hit = if let Some(payload) = single_line_payload { PrecheckHit::Payload(payload) } else { PrecheckHit::Availability(status_code) }; - Ok((hit, completion)) + Ok(hit) } async fn read_complete_precheck_hit( @@ -295,33 +290,26 @@ async fn read_complete_precheck_hit( status_code: StatusCode, conn: &mut crate::pool::ConnectionGuard, classified: crate::session::backend::ClassifiedResponse, -) -> Result< - ( - PrecheckHit, - crate::session::backend::BackendResponseComplete, - ), - (), -> { +) -> Result { let mut response = deps .cache .stores_payload_responses() .then(crate::pool::ChunkedResponse::default); - let completion = if let Some(response) = &mut response { - let (retained, completion) = classified + if let Some(response) = &mut response { + let retained = classified .capture_isolated_chunked_optional(conn.stream_mut(), &deps.buffer_pool, response) .await .map_err(|_| ())?; if !retained { response.clear(); - return Ok((PrecheckHit::Availability(status_code), completion)); + return Ok(PrecheckHit::Availability(status_code)); } - completion } else { classified .observe_isolated(conn.stream_mut()) .await - .map_err(|_| ())? + .map_err(|_| ())?; }; let hit = if let Some(response) = response { @@ -329,7 +317,7 @@ async fn read_complete_precheck_hit( } else { PrecheckHit::Availability(status_code) }; - Ok((hit, completion)) + Ok(hit) } fn classify_precheck_result( From eceb7658b5ee2f3d753a9b7755c862f1084bf99e Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Tue, 15 Sep 2026 12:23:04 -0600 Subject: [PATCH 11/23] Bind response transfers to their connection guard --- src/client/mod.rs | 3 +- src/session/handlers/command_execution.rs | 36 ++++-------- src/session/multiline_framing.rs | 71 ++++++++++++++--------- src/session/precheck.rs | 9 +-- 4 files changed, 59 insertions(+), 60 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 9762697d..4a83a4bf 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -191,7 +191,7 @@ impl NntpClient { Self::validate_response(status_code)?; let captured = match response - .capture_isolated(conn.stream_mut(), &self.buffer_pool) + .capture_isolated(&mut conn, &self.buffer_pool) .await { Ok(result) => result, @@ -200,7 +200,6 @@ impl NntpClient { return Err(error); } }; - conn.mark_response_complete(); let _reusable = conn.complete_success(); Ok(captured) } diff --git a/src/session/handlers/command_execution.rs b/src/session/handlers/command_execution.rs index 90713e07..0bbd48e6 100644 --- a/src/session/handlers/command_execution.rs +++ b/src/session/handlers/command_execution.rs @@ -316,10 +316,9 @@ impl ClientSession { .await; } buffer - .observe(conn.stream_mut(), &self.buffer_pool, backend_id) + .observe(&mut conn, &self.buffer_pool, backend_id) .await .map_err(SessionError::from)?; - conn.mark_response_complete(); self.release_or_reuse_connection( conn, backend_id, @@ -429,7 +428,7 @@ impl ClientSession { let status_code = response.status_code(); match response - .observe(conn.stream_mut(), &self.buffer_pool, backend_id) + .observe(&mut conn, &self.buffer_pool, backend_id) .await { Ok(()) => {} @@ -438,7 +437,6 @@ impl ClientSession { return RetryStatProbeOutcome::Unavailable(backend_id); } } - conn.mark_response_complete(); let _ = conn.complete_success(); if status_code.is_some_and(|status| status.as_u16() == 430) { @@ -558,17 +556,13 @@ impl ClientSession { } }; let status_code = response.status_code(); - match response - .observe(conn.stream_mut(), &buffer_pool, backend_id) - .await - { + match response.observe(&mut conn, &buffer_pool, backend_id).await { Ok(()) => {} Err(_) => { conn.fail_backend(); return; } } - conn.mark_response_complete(); let _ = conn.complete_success(); if status_code.is_some_and(|status| status.as_u16() == 430) @@ -754,13 +748,7 @@ impl ClientSession { "Writing backend response to client" ); let bytes_written = match self - .write_response_to_client( - conn.stream_mut(), - client_write, - backend, - backend_bytes, - params, - ) + .write_response_to_client(&mut conn, client_write, backend, backend_bytes, params) .await { Ok(result) => result, @@ -773,7 +761,6 @@ impl ClientSession { )); } }; - conn.mark_response_complete(); if let Err(e) = client_write.flush().await { return Err(self.handle_response_transfer_error( conn, @@ -1153,10 +1140,9 @@ impl ClientSession { Ok((probe_response, probe_timings)) } else { probe_response - .observe(guard.stream_mut(), &self.buffer_pool, backend.backend_id()) + .observe(&mut guard, &self.buffer_pool, backend.backend_id()) .await .map_err(SessionError::from)?; - guard.mark_response_complete(); let _ = probe_timings; self.execute_and_read_response(guard.stream_mut(), backend, request) .await @@ -1260,10 +1246,9 @@ impl ClientSession { .await; } buffer - .observe(conn.stream_mut(), &self.buffer_pool, backend_id) + .observe(conn, &self.buffer_pool, backend_id) .await .map_err(AlreadySentResponseError::Transfer)?; - conn.mark_response_complete(); self.send_430_to_client(client_write, backend_to_client_bytes) .await .map_err(|error| { @@ -1281,10 +1266,9 @@ impl ClientSession { status_code, }; let bytes_written = self - .write_response_to_client(conn.stream_mut(), client_write, backend, buffer, params) + .write_response_to_client(&mut *conn, client_write, backend, buffer, params) .await .map_err(AlreadySentResponseError::Transfer)?; - conn.mark_response_complete(); client_write.flush().await.map_err(|error| { AlreadySentResponseError::Transfer(classify_response_write_err(error)) })?; @@ -1333,7 +1317,7 @@ impl ClientSession { /// without string/downcast inspection. async fn write_response_to_client( &self, - pooled_conn: &mut crate::stream::ConnectionStream, + pooled_conn: &mut crate::pool::ConnectionGuard, client_write: &mut W, backend: &ArticleBackend, backend_bytes: crate::session::backend::ClassifiedResponse, @@ -1449,7 +1433,7 @@ impl ClientSession { async fn write_response_without_retention( &self, - pooled_conn: &mut crate::stream::ConnectionStream, + pooled_conn: &mut crate::pool::ConnectionGuard, client_write: &mut W, backend_id: BackendId, backend_bytes: crate::session::backend::ClassifiedResponse, @@ -1465,7 +1449,7 @@ impl ClientSession { #[allow(clippy::too_many_arguments)] async fn write_response_with_retention( &self, - pooled_conn: &mut crate::stream::ConnectionStream, + pooled_conn: &mut crate::pool::ConnectionGuard, client_write: &mut W, backend_id: BackendId, backend_bytes: crate::session::backend::ClassifiedResponse, diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index b4a197fd..2cbcd87a 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -657,23 +657,25 @@ impl ClassifiedResponse { pub(crate) async fn capture_isolated( mut self, - conn: &mut crate::stream::ConnectionStream, + conn: &mut crate::pool::ConnectionGuard, pool: &crate::pool::BufferPool, ) -> anyhow::Result { match self.frame { Ok(ResponseFrame::SingleLine { framed, .. }) => { framed.require_isolated()?; + conn.mark_response_complete(); Ok(self.buffer) } Ok(ResponseFrame::Multiline { .. }) => { let mut capture = pool.acquire_capture(); - IsolatedMultilineResponse::begin(conn, &mut self.buffer) + IsolatedMultilineResponse::begin(conn.stream_mut(), &mut self.buffer) .map_err(isolated_multiline_error)? .capture_into(&mut capture) .await?; // Framing owns the already-proven terminator. The public // article parser receives the status line and payload only. strip_captured_terminator(&mut capture)?; + conn.mark_response_complete(); Ok(capture) } Err(error) => anyhow::bail!("cannot capture invalid backend response: {error:?}"), @@ -682,7 +684,7 @@ impl ClassifiedResponse { pub(crate) async fn capture_isolated_chunked_optional( mut self, - conn: &mut crate::stream::ConnectionStream, + conn: &mut crate::pool::ConnectionGuard, pool: &crate::pool::BufferPool, captured: &mut crate::pool::ChunkedResponse, ) -> anyhow::Result { @@ -691,24 +693,25 @@ impl ClassifiedResponse { _ => anyhow::bail!("isolated multiline capture requires a multiline response"), } let retained = capture_isolated_multiline_response_chunked_optional( - conn, + conn.stream_mut(), &mut self.buffer, pool, captured, ) .await .map_err(|error| anyhow::anyhow!("backend multiline response capture failed: {error:?}"))?; + conn.mark_response_complete(); Ok(retained) } pub(crate) async fn observe_isolated( mut self, - conn: &mut crate::stream::ConnectionStream, + conn: &mut crate::pool::ConnectionGuard, ) -> anyhow::Result<()> { match self.frame { Ok(ResponseFrame::SingleLine { framed, .. }) => framed.require_isolated()?, Ok(ResponseFrame::Multiline { .. }) => { - observe_isolated_multiline_response(conn, &mut self.buffer) + observe_isolated_multiline_response(conn.stream_mut(), &mut self.buffer) .await .map_err(|error| { anyhow::anyhow!("backend multiline response drain failed: {error:?}") @@ -716,6 +719,7 @@ impl ClassifiedResponse { } Err(error) => anyhow::bail!("cannot observe invalid backend response: {error:?}"), } + conn.mark_response_complete(); Ok(()) } @@ -755,37 +759,47 @@ impl ClassifiedResponse { pub(crate) async fn write( mut self, - conn: &mut crate::stream::ConnectionStream, + conn: &mut crate::pool::ConnectionGuard, writer: &mut W, pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, ) -> Result { - let stats = self.stream(conn, pool, backend_id)?.write(writer).await?; + let stats = self + .stream(conn.stream_mut(), pool, backend_id)? + .write(writer) + .await?; stats.record(); + conn.mark_response_complete(); Ok(stats.bytes_written_u64()) } pub(crate) async fn observe( mut self, - conn: &mut crate::stream::ConnectionStream, + conn: &mut crate::pool::ConnectionGuard, pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - self.stream(conn, pool, backend_id)?.observe().await?; + self.stream(conn.stream_mut(), pool, backend_id)? + .observe() + .await?; + conn.mark_response_complete(); Ok(()) } pub(crate) async fn capture_and_write( mut self, - conn: &mut crate::stream::ConnectionStream, + conn: &mut crate::pool::ConnectionGuard, writer: &mut W, captured: &mut crate::pool::ChunkedResponse, pool: &crate::pool::BufferPool, backend_id: crate::types::BackendId, ) -> Result<(u64, bool), crate::session::response_transfer::ResponseTransferError> { - self.stream(conn, pool, backend_id)? + let result = self + .stream(conn.stream_mut(), pool, backend_id)? .capture_and_write(writer, captured, MAX_CAPTURED_MULTILINE_RESPONSE_BYTES) - .await + .await?; + conn.mark_response_complete(); + Ok(result) } } @@ -1323,8 +1337,14 @@ async fn write_response( ClassifiedResponse::read(conn, request, buffer) .await .map_err(crate::session::response_transfer::ResponseTransferError::Io)? - .write(conn, writer, pool, backend_id) + .stream(conn, pool, backend_id) + .expect("classified response should stream") + .write(writer) .await + .map(|stats| { + stats.record(); + stats.bytes_written_u64() + }) } impl MultilineFramer { @@ -1898,8 +1918,6 @@ mod tests { response.completion_proof().is_err(), "an isolated request must not authorize reuse with a following reply in its buffer" ); - let mut conn = loopback_connection_stream().await; - assert!(response.capture_isolated(&mut conn, &pool).await.is_err()); } #[tokio::test] @@ -1931,19 +1949,17 @@ mod tests { buffer.copy_from_slice(&wire[..split]); let tail = [wire[split..].as_ref(), next.as_slice()].concat(); let mut conn = mock_backend_conn(vec![tail]).await; - let response = ClassifiedResponse::read(&mut conn, &request, buffer) + let mut response = ClassifiedResponse::read(&mut conn, &request, buffer) .await .unwrap(); let mut writer = RecordingWriter::default(); let bytes = response - .write( - &mut conn, - &mut writer, - &pool, - crate::types::BackendId::from_index(0), - ) + .stream(&mut conn, &pool, crate::types::BackendId::from_index(0)) + .unwrap() + .write(&mut writer) .await - .unwrap(); + .unwrap() + .bytes_written_u64(); assert_eq!( bytes, wire.len() as u64, @@ -3521,10 +3537,13 @@ mod contracts { pool: &crate::pool::BufferPool, writer: &mut (impl AsyncWrite + Unpin), ) { - let response = ClassifiedResponse::read(conn, request, buffer) + let mut response = ClassifiedResponse::read(conn, request, buffer) .await .expect("classified response"); - let transfer = response.write(conn, writer, pool, crate::types::BackendId::from_index(0)); + let transfer = response + .stream(conn, pool, crate::types::BackendId::from_index(0)) + .expect("classified response should stream") + .write(writer); #[cfg(response_contract = "response_twice")] let _conflicting = response.status_code(); std::hint::black_box(transfer); diff --git a/src/session/precheck.rs b/src/session/precheck.rs index 82621590..e8dd4586 100644 --- a/src/session/precheck.rs +++ b/src/session/precheck.rs @@ -253,7 +253,6 @@ async fn execute_backend_query( let result = classify_precheck_result(deps, backend, status_code, timings, response); - conn.mark_response_complete(); let _ = conn.complete_success(); Ok(result) } @@ -277,6 +276,7 @@ async fn build_precheck_hit( } response.completion_proof().map_err(|_| ())?; + conn.mark_response_complete(); let hit = if let Some(payload) = single_line_payload { PrecheckHit::Payload(payload) } else { @@ -298,7 +298,7 @@ async fn read_complete_precheck_hit( if let Some(response) = &mut response { let retained = classified - .capture_isolated_chunked_optional(conn.stream_mut(), &deps.buffer_pool, response) + .capture_isolated_chunked_optional(conn, &deps.buffer_pool, response) .await .map_err(|_| ())?; if !retained { @@ -306,10 +306,7 @@ async fn read_complete_precheck_hit( return Ok(PrecheckHit::Availability(status_code)); } } else { - classified - .observe_isolated(conn.stream_mut()) - .await - .map_err(|_| ())?; + classified.observe_isolated(conn).await.map_err(|_| ())?; }; let hit = if let Some(response) = response { From 30978c02455295ef5cb31840679c2c6028f18ef3 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 11:42:32 -0600 Subject: [PATCH 12/23] Bind single-line completion to connection guard --- src/client/mod.rs | 3 +-- src/session/multiline_framing.rs | 18 +++++++++++++++--- src/session/precheck.rs | 3 +-- 3 files changed, 17 insertions(+), 7 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 4a83a4bf..2f711546 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -153,8 +153,7 @@ impl NntpClient { let result = Self::parse_stat_response(status_code); if result.is_ok() { - response.completion_proof()?; - conn.mark_response_complete(); + response.complete_single_line(&mut conn)?; let _reusable = conn.complete_success(); } result diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 2cbcd87a..7562ca14 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -733,7 +733,7 @@ impl ClassifiedResponse { } } - pub(crate) fn completion_proof(&self) -> anyhow::Result<()> { + pub(crate) fn require_isolated_single_line(&self) -> anyhow::Result<()> { match &self.frame { Ok(ResponseFrame::SingleLine { framed, .. }) => { framed.require_isolated()?; @@ -742,10 +742,22 @@ impl ClassifiedResponse { Ok(ResponseFrame::Multiline { .. }) => { anyhow::bail!("multiline response requires framer-owned completion") } - Err(_) => anyhow::bail!("cannot prove completion for an invalid backend response"), + Err(_) => anyhow::bail!("cannot complete an invalid backend response"), } } + /// Complete a successfully isolated single-line response on its owning + /// connection. Keeping the isolation check and guard transition together + /// prevents callers from pairing an independent proof with a connection. + pub(crate) fn complete_single_line( + &self, + conn: &mut crate::pool::ConnectionGuard, + ) -> anyhow::Result<()> { + self.require_isolated_single_line()?; + conn.mark_response_complete(); + Ok(()) + } + fn stream<'a>( &'a mut self, conn: &'a mut crate::stream::ConnectionStream, @@ -1915,7 +1927,7 @@ mod tests { Some(b"223 0 \r\n".as_slice()) ); assert!( - response.completion_proof().is_err(), + response.require_isolated_single_line().is_err(), "an isolated request must not authorize reuse with a following reply in its buffer" ); } diff --git a/src/session/precheck.rs b/src/session/precheck.rs index e8dd4586..47aeadbd 100644 --- a/src/session/precheck.rs +++ b/src/session/precheck.rs @@ -275,8 +275,7 @@ async fn build_precheck_hit( return read_complete_precheck_hit(deps, status_code, conn, response).await; } - response.completion_proof().map_err(|_| ())?; - conn.mark_response_complete(); + response.complete_single_line(conn).map_err(|_| ())?; let hit = if let Some(payload) = single_line_payload { PrecheckHit::Payload(payload) } else { From 444cce39c7602f07813378600cc34bf6c962414b Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 11:43:52 -0600 Subject: [PATCH 13/23] Consume single-line completion state --- src/session/multiline_framing.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 7562ca14..11392537 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -750,7 +750,7 @@ impl ClassifiedResponse { /// connection. Keeping the isolation check and guard transition together /// prevents callers from pairing an independent proof with a connection. pub(crate) fn complete_single_line( - &self, + self, conn: &mut crate::pool::ConnectionGuard, ) -> anyhow::Result<()> { self.require_isolated_single_line()?; From 10d90bbc6fbb40e5388adc98ba5e9a4f91b0bf00 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 11:54:17 -0600 Subject: [PATCH 14/23] Keep packed suffix with framed response window --- src/session/multiline_framing.rs | 52 ++++++-------------------------- 1 file changed, 10 insertions(+), 42 deletions(-) diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 11392537..8a3b697e 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -136,22 +136,13 @@ impl CompleteResponseWindow { writer: &mut W, io_buffer: &mut crate::pool::PooledBuffer, conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, total_len: usize, ) -> Result where W: AsyncWrite + Unpin, { - write_response_chunk_preserving_suffix_on_error( - writer, - io_buffer, - conn, - pool, - total_len, - &self.response, - &self.next_response_input, - ) - .await + write_response_chunk_preserving_suffix_on_error(writer, io_buffer, conn, total_len, self) + .await } fn extend_capture_from(&self, source: &[u8], capture: &mut crate::pool::PooledBuffer) { @@ -214,21 +205,6 @@ impl CompleteResponseWindow { } } -fn queue_packed_next_response_input( - io_buffer: &mut crate::pool::PooledBuffer, - conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, - next_response_input: &Range, - total_len: usize, -) -> Result<(), crate::session::response_transfer::ResponseTransferError> { - if next_response_input.start < total_len { - let old = std::mem::replace(io_buffer, pool.acquire()); - conn.queue_pooled_pending_bytes_first(old, next_response_input.clone()) - .map_err(crate::session::response_transfer::ResponseTransferError::Io)?; - } - Ok(()) -} - /// Reuse the framer-owned packed suffix, or acquire an empty read buffer. pub(crate) fn take_queued_input_or_acquire_empty( conn: &mut crate::stream::ConnectionStream, @@ -245,24 +221,22 @@ async fn write_response_chunk_preserving_suffix_on_error( writer: &mut W, io_buffer: &mut crate::pool::PooledBuffer, conn: &mut crate::stream::ConnectionStream, - pool: &crate::pool::BufferPool, total_len: usize, - response: &Range, - next_response_input: &Range, + window: &CompleteResponseWindow, ) -> Result where W: AsyncWrite + Unpin, { - let chunk = &io_buffer[..total_len][response.clone()]; + let chunk = &io_buffer[..total_len][window.response.clone()]; if let Err(err) = writer.write_all(chunk).await { - queue_packed_next_response_input(io_buffer, conn, pool, next_response_input, total_len)?; + window.queue_next_response_input(&io_buffer[..total_len], conn)?; return Err( crate::session::response_transfer::ResponseTransferError::ClientDisconnect(err), ); } let mut stats = ResponseWriteStats::default(); stats.add_chunk(chunk.len()); - queue_packed_next_response_input(io_buffer, conn, pool, next_response_input, total_len)?; + window.queue_next_response_input(&io_buffer[..total_len], conn)?; Ok(stats) } @@ -1217,7 +1191,7 @@ impl<'a> StreamingResponse<'a> { ResponseWindow::Complete(chunk) => { let total_len = self.io_buffer.initialized(); stats += chunk - .write_from(writer, self.io_buffer, self.conn, self.pool, total_len) + .write_from(writer, self.io_buffer, self.conn, total_len) .await?; return Ok(stats); } @@ -2458,7 +2432,7 @@ mod tests { let mut writer = Vec::new(); let written = framed - .write_from(&mut writer, &mut io_buffer, &mut conn, &pool, chunk.len()) + .write_from(&mut writer, &mut io_buffer, &mut conn, chunk.len()) .await .expect("complete response should write"); @@ -2515,7 +2489,7 @@ mod tests { let mut writer = FailingWriter; let err = framed - .write_from(&mut writer, &mut io_buffer, &mut conn, &pool, chunk.len()) + .write_from(&mut writer, &mut io_buffer, &mut conn, chunk.len()) .await; assert!(matches!( @@ -2581,13 +2555,7 @@ mod tests { let mut writer = FailingWriter; let err = complete - .write_from( - &mut writer, - &mut io_buffer, - &mut conn, - &pool, - continuation.len(), - ) + .write_from(&mut writer, &mut io_buffer, &mut conn, continuation.len()) .await; assert!(matches!( From f39e8fec74bb22c9e5278e3941108a5479039769 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 11:56:06 -0600 Subject: [PATCH 15/23] Preserve pooled packed response ownership --- src/session/multiline_framing.rs | 41 +++++++++++++++++++++++++------- 1 file changed, 33 insertions(+), 8 deletions(-) diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 8a3b697e..40b2e0be 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -136,13 +136,16 @@ impl CompleteResponseWindow { writer: &mut W, io_buffer: &mut crate::pool::PooledBuffer, conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, total_len: usize, ) -> Result where W: AsyncWrite + Unpin, { - write_response_chunk_preserving_suffix_on_error(writer, io_buffer, conn, total_len, self) - .await + write_response_chunk_preserving_suffix_on_error( + writer, io_buffer, conn, pool, total_len, self, + ) + .await } fn extend_capture_from(&self, source: &[u8], capture: &mut crate::pool::PooledBuffer) { @@ -171,6 +174,21 @@ impl CompleteResponseWindow { Ok(()) } + fn queue_pooled_next_response_input( + &self, + io_buffer: &mut crate::pool::PooledBuffer, + conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, + total_len: usize, + ) -> Result<(), crate::session::response_transfer::ResponseTransferError> { + if self.next_response_input.start < total_len { + let old = std::mem::replace(io_buffer, pool.acquire()); + conn.queue_pooled_pending_bytes_first(old, self.next_response_input.clone()) + .map_err(crate::session::response_transfer::ResponseTransferError::Io)?; + } + Ok(()) + } + fn push_from_buffer( &self, io_buffer: &mut crate::pool::PooledBuffer, @@ -221,6 +239,7 @@ async fn write_response_chunk_preserving_suffix_on_error( writer: &mut W, io_buffer: &mut crate::pool::PooledBuffer, conn: &mut crate::stream::ConnectionStream, + pool: &crate::pool::BufferPool, total_len: usize, window: &CompleteResponseWindow, ) -> Result @@ -229,14 +248,14 @@ where { let chunk = &io_buffer[..total_len][window.response.clone()]; if let Err(err) = writer.write_all(chunk).await { - window.queue_next_response_input(&io_buffer[..total_len], conn)?; + window.queue_pooled_next_response_input(io_buffer, conn, pool, total_len)?; return Err( crate::session::response_transfer::ResponseTransferError::ClientDisconnect(err), ); } let mut stats = ResponseWriteStats::default(); stats.add_chunk(chunk.len()); - window.queue_next_response_input(&io_buffer[..total_len], conn)?; + window.queue_pooled_next_response_input(io_buffer, conn, pool, total_len)?; Ok(stats) } @@ -1191,7 +1210,7 @@ impl<'a> StreamingResponse<'a> { ResponseWindow::Complete(chunk) => { let total_len = self.io_buffer.initialized(); stats += chunk - .write_from(writer, self.io_buffer, self.conn, total_len) + .write_from(writer, self.io_buffer, self.conn, self.pool, total_len) .await?; return Ok(stats); } @@ -2432,7 +2451,7 @@ mod tests { let mut writer = Vec::new(); let written = framed - .write_from(&mut writer, &mut io_buffer, &mut conn, chunk.len()) + .write_from(&mut writer, &mut io_buffer, &mut conn, &pool, chunk.len()) .await .expect("complete response should write"); @@ -2489,7 +2508,7 @@ mod tests { let mut writer = FailingWriter; let err = framed - .write_from(&mut writer, &mut io_buffer, &mut conn, chunk.len()) + .write_from(&mut writer, &mut io_buffer, &mut conn, &pool, chunk.len()) .await; assert!(matches!( @@ -2555,7 +2574,13 @@ mod tests { let mut writer = FailingWriter; let err = complete - .write_from(&mut writer, &mut io_buffer, &mut conn, continuation.len()) + .write_from( + &mut writer, + &mut io_buffer, + &mut conn, + &pool, + continuation.len(), + ) .await; assert!(matches!( From 1c223eddbf2298aa25955ddcffb34b811e510012 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 11:59:57 -0600 Subject: [PATCH 16/23] Isolate response metrics in parallel tests --- src/pool/buffer.rs | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index ae2d070d..a4c5a795 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -629,14 +629,26 @@ pub struct HotPathAllocationMetricsSnapshot { } fn response_write_metrics_enabled() -> bool { + #[cfg(test)] + { + return TEST_RESPONSE_WRITE_METRICS_ENABLED.with(std::cell::Cell::get); + } + + #[cfg(not(test))] response_write_metrics_enabled_flag().load(Ordering::Relaxed) } #[cfg(test)] pub(crate) fn set_response_write_metrics_enabled(enabled: bool) { - response_write_metrics_enabled_flag().store(enabled, Ordering::Relaxed); + TEST_RESPONSE_WRITE_METRICS_ENABLED.with(|value| value.set(enabled)); +} + +#[cfg(test)] +thread_local! { + static TEST_RESPONSE_WRITE_METRICS_ENABLED: std::cell::Cell = const { std::cell::Cell::new(false) }; } +#[cfg(not(test))] fn response_write_metrics_enabled_flag() -> &'static AtomicBool { static ENABLED: OnceLock = OnceLock::new(); ENABLED.get_or_init(|| { From dca05538ba11f0c9d6b12178023d1b35b6b4aee6 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 12:22:11 -0600 Subject: [PATCH 17/23] Characterize proxy article fixtures --- src/protocol/article/mod.rs | 60 +++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/src/protocol/article/mod.rs b/src/protocol/article/mod.rs index dc81c859..e2a437e9 100644 --- a/src/protocol/article/mod.rs +++ b/src/protocol/article/mod.rs @@ -363,4 +363,64 @@ mod tests { assert!(decoded.is_none()); } + + #[test] + fn compatibility_fixture_matrix_records_proxy_wire_contract() { + let fixtures = [ + ( + b"220 0 \r\nSubject: fixture\r\n\r\nbody\r\n".as_slice(), + Some(0), + true, + true, + ), + ( + b"221 0 \r\nSubject: fixture\r\nFrom: test@example.com\r\n" + .as_slice(), + Some(0), + true, + false, + ), + ( + b"222 0 \r\nbody\r\n".as_slice(), + Some(0), + false, + true, + ), + ( + b"223 0 \r\n".as_slice(), + Some(0), + false, + false, + ), + ( + b"222 0 \r\n".as_slice(), + Some(0), + false, + true, + ), + ]; + + for (wire, article_number, has_headers, has_body) in fixtures { + let article = Article::parse(wire, false).expect("proxy fixture remains accepted"); + assert_eq!(article.article_number, article_number); + assert_eq!(article.headers.is_some(), has_headers); + assert_eq!(article.body.is_some(), has_body); + } + } + + #[test] + fn compatibility_fixture_matrix_keeps_lenient_article_number_behavior() { + for number in [b"not-a-number".as_slice(), b"18446744073709551616"] { + let wire = [ + b"220 ".as_slice(), + number, + b" \r\nSubject: fixture\r\n\r\nbody\r\n", + ] + .concat(); + + let article = Article::parse(&wire, false).expect("proxy parser is permissive here"); + assert_eq!(article.article_number, None); + assert_eq!(article.message_id.as_str(), ""); + } + } } From 697e7215f8f93ea20c38780e27bac199ef4368c8 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 12:22:15 -0600 Subject: [PATCH 18/23] Handle NNTP dot stuffing before yEnc decode --- benches/yenc_decoding.rs | 13 +++++++++++++ src/protocol/article/yenc.rs | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+) diff --git a/benches/yenc_decoding.rs b/benches/yenc_decoding.rs index c5e1245f..38bcea68 100644 --- a/benches/yenc_decoding.rs +++ b/benches/yenc_decoding.rs @@ -4,6 +4,7 @@ //! we're getting expected performance from it. use divan::{Bencher, black_box}; +use nntp_proxy::protocol::article::yenc::decode_yenc_line; fn main() { divan::main(); @@ -150,4 +151,16 @@ mod yenc_crate_benches { let data = b"r\x8f\x96\x96\x99VJ\xa3o\x98\x8dK\r\n"; bencher.bench(|| black_box(yenc::decode_buffer(black_box(data)).unwrap())); } + + #[divan::bench(sample_count = 1000)] + fn yenc_decode_wire_stuffed_line(bencher: Bencher) { + let data = b"..r\x8f\x96\x96\x99VJ\xa3o\x98\x8dK\r\n"; + bencher.bench(|| black_box(decode_yenc_line(black_box(data)))); + } + + #[divan::bench(sample_count = 1000)] + fn yenc_decode_unstuffed_line(bencher: Bencher) { + let data = b"r\x8f\x96\x96\x99VJ\xa3o\x98\x8dK\r\n"; + bencher.bench(|| black_box(decode_yenc_line(black_box(data)))); + } } diff --git a/src/protocol/article/yenc.rs b/src/protocol/article/yenc.rs index 1acbcd24..d25eb114 100644 --- a/src/protocol/article/yenc.rs +++ b/src/protocol/article/yenc.rs @@ -18,6 +18,14 @@ use std::io::{BufRead, BufReader}; #[inline] #[must_use] pub fn decode_yenc_line(input: &[u8]) -> Vec { + // NNTP dot-stuffs a data line that begins with a dot. Remove only the + // protocol-added dot; a second leading dot is part of the yEnc payload. + let input = if input.starts_with(b"..") { + &input[1..] + } else { + input + }; + // yenc crate returns Result but decoding never actually fails // (it just returns empty/partial on malformed input) yenc::decode_buffer(input).unwrap_or_default() @@ -256,6 +264,31 @@ mod tests { assert!(result.is_ok()); } + #[test] + fn decode_yenc_line_removes_one_wire_stuffed_leading_dot() { + let once_stuffed = decode_yenc_line(b"..encoded"); + let twice_stuffed = decode_yenc_line(b"...encoded"); + + assert_eq!(once_stuffed, yenc::decode_buffer(b".encoded").unwrap()); + assert_eq!(twice_stuffed, yenc::decode_buffer(b"..encoded").unwrap()); + } + + #[test] + fn stuffed_leading_dots_contribute_to_the_yenc_checksum_once() { + let mut actual = crc32fast::Hasher::new(); + decode_and_checksum(b"..\r\n", &mut actual); + + let mut expected = crc32fast::Hasher::new(); + expected.update(&decode_yenc_line(b"..")); + assert_eq!(actual.finalize(), expected.finalize()); + + let mut two_bytes = crc32fast::Hasher::new(); + decode_and_checksum(b"...\r\n", &mut two_bytes); + let mut expected_two_bytes = crc32fast::Hasher::new(); + expected_two_bytes.update(&decode_yenc_line(b"...")); + assert_eq!(two_bytes.finalize(), expected_two_bytes.finalize()); + } + #[test] fn test_missing_yend() { let data = b"=ybegin part=1 size=1000 name=test.bin\r\n\ From 7cb5a269e127cab5c47f776b6301afb9bbb68f90 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 12:25:22 -0600 Subject: [PATCH 19/23] Keep yEnc benchmark on public module path --- benches/yenc_decoding.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/benches/yenc_decoding.rs b/benches/yenc_decoding.rs index 38bcea68..20c9276d 100644 --- a/benches/yenc_decoding.rs +++ b/benches/yenc_decoding.rs @@ -4,7 +4,7 @@ //! we're getting expected performance from it. use divan::{Bencher, black_box}; -use nntp_proxy::protocol::article::yenc::decode_yenc_line; +use nntp_proxy::protocol::yenc::decode_yenc_line; fn main() { divan::main(); @@ -40,7 +40,9 @@ fn generate_yenc_with_escapes(length: usize, escape_freq: usize) -> Vec { } mod yenc_crate_benches { - use super::{Bencher, black_box, generate_yenc_data, generate_yenc_with_escapes}; + use super::{ + Bencher, black_box, decode_yenc_line, generate_yenc_data, generate_yenc_with_escapes, + }; #[divan::bench(sample_count = 1000)] fn yenc_decode_128_bytes(bencher: Bencher) { From 22d63591e5d024d2722a30ad65487e216d1dcf04 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 12:28:20 -0600 Subject: [PATCH 20/23] Characterize proxy wire sections --- src/protocol/article/mod.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/protocol/article/mod.rs b/src/protocol/article/mod.rs index e2a437e9..ab201947 100644 --- a/src/protocol/article/mod.rs +++ b/src/protocol/article/mod.rs @@ -423,4 +423,24 @@ mod tests { assert_eq!(article.message_id.as_str(), ""); } } + + #[test] + fn compatibility_fixture_matrix_preserves_proxy_wire_sections() { + let folded = b"220 0 \r\nSubject: first\r\n second\r\n\r\nbody\r\n"; + let folded_article = Article::parse(folded, false).unwrap(); + assert_eq!( + folded_article.headers.unwrap().get("Subject"), + Some(&b"first"[..]) + ); + + let stuffed = b"222 0 \r\n..wire-dot\r\n"; + let stuffed_article = Article::parse(stuffed, false).unwrap(); + assert_eq!(stuffed_article.body, Some(&b"..wire-dot\r\n"[..])); + + let binary = b"222 0 \r\nbinary\0body\r\n"; + assert!(Article::parse(binary, false).is_ok()); + + let bare_lf = b"222 0 \r\nbody\nnext\r\n"; + assert!(Article::parse(bare_lf, false).is_ok()); + } } From 7573d02969a55881c2b5609de6143aabe2ec4c7e Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 12:36:28 -0600 Subject: [PATCH 21/23] Fix test-only metrics imports --- src/pool/buffer.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index a4c5a795..eb392acd 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -7,7 +7,9 @@ use std::ops::{Deref, Range}; use std::pin::Pin; use std::sync::Arc; use std::sync::OnceLock; -use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}; +#[cfg(not(test))] +use std::sync::atomic::AtomicBool; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use std::time::Instant; use tokio::io::{AsyncRead, AsyncWriteExt, ReadBuf}; use tracing::{debug, info, warn}; @@ -631,7 +633,7 @@ pub struct HotPathAllocationMetricsSnapshot { fn response_write_metrics_enabled() -> bool { #[cfg(test)] { - return TEST_RESPONSE_WRITE_METRICS_ENABLED.with(std::cell::Cell::get); + TEST_RESPONSE_WRITE_METRICS_ENABLED.with(std::cell::Cell::get) } #[cfg(not(test))] From 203ce6cc5f4671585387e52cba785780fb2d282b Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 13:01:46 -0600 Subject: [PATCH 22/23] Carry multiline payload boundaries through framing --- src/pool/buffer.rs | 14 --- src/session/multiline_framing.rs | 149 ++++++++++++++++++++++--------- 2 files changed, 106 insertions(+), 57 deletions(-) diff --git a/src/pool/buffer.rs b/src/pool/buffer.rs index eb392acd..69feccd9 100644 --- a/src/pool/buffer.rs +++ b/src/pool/buffer.rs @@ -188,11 +188,6 @@ impl BufferStorage { self.visible_from = start; } - fn truncate_visible(&mut self, len: usize) { - assert!(len <= self.initialized_len()); - self.bytes.truncate(self.visible_from + len); - } - fn compact_visible(&mut self) { if self.visible_from == 0 { return; @@ -395,15 +390,6 @@ impl PooledBuffer { self.buffer.clear(); } - /// Truncate initialized bytes without moving the backing allocation. - pub(crate) fn truncate_initialized(&mut self, len: usize) { - assert!( - len <= self.initialized(), - "truncate length exceeds initialized bytes" - ); - self.buffer.truncate_visible(len); - } - /// Append data to the buffer (accumulator mode) /// /// Used when `PooledBuffer` is acquired from capture pool for accumulating diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index 40b2e0be..df02218e 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -127,6 +127,13 @@ impl CompleteMultilinePayloadSplit { #[derive(Debug, PartialEq, Eq)] struct CompleteResponseWindow { response: Range, + /// The response bytes intended for isolated capture. + /// + /// Multiline framing proves the terminator while producing this window, + /// so capture can use the already-known payload boundary without scanning + /// the completed bytes a second time. Single-line responses use their + /// complete response range here. + capture: Range, next_response_input: Range, } @@ -152,6 +159,10 @@ impl CompleteResponseWindow { capture.extend_from_slice(&source[self.response.clone()]); } + fn extend_payload_capture_from(&self, source: &[u8], capture: &mut crate::pool::PooledBuffer) { + capture.extend_from_slice(&source[self.capture.clone()]); + } + fn push_isolated_buffer_to( &self, response: &mut crate::pool::ChunkedResponse, @@ -295,9 +306,15 @@ impl WindowEnd { } } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct CompleteChunk { + consumed: ChunkConsumed, + payload_end: ChunkConsumed, +} + #[derive(Debug, PartialEq, Eq)] enum ChunkProgress { - Complete(ChunkConsumed), + Complete(CompleteChunk), Incomplete, } @@ -306,10 +323,11 @@ impl ChunkProgress { // own buffer-bound append result, never from an application caller. fn in_window(self, origin: WindowEnd, window_len: usize) -> ResponseWindow { match self { - Self::Complete(consumed) => { - let end = origin.after_chunk(consumed); + Self::Complete(complete) => { + let end = origin.after_chunk(complete.consumed); ResponseWindow::Complete(CompleteResponseWindow { response: 0..end.0, + capture: 0..origin.after_chunk(complete.payload_end).0, next_response_input: end.0..window_len, }) } @@ -663,11 +681,8 @@ impl ClassifiedResponse { let mut capture = pool.acquire_capture(); IsolatedMultilineResponse::begin(conn.stream_mut(), &mut self.buffer) .map_err(isolated_multiline_error)? - .capture_into(&mut capture) + .capture_payload_into(&mut capture) .await?; - // Framing owns the already-proven terminator. The public - // article parser receives the status line and payload only. - strip_captured_terminator(&mut capture)?; conn.mark_response_complete(); Ok(capture) } @@ -808,17 +823,6 @@ impl ClassifiedResponse { } } -fn strip_captured_terminator(capture: &mut crate::pool::PooledBuffer) -> anyhow::Result<()> { - let len = capture.initialized(); - anyhow::ensure!(len >= TERMINATOR.len(), "captured response is too short"); - anyhow::ensure!( - capture[len - TERMINATOR.len()..] == TERMINATOR[..], - "captured response does not end at the framed terminator" - ); - capture.truncate_initialized(len - TERMINATOR.len()); - Ok(()) -} - pub(crate) fn unpacked_single_line_response<'a>( request: &crate::protocol::RequestContext, bytes: &'a [u8], @@ -947,11 +951,30 @@ impl<'a> IsolatedMultilineResponse<'a> { }) } - async fn capture_into(mut self, capture: &mut crate::pool::PooledBuffer) -> anyhow::Result<()> { + #[cfg(test)] + async fn capture_into(self, capture: &mut crate::pool::PooledBuffer) -> anyhow::Result<()> { + self.capture_into_mode::(capture).await + } + + async fn capture_payload_into( + self, + capture: &mut crate::pool::PooledBuffer, + ) -> anyhow::Result<()> { + self.capture_into_mode::(capture).await + } + + async fn capture_into_mode( + mut self, + capture: &mut crate::pool::PooledBuffer, + ) -> anyhow::Result<()> { loop { match &self.frame { ResponseWindow::Complete(chunk) => { - chunk.extend_capture_from(self.io_buffer, capture); + if INCLUDE_TERMINATOR { + chunk.extend_capture_from(self.io_buffer, capture); + } else { + chunk.extend_payload_capture_from(self.io_buffer, capture); + } return Ok(()); } ResponseWindow::Incomplete(chunk) => { @@ -1087,6 +1110,7 @@ impl<'a> StreamingResponse<'a> { ResponseFrame::SingleLine { framed, .. } => ( ResponseShape::SingleLine, ResponseWindow::Complete(CompleteResponseWindow { + capture: framed.response.clone(), response: framed.response, next_response_input: framed.next_response_input, }), @@ -1401,9 +1425,32 @@ impl MultilineFramer { chunk: &[u8], suffix_policy: PackedPendingBytesPolicy, ) -> Result { + if let Some(end) = self.find_spanning_terminator(chunk) { + if end == chunk.len() + || matches!(suffix_policy, PackedPendingBytesPolicy::AllowIfStatusPrefix) + && plausible_status_prefix(&chunk[end..]) + { + return Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(end), + payload_end: ChunkConsumed(0), + })); + } + if matches!(suffix_policy, PackedPendingBytesPolicy::Reject) { + return Err(FramingError::UnexpectedTrailingResponseBytes); + } + } for end in self.terminator_ends(chunk) { + if self.find_spanning_terminator(chunk) == Some(end) { + continue; + } if end == chunk.len() { - return Ok(ChunkProgress::Complete(ChunkConsumed(end))); + return Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(end), + payload_end: ChunkConsumed( + end.checked_sub(TERMINATOR.len()) + .expect("complete multiline response includes its terminator"), + ), + })); } match suffix_policy { @@ -1413,7 +1460,13 @@ impl MultilineFramer { PackedPendingBytesPolicy::AllowIfStatusPrefix if plausible_status_prefix(&chunk[end..]) => { - return Ok(ChunkProgress::Complete(ChunkConsumed(end))); + return Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(end), + payload_end: ChunkConsumed( + end.checked_sub(TERMINATOR.len()) + .expect("complete multiline response includes its terminator"), + ), + })); } PackedPendingBytesPolicy::AllowIfStatusPrefix => {} } @@ -1557,8 +1610,8 @@ impl PendingRequestFrame { match framer .split_chunk(&chunk[end..], PackedPendingBytesPolicy::AllowIfStatusPrefix) { - Ok(ChunkProgress::Complete(consumed)) => Some(FramedResponseForRequest { - response: offset..WindowEnd(end).after_chunk(consumed).0, + Ok(ChunkProgress::Complete(complete)) => Some(FramedResponseForRequest { + response: offset..WindowEnd(end).after_chunk(complete.consumed).0, }), Ok(ChunkProgress::Incomplete) => { self.state = PendingRequestFrameState::ReadingMultiline { framer }; @@ -1574,8 +1627,8 @@ impl PendingRequestFrame { &chunk[offset..], PackedPendingBytesPolicy::AllowIfStatusPrefix, ) { - Ok(ChunkProgress::Complete(consumed)) => Some(FramedResponseForRequest { - response: offset..WindowEnd(offset).after_chunk(consumed).0, + Ok(ChunkProgress::Complete(complete)) => Some(FramedResponseForRequest { + response: offset..WindowEnd(offset).after_chunk(complete.consumed).0, }), Ok(ChunkProgress::Incomplete) => None, Err(_) => Some(FramedResponseForRequest { @@ -1642,8 +1695,8 @@ fn complete_multiline_payload_split(payload: &[u8]) -> Option { - let response_end = consumed.0; + Ok(ChunkProgress::Complete(complete)) => { + let response_end = complete.consumed.0; let body_end = response_end.checked_sub(TERMINATOR.len())?; let terminator_start = response_end.checked_sub(3)?; Some(CompleteMultilinePayloadSplit::new( @@ -1710,10 +1763,10 @@ pub fn benchmark_incremental_multiline_frame(response: &[u8], chunk_size: usize) let chunk_size = chunk_size.max(1); for chunk in response.chunks(chunk_size) { - if let Ok(ChunkProgress::Complete(consumed)) = + if let Ok(ChunkProgress::Complete(complete)) = framer.split_chunk(chunk, PackedPendingBytesPolicy::AllowIfStatusPrefix) { - return consumed.0; + return complete.consumed.0; } } @@ -1730,11 +1783,11 @@ pub fn benchmark_stateless_multiline_frame(response: &[u8], chunk_size: usize) - for chunk in response.chunks(chunk_size) { received += chunk.len(); let mut framer = MultilineFramer::default(); - if let Ok(ChunkProgress::Complete(consumed)) = framer.split_chunk( + if let Ok(ChunkProgress::Complete(complete)) = framer.split_chunk( &response[..received], PackedPendingBytesPolicy::AllowIfStatusPrefix, ) { - return consumed.0; + return complete.consumed.0; } } @@ -2442,6 +2495,7 @@ mod tests { let response_len = b"220 article\r\nbody\r\n.\r\n".len(); let framed = CompleteResponseWindow { response: 0..response_len, + capture: 0..response_len, next_response_input: response_len..chunk.len(), }; let pool = make_pool(); @@ -2499,6 +2553,7 @@ mod tests { let response_len = b"220 article\r\nbody\r\n.\r\n".len(); let framed = CompleteResponseWindow { response: 0..response_len, + capture: 0..response_len, next_response_input: response_len..chunk.len(), }; let pool = make_pool(); @@ -3302,9 +3357,10 @@ mod tests { assert_eq!( split, - Ok(ChunkProgress::Complete(ChunkConsumed( - b"220 article\r\nbody\r\n.\r\n".len() - ))) + Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(b"220 article\r\nbody\r\n.\r\n".len()), + payload_end: ChunkConsumed(b"220 article\r\nbody".len()), + })) ); } @@ -3317,9 +3373,10 @@ mod tests { assert_eq!( split, - Ok(ChunkProgress::Complete(ChunkConsumed( - b"220 article\r\nbody\r\n.\r\n".len() - ))) + Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(b"220 article\r\nbody\r\n.\r\n".len()), + payload_end: ChunkConsumed(b"220 article\r\nbody".len()), + })) ); } @@ -3342,7 +3399,10 @@ mod tests { assert_eq!( split, - Ok(ChunkProgress::Complete(ChunkConsumed(chunk.len()))) + Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(chunk.len()), + payload_end: ChunkConsumed(chunk.len() - TERMINATOR.len()), + })) ); } @@ -3355,7 +3415,10 @@ mod tests { assert_eq!( split, - Ok(ChunkProgress::Complete(ChunkConsumed(b".\r\n".len()))) + Ok(ChunkProgress::Complete(CompleteChunk { + consumed: ChunkConsumed(b".\r\n".len()), + payload_end: ChunkConsumed(0), + })) ); } @@ -3373,7 +3436,7 @@ mod tests { ) .expect("first push should not reject a valid response"); let actual = match first { - ChunkProgress::Complete(consumed) => Some(consumed.0), + ChunkProgress::Complete(complete) => Some(complete.consumed.0), ChunkProgress::Incomplete => framer .split_chunk( &response[split..], @@ -3381,7 +3444,7 @@ mod tests { ) .ok() .and_then(|result| match result { - ChunkProgress::Complete(consumed) => Some(split + consumed.0), + ChunkProgress::Complete(complete) => Some(split + complete.consumed.0), ChunkProgress::Incomplete => None, }), }; @@ -3395,7 +3458,7 @@ mod tests { let ChunkProgress::Complete(rescanned) = rescanned_end else { panic!("rescan did not complete for split={split}"); }; - assert_eq!(actual, Some(rescanned.0), "split={split}"); + assert_eq!(actual, Some(rescanned.consumed.0), "split={split}"); } } From 7e22da6ebc39c5c58d1983b19fe8e65e82705099 Mon Sep 17 00:00:00 2001 From: Mika Cohen Date: Wed, 16 Sep 2026 13:16:12 -0600 Subject: [PATCH 23/23] Avoid duplicate spanning terminator scans --- src/session/multiline_framing.rs | 20 +++----------------- 1 file changed, 3 insertions(+), 17 deletions(-) diff --git a/src/session/multiline_framing.rs b/src/session/multiline_framing.rs index df02218e..3e2485d2 100644 --- a/src/session/multiline_framing.rs +++ b/src/session/multiline_framing.rs @@ -1439,10 +1439,7 @@ impl MultilineFramer { return Err(FramingError::UnexpectedTrailingResponseBytes); } } - for end in self.terminator_ends(chunk) { - if self.find_spanning_terminator(chunk) == Some(end) { - continue; - } + for end in terminator_ends_in_chunk(chunk) { if end == chunk.len() { return Ok(ChunkProgress::Complete(CompleteChunk { consumed: ChunkConsumed(end), @@ -1503,17 +1500,6 @@ impl MultilineFramer { .or_else(|| find_terminator_end_from(chunk, 0)) } - /// Return every complete terminator end offset touching the current chunk. - #[must_use] - fn terminator_ends(&self, chunk: &[u8]) -> smallvec::SmallVec<[usize; 2]> { - let mut ends = smallvec::SmallVec::new(); - if let Some(end) = self.find_spanning_terminator(chunk) { - ends.push(end); - } - ends.extend(terminator_ends(chunk)); - ends - } - /// Return the earliest terminator end offset, updating rolling state on miss. /// /// If no terminator is found, this appends `chunk` into the framer's rolling @@ -1740,7 +1726,7 @@ fn find_terminator_end(data: &[u8]) -> Option { } #[inline] -fn terminator_ends(data: &[u8]) -> impl Iterator + '_ { +fn terminator_ends_in_chunk(data: &[u8]) -> impl Iterator + '_ { TERMINATOR_FINDER .find_iter(data) .map(|found| found + TERMINATOR.len()) @@ -1799,7 +1785,7 @@ pub fn benchmark_stateless_multiline_frame(response: &[u8], chunk_size: usize) - fn find_terminator_end_from(data: &[u8], start: usize) -> Option { let data = data.get(start..)?; let mut first = None; - for end in terminator_ends(data) { + for end in terminator_ends_in_chunk(data) { first.get_or_insert(start + end); } first