diff --git a/src/py/client/streams/coalescing.rs b/src/py/client/streams/coalescing.rs new file mode 100644 index 0000000..29dbd18 --- /dev/null +++ b/src/py/client/streams/coalescing.rs @@ -0,0 +1,81 @@ +use super::constants::{MAX_READY_FRAME_COALESCE_COUNT, READY_FRAME_COALESCE_TARGET_BYTES}; +use bytes::{Bytes, BytesMut}; +use hyper::body::{Body, Incoming}; +use std::pin::Pin; +use std::task::{Context, Poll, Waker}; + +pub(super) enum ReadyFrameDrain { + Eof, + Error(String), + Stopped, +} + +enum ReadyBodyFrame { + Data(Bytes), + Eof, + Pending, + Error(String), +} + +pub(super) fn drain_ready_data_frames(chunk: &mut Bytes, body: &mut Incoming) -> ReadyFrameDrain { + drain_ready_data_frames_with(chunk, || poll_ready_body_frame(body)) +} + +fn poll_ready_body_frame(body: &mut Incoming) -> ReadyBodyFrame { + let waker = Waker::noop(); + let mut context = Context::from_waker(waker); + loop { + let frame = match Pin::new(&mut *body).poll_frame(&mut context) { + Poll::Ready(frame) => frame, + Poll::Pending => return ReadyBodyFrame::Pending, + }; + let Some(frame) = frame else { + return ReadyBodyFrame::Eof; + }; + let frame = match frame { + Ok(frame) => frame, + Err(err) => return ReadyBodyFrame::Error(err.to_string()), + }; + let Ok(data) = frame.into_data() else { + continue; + }; + return ReadyBodyFrame::Data(data); + } +} + +fn drain_ready_data_frames_with( + chunk: &mut Bytes, + mut poll_ready_body_frame: impl FnMut() -> ReadyBodyFrame, +) -> ReadyFrameDrain { + let mut coalesced_frames = 1; + let mut merged: Option = None; + let outcome = loop { + let chunk_len = merged.as_ref().map_or(chunk.len(), BytesMut::len); + if chunk_len >= READY_FRAME_COALESCE_TARGET_BYTES + || coalesced_frames >= MAX_READY_FRAME_COALESCE_COUNT + { + break ReadyFrameDrain::Stopped; + } + match poll_ready_body_frame() { + ReadyBodyFrame::Data(data) => { + let buffer = merged.get_or_insert_with(|| { + let mut buffer = BytesMut::with_capacity(chunk.len() + data.len()); + buffer.extend_from_slice(chunk); + buffer + }); + buffer.extend_from_slice(&data); + coalesced_frames += 1; + } + ReadyBodyFrame::Eof => break ReadyFrameDrain::Eof, + ReadyBodyFrame::Error(error) => break ReadyFrameDrain::Error(error), + ReadyBodyFrame::Pending => break ReadyFrameDrain::Stopped, + } + }; + if let Some(merged) = merged { + *chunk = merged.freeze(); + } + outcome +} + +#[cfg(test)] +mod tests; diff --git a/src/py/client/streams/coalescing/tests.rs b/src/py/client/streams/coalescing/tests.rs new file mode 100644 index 0000000..0cfaeb0 --- /dev/null +++ b/src/py/client/streams/coalescing/tests.rs @@ -0,0 +1,160 @@ +use super::{drain_ready_data_frames_with, ReadyBodyFrame, ReadyFrameDrain}; +use crate::py::client::streams::constants::{ + MAX_READY_FRAME_COALESCE_COUNT, READY_FRAME_COALESCE_TARGET_BYTES, +}; +use bytes::Bytes; +use pyo3::prelude::*; +use std::collections::VecDeque; +use std::sync::Once; + +fn initialize_python() { + static PYTHON: Once = Once::new(); + PYTHON.call_once(Python::initialize); +} + +#[test] +fn ready_frame_coalescing_merges_ready_data_until_eof() { + let mut chunk = Bytes::from_static(b"first"); + let mut frames = VecDeque::from([ + ReadyBodyFrame::Data(Bytes::from_static(b"second")), + ReadyBodyFrame::Data(Bytes::from_static(b"third")), + ReadyBodyFrame::Eof, + ]); + + let outcome = drain_ready_data_frames_with(&mut chunk, || frames.pop_front().unwrap()); + + assert!(matches!(outcome, ReadyFrameDrain::Eof)); + assert_eq!(chunk.as_ref(), b"firstsecondthird"); +} + +#[test] +fn ready_frame_coalescing_stops_on_pending_without_consuming_later_frames() { + let mut chunk = Bytes::from_static(b"first"); + let mut frames = VecDeque::from([ + ReadyBodyFrame::Data(Bytes::from_static(b"second")), + ReadyBodyFrame::Pending, + ReadyBodyFrame::Data(Bytes::from_static(b"third")), + ]); + + let outcome = drain_ready_data_frames_with(&mut chunk, || frames.pop_front().unwrap()); + + assert!(matches!(outcome, ReadyFrameDrain::Stopped)); + assert_eq!(chunk.as_ref(), b"firstsecond"); + assert!(matches!(frames.pop_front(), Some(ReadyBodyFrame::Data(_)))); +} + +#[test] +fn ready_frame_coalescing_returns_collected_data_before_deferred_error() { + let mut chunk = Bytes::from_static(b"first"); + let mut frames = VecDeque::from([ + ReadyBodyFrame::Data(Bytes::from_static(b"second")), + ReadyBodyFrame::Error("broken body".to_string()), + ]); + + let outcome = drain_ready_data_frames_with(&mut chunk, || frames.pop_front().unwrap()); + + assert!(matches!(outcome, ReadyFrameDrain::Error(error) if error == "broken body")); + assert_eq!(chunk.as_ref(), b"firstsecond"); +} + +#[test] +fn ready_frame_coalescing_respects_count_and_byte_targets() { + let mut chunk = Bytes::new(); + let mut poll_count = 0; + let outcome = drain_ready_data_frames_with(&mut chunk, || { + poll_count += 1; + ReadyBodyFrame::Data(Bytes::from_static(b"x")) + }); + + assert!(matches!(outcome, ReadyFrameDrain::Stopped)); + assert_eq!(poll_count, MAX_READY_FRAME_COALESCE_COUNT - 1); + assert_eq!(chunk.len(), MAX_READY_FRAME_COALESCE_COUNT - 1); + + let mut chunk = Bytes::from(vec![b'a'; READY_FRAME_COALESCE_TARGET_BYTES - 1]); + let mut poll_count = 0; + let outcome = drain_ready_data_frames_with(&mut chunk, || { + poll_count += 1; + ReadyBodyFrame::Data(Bytes::from_static(b"xy")) + }); + + assert!(matches!(outcome, ReadyFrameDrain::Stopped)); + assert_eq!(poll_count, 1); + assert_eq!(chunk.len(), READY_FRAME_COALESCE_TARGET_BYTES + 1); +} + +#[test] +#[ignore = "manual release-mode boundary measurement; no timing assertion"] +fn measure_streaming_boundary_conversion() { + use pyo3::types::PyBytes; + use std::hint::black_box; + use std::time::Instant; + + initialize_python(); + Python::attach(|py| { + for size in [4096, 65_536, 262_144] { + for frames in [1, 2] { + let data = Bytes::from(vec![b'x'; size / frames]); + let iterations = u32::try_from(256 * 1024 * 1024 / size).unwrap(); + for warmup in [true, false] { + let start = Instant::now(); + for _ in 0..iterations { + // Match the initial frame handoff in read_next_chunk. + let mut chunk = black_box(&data).clone(); + let mut remaining = frames - 1; + let outcome = drain_ready_data_frames_with(&mut chunk, || { + if remaining == 0 { + ReadyBodyFrame::Pending + } else { + remaining -= 1; + ReadyBodyFrame::Data(data.clone()) + } + }); + black_box(outcome); + black_box(PyBytes::new(py, &chunk)); + } + if !warmup { + println!( + "boundary size={size} frames={frames} iterations={iterations} ns_per_chunk={:.1}", + start.elapsed().as_secs_f64() * 1e9 / f64::from(iterations), + ); + } + } + } + } + }); +} + +#[test] +fn single_ready_frame_keeps_owned_storage_until_python_conversion() { + for frame in [ + ReadyBodyFrame::Pending, + ReadyBodyFrame::Eof, + ReadyBodyFrame::Error("broken body".to_string()), + ] { + let owner = Bytes::from(vec![b'x'; 4096]); + let mut chunk = owner.slice(10..20); + let original = chunk.as_ptr(); + let mut frame = Some(frame); + let _ = drain_ready_data_frames_with(&mut chunk, || frame.take().unwrap()); + drop(owner); + + assert_eq!(chunk.as_ptr(), original); + assert_eq!(chunk.as_ref(), &[b'x'; 10]); + } +} + +#[test] +fn ready_frame_at_byte_target_is_not_polled_or_copied() { + for size in [ + READY_FRAME_COALESCE_TARGET_BYTES, + READY_FRAME_COALESCE_TARGET_BYTES + 1, + ] { + let mut chunk = Bytes::from(vec![b'x'; size]); + let original = chunk.as_ptr(); + let outcome = drain_ready_data_frames_with(&mut chunk, || panic!("must not poll")); + + assert!(matches!(outcome, ReadyFrameDrain::Stopped)); + assert_eq!(chunk.as_ptr(), original); + assert_eq!(chunk.len(), size); + } +} diff --git a/src/py/client/streams/mod.rs b/src/py/client/streams/mod.rs index 0e2429f..a0d1aff 100644 --- a/src/py/client/streams/mod.rs +++ b/src/py/client/streams/mod.rs @@ -1,4 +1,5 @@ mod callback; +mod coalescing; mod constants; mod parts; mod read; diff --git a/src/py/client/streams/state.rs b/src/py/client/streams/state.rs index 376e3c7..a27fd57 100644 --- a/src/py/client/streams/state.rs +++ b/src/py/client/streams/state.rs @@ -1,4 +1,4 @@ -use super::constants::{MAX_READY_FRAME_COALESCE_COUNT, READY_FRAME_COALESCE_TARGET_BYTES}; +use super::coalescing::{drain_ready_data_frames, ReadyFrameDrain}; use super::read::next_stream_body_frame; use super::registry::StreamRegistry; use crate::core::client::{ConnectionAbortReason, ConnectionUseGuard}; @@ -10,14 +10,11 @@ use crate::py::client::async_requests::RequestCompletion; use crate::py::client::future::cancel_python_future; use crate::py::client::lifecycle::ResponseBodyLifecycle; use crate::py::client::timeout_diagnostics::response_body_transport_error; -use bytes::{Bytes, BytesMut}; -use hyper::body::Body; +use bytes::Bytes; use hyper::body::Incoming; use pyo3::prelude::*; use pyo3::types::PyAny; -use std::pin::Pin; use std::sync::{Arc, Mutex, MutexGuard}; -use std::task::{Context, Poll, Waker}; use std::time::Duration; use tokio::task::AbortHandle; @@ -82,19 +79,6 @@ pub(super) struct ActiveStreamRead { notification: StreamReadNotification, } -enum ReadyBodyFrame { - Data(Bytes), - Eof, - Pending, - Error(String), -} - -enum ReadyFrameDrain { - Eof, - Error(String), - Stopped, -} - #[derive(Clone, Copy)] pub(super) enum ReadyFrameCoalescing { Disabled, @@ -395,7 +379,7 @@ impl StreamReadGuard { let mut chunk = data; if self.ready_frame_coalescing.is_enabled() { - match self.drain_ready_data_frames(&mut chunk) { + match drain_ready_data_frames(&mut chunk, self.body_mut()) { ReadyFrameDrain::Eof => { self.finish_success_from_read(); return Ok(Some(chunk)); @@ -412,32 +396,6 @@ impl StreamReadGuard { } } - fn drain_ready_data_frames(&mut self, chunk: &mut Bytes) -> ReadyFrameDrain { - drain_ready_data_frames_with(chunk, || self.poll_ready_body_frame()) - } - - fn poll_ready_body_frame(&mut self) -> ReadyBodyFrame { - let waker = Waker::noop(); - let mut context = Context::from_waker(waker); - loop { - let frame = match Pin::new(self.body_mut()).poll_frame(&mut context) { - Poll::Ready(frame) => frame, - Poll::Pending => return ReadyBodyFrame::Pending, - }; - let Some(frame) = frame else { - return ReadyBodyFrame::Eof; - }; - let frame = match frame { - Ok(frame) => frame, - Err(err) => return ReadyBodyFrame::Error(err.to_string()), - }; - let Ok(data) = frame.into_data() else { - continue; - }; - return ReadyBodyFrame::Data(data); - } - } - fn body_mut(&mut self) -> &mut Incoming { self.body .as_mut() @@ -494,40 +452,6 @@ impl Drop for StreamReadGuard { } } -fn drain_ready_data_frames_with( - chunk: &mut Bytes, - mut poll_ready_body_frame: impl FnMut() -> ReadyBodyFrame, -) -> ReadyFrameDrain { - let mut coalesced_frames = 1; - let mut merged: Option = None; - let outcome = loop { - let chunk_len = merged.as_ref().map_or(chunk.len(), BytesMut::len); - if chunk_len >= READY_FRAME_COALESCE_TARGET_BYTES - || coalesced_frames >= MAX_READY_FRAME_COALESCE_COUNT - { - break ReadyFrameDrain::Stopped; - } - match poll_ready_body_frame() { - ReadyBodyFrame::Data(data) => { - let buffer = merged.get_or_insert_with(|| { - let mut buffer = BytesMut::with_capacity(chunk.len() + data.len()); - buffer.extend_from_slice(chunk); - buffer - }); - buffer.extend_from_slice(&data); - coalesced_frames += 1; - } - ReadyBodyFrame::Eof => break ReadyFrameDrain::Eof, - ReadyBodyFrame::Error(error) => break ReadyFrameDrain::Error(error), - ReadyBodyFrame::Pending => break ReadyFrameDrain::Stopped, - } - }; - if let Some(merged) = merged { - *chunk = merged.freeze(); - } - outcome -} - #[derive(Clone, Copy)] enum StreamFinish { Success, diff --git a/src/py/client/streams/state_tests.rs b/src/py/client/streams/state_tests.rs index fb4a95c..f23764c 100644 --- a/src/py/client/streams/state_tests.rs +++ b/src/py/client/streams/state_tests.rs @@ -1,18 +1,10 @@ -use super::{ - drain_ready_data_frames_with, ActiveStreamRead, ReadyBodyFrame, ReadyFrameDrain, StreamState, - StreamStateFields, StreamStateInner, -}; +use super::{ActiveStreamRead, StreamState, StreamStateFields, StreamStateInner}; use crate::core::metrics::{Metrics, ResponseBodyLifecycleOutcome}; use crate::py::client::async_requests::RequestCompletion; use crate::py::client::lifecycle::ResponseBodyLifecycle; -use crate::py::client::streams::constants::{ - MAX_READY_FRAME_COALESCE_COUNT, READY_FRAME_COALESCE_TARGET_BYTES, -}; use crate::py::client::streams::StreamRegistry; -use bytes::Bytes; use pyo3::prelude::*; use pyo3::types::PyAny; -use std::collections::VecDeque; use std::sync::{Arc, Mutex, Once}; use std::time::Duration; use tokio::runtime::{Builder, Runtime}; @@ -168,150 +160,3 @@ fn cancel_after_success_does_not_reclassify_stream() { assert_eq!(snapshot.response_body_aborted, 0); assert_eq!(snapshot.response_body_closed, 1); } - -#[test] -fn ready_frame_coalescing_merges_ready_data_until_eof() { - let mut chunk = Bytes::from_static(b"first"); - let mut frames = VecDeque::from([ - ReadyBodyFrame::Data(Bytes::from_static(b"second")), - ReadyBodyFrame::Data(Bytes::from_static(b"third")), - ReadyBodyFrame::Eof, - ]); - - let outcome = drain_ready_data_frames_with(&mut chunk, || frames.pop_front().unwrap()); - - assert!(matches!(outcome, ReadyFrameDrain::Eof)); - assert_eq!(chunk.as_ref(), b"firstsecondthird"); -} - -#[test] -fn ready_frame_coalescing_stops_on_pending_without_consuming_later_frames() { - let mut chunk = Bytes::from_static(b"first"); - let mut frames = VecDeque::from([ - ReadyBodyFrame::Data(Bytes::from_static(b"second")), - ReadyBodyFrame::Pending, - ReadyBodyFrame::Data(Bytes::from_static(b"third")), - ]); - - let outcome = drain_ready_data_frames_with(&mut chunk, || frames.pop_front().unwrap()); - - assert!(matches!(outcome, ReadyFrameDrain::Stopped)); - assert_eq!(chunk.as_ref(), b"firstsecond"); - assert!(matches!(frames.pop_front(), Some(ReadyBodyFrame::Data(_)))); -} - -#[test] -fn ready_frame_coalescing_returns_collected_data_before_deferred_error() { - let mut chunk = Bytes::from_static(b"first"); - let mut frames = VecDeque::from([ - ReadyBodyFrame::Data(Bytes::from_static(b"second")), - ReadyBodyFrame::Error("broken body".to_string()), - ]); - - let outcome = drain_ready_data_frames_with(&mut chunk, || frames.pop_front().unwrap()); - - assert!(matches!(outcome, ReadyFrameDrain::Error(error) if error == "broken body")); - assert_eq!(chunk.as_ref(), b"firstsecond"); -} - -#[test] -fn ready_frame_coalescing_respects_count_and_byte_targets() { - let mut chunk = Bytes::new(); - let mut poll_count = 0; - let outcome = drain_ready_data_frames_with(&mut chunk, || { - poll_count += 1; - ReadyBodyFrame::Data(Bytes::from_static(b"x")) - }); - - assert!(matches!(outcome, ReadyFrameDrain::Stopped)); - assert_eq!(poll_count, MAX_READY_FRAME_COALESCE_COUNT - 1); - assert_eq!(chunk.len(), MAX_READY_FRAME_COALESCE_COUNT - 1); - - let mut chunk = Bytes::from(vec![b'a'; READY_FRAME_COALESCE_TARGET_BYTES - 1]); - let mut poll_count = 0; - let outcome = drain_ready_data_frames_with(&mut chunk, || { - poll_count += 1; - ReadyBodyFrame::Data(Bytes::from_static(b"xy")) - }); - - assert!(matches!(outcome, ReadyFrameDrain::Stopped)); - assert_eq!(poll_count, 1); - assert_eq!(chunk.len(), READY_FRAME_COALESCE_TARGET_BYTES + 1); -} - -#[test] -#[ignore = "manual release-mode boundary measurement; no timing assertion"] -fn measure_streaming_boundary_conversion() { - use pyo3::types::PyBytes; - use std::hint::black_box; - use std::time::Instant; - - initialize_python(); - Python::attach(|py| { - for size in [4096, 65_536, 262_144] { - for frames in [1, 2] { - let data = Bytes::from(vec![b'x'; size / frames]); - let iterations = u32::try_from(256 * 1024 * 1024 / size).unwrap(); - for warmup in [true, false] { - let start = Instant::now(); - for _ in 0..iterations { - // Match the initial frame handoff in read_next_chunk. - let mut chunk = black_box(&data).clone(); - let mut remaining = frames - 1; - let outcome = drain_ready_data_frames_with(&mut chunk, || { - if remaining == 0 { - ReadyBodyFrame::Pending - } else { - remaining -= 1; - ReadyBodyFrame::Data(data.clone()) - } - }); - black_box(outcome); - black_box(PyBytes::new(py, &chunk)); - } - if !warmup { - println!( - "boundary size={size} frames={frames} iterations={iterations} ns_per_chunk={:.1}", - start.elapsed().as_secs_f64() * 1e9 / f64::from(iterations), - ); - } - } - } - } - }); -} - -#[test] -fn single_ready_frame_keeps_owned_storage_until_python_conversion() { - for frame in [ - ReadyBodyFrame::Pending, - ReadyBodyFrame::Eof, - ReadyBodyFrame::Error("broken body".to_string()), - ] { - let owner = Bytes::from(vec![b'x'; 4096]); - let mut chunk = owner.slice(10..20); - let original = chunk.as_ptr(); - let mut frame = Some(frame); - let _ = drain_ready_data_frames_with(&mut chunk, || frame.take().unwrap()); - drop(owner); - - assert_eq!(chunk.as_ptr(), original); - assert_eq!(chunk.as_ref(), &[b'x'; 10]); - } -} - -#[test] -fn ready_frame_at_byte_target_is_not_polled_or_copied() { - for size in [ - READY_FRAME_COALESCE_TARGET_BYTES, - READY_FRAME_COALESCE_TARGET_BYTES + 1, - ] { - let mut chunk = Bytes::from(vec![b'x'; size]); - let original = chunk.as_ptr(); - let outcome = drain_ready_data_frames_with(&mut chunk, || panic!("must not poll")); - - assert!(matches!(outcome, ReadyFrameDrain::Stopped)); - assert_eq!(chunk.as_ptr(), original); - assert_eq!(chunk.len(), size); - } -}