diff --git a/benches/end_to_end_proxy.rs b/benches/end_to_end_proxy.rs index e8ebea48..379c351c 100644 --- a/benches/end_to_end_proxy.rs +++ b/benches/end_to_end_proxy.rs @@ -148,17 +148,17 @@ struct BenchProxy { } impl BenchProxy { - async fn start(body_len: usize, cache: Option) -> Self { + async fn start(body_len: usize, cache: Option, mode: RoutingMode) -> Self { let backend_listener = bind_localhost().await; let backend_port = backend_listener.local_addr().unwrap().port(); spawn_backend(backend_listener, article_response(body_len).into()); let proxy_listener = bind_localhost().await; let proxy_addr = proxy_listener.local_addr().unwrap(); - let proxy = NntpProxy::new(bench_config(backend_port, cache), RoutingMode::PerCommand) + let proxy = NntpProxy::new(bench_config(backend_port, cache), mode) .await .unwrap(); - spawn_proxy(proxy_listener, proxy, RoutingMode::PerCommand); + spawn_proxy(proxy_listener, proxy, mode); let mut stream = TcpStream::connect(proxy_addr).await.unwrap(); stream.set_nodelay(true).unwrap(); @@ -172,7 +172,7 @@ impl BenchProxy { Self { stream, - response_buffer: vec![0; body_len + 512].into_boxed_slice(), + response_buffer: vec![0; body_len.saturating_mul(2) + 1024].into_boxed_slice(), response_len: article_response(body_len).len(), } } @@ -189,6 +189,19 @@ impl BenchProxy { ) .await } + + async fn article_window_roundtrip(&mut self) -> usize { + self.stream + .write_all(b"ARTICLE \r\nARTICLE \r\n") + .await + .unwrap(); + read_exact_response_into( + &mut self.stream, + &mut self.response_buffer, + self.response_len.saturating_mul(2), + ) + .await + } } async fn read_line(stream: &mut TcpStream, buffer: &mut Vec) -> usize { @@ -219,7 +232,7 @@ async fn read_exact_response_into( fn bench_roundtrip(bencher: Bencher, body_len: usize, cache: Option, warm_cache: bool) { let rt = Builder::new_current_thread().enable_all().build().unwrap(); - let mut proxy = rt.block_on(BenchProxy::start(body_len, cache)); + let mut proxy = rt.block_on(BenchProxy::start(body_len, cache, RoutingMode::PerCommand)); if warm_cache { let bytes = rt.block_on(proxy.article_roundtrip()); assert!(bytes > body_len); @@ -232,6 +245,25 @@ fn bench_roundtrip(bencher: Bencher, body_len: usize, cache: Option, warm }); } +fn bench_window(bencher: Bencher, body_len: usize, mode: RoutingMode) { + let rt = Builder::new_current_thread().enable_all().build().unwrap(); + let mut proxy = rt.block_on(BenchProxy::start(body_len, None, mode)); + bencher + .counter(divan::counter::BytesCount::new(body_len.saturating_mul(2))) + .bench_local(|| { + let bytes = rt.block_on(proxy.article_window_roundtrip()); + black_box(bytes) + }); +} + +fn bench_stateful_window(bencher: Bencher, body_len: usize) { + bench_window(bencher, body_len, RoutingMode::Stateful); +} + +fn bench_per_command_pair(bencher: Bencher, body_len: usize) { + bench_window(bencher, body_len, RoutingMode::PerCommand); +} + mod backend_roundtrip { use super::{ARTICLE_64K, ARTICLE_768K, ARTICLE_AVG_788K, ARTICLE_LARGE_1536K}; use super::{Bencher, bench_roundtrip}; @@ -316,3 +348,17 @@ mod cache_miss_roundtrip { ); } } + +mod stateful_upstream_window { + use super::{ARTICLE_64K, Bencher, bench_per_command_pair, bench_stateful_window}; + + #[divan::bench(sample_count = 50, sample_size = 5)] + fn article_pair_64k(bencher: Bencher) { + bench_stateful_window(bencher, ARTICLE_64K); + } + + #[divan::bench(sample_count = 50, sample_size = 5)] + fn per_command_pair_64k(bencher: Bencher) { + bench_per_command_pair(bencher, ARTICLE_64K); + } +} diff --git a/src/session/handlers/stateful.rs b/src/session/handlers/stateful.rs index 4dd87e07..976f8024 100644 --- a/src/session/handlers/stateful.rs +++ b/src/session/handlers/stateful.rs @@ -14,6 +14,8 @@ use tracing::{debug, error, warn}; use crate::constants::buffer::READER_CAPACITY; +const MAX_UPSTREAM_PIPELINE_DEPTH: usize = 16; + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(in crate::session) enum StatefulConnectionDisposition { RetireClient, @@ -257,10 +259,6 @@ impl ClientSession { .write_wire_to(backend_write) .await .map_err(|error| StatefulLoopError::backend(error.into()))?; - backend_write - .flush() - .await - .map_err(|error| StatefulLoopError::backend(error.into()))?; state.add_client_to_backend(request.request_wire_len().get()); state.mark_backend_request_sent(request.kind()); } @@ -446,6 +444,10 @@ impl ClientSession { BW: tokio::io::AsyncWrite + Unpin, { let mut command_reader = StatefulCommandReader::new(); + // Keep a small upstream window open while the client is producing + // commands. Replies are still consumed and ordered by the framer below; + // this counter only controls when buffered request bytes are flushed. + let mut pending_backend_writes = 0usize; let exit = loop { // Periodic metrics flush @@ -470,6 +472,13 @@ impl ClientSession { } if matches!(state.read_mode(), StatefulReadMode::DrainBackendReplies) { + if pending_backend_writes != 0 { + backend_write + .flush() + .await + .map_err(|error| StatefulLoopError::backend(error.into()))?; + pending_backend_writes = 0; + } match self.read_stateful_backend_bytes(&mut backend_read).await { Ok(Some((buffer, len))) => { self.write_stateful_backend_bytes( @@ -537,6 +546,15 @@ impl ClientSession { &mut state, ) .await?; + pending_backend_writes = + pending_backend_writes.saturating_add(1); + if pending_backend_writes >= MAX_UPSTREAM_PIPELINE_DEPTH { + backend_write + .flush() + .await + .map_err(|error| StatefulLoopError::backend(error.into()))?; + pending_backend_writes = 0; + } } else { // Auth path let auth_result = common::handle_stateful_auth_check( @@ -567,6 +585,13 @@ impl ClientSession { // Backend → Client result = self.read_stateful_backend_bytes(&mut backend_read) => { + if pending_backend_writes != 0 { + backend_write + .flush() + .await + .map_err(|error| StatefulLoopError::backend(error.into()))?; + pending_backend_writes = 0; + } match result { Ok(Some((buffer, len))) => { // Complete output after select! returns so client readiness cannot @@ -810,6 +835,63 @@ mod tests { backend.await.unwrap(); } + #[tokio::test] + async fn upstream_window_releases_two_commands_before_backend_replies() { + let session = test_session(); + let state = SessionLoopState::new(false); + let (mut client_end, proxy_client_end) = tokio::io::duplex(4096); + let (backend_end, proxy_backend_end) = tokio::io::duplex(4096); + + let backend = tokio::spawn(async move { + let (backend_read, mut backend_write) = tokio::io::split(backend_end); + let mut reader = BufReader::new(backend_read); + let mut first = String::new(); + let mut second = String::new(); + reader.read_line(&mut first).await.unwrap(); + reader.read_line(&mut second).await.unwrap(); + assert_eq!(first, "DATE\r\n"); + assert_eq!(second, "DATE\r\n"); + backend_write + .write_all(b"111 first\r\n111 second\r\n") + .await + .unwrap(); + }); + + client_end.write_all(b"DATE\r\nDATE\r\n").await.unwrap(); + let (proxy_client_read, proxy_client_write) = tokio::io::split(proxy_client_end); + let client_reader = BufReader::new(proxy_client_read); + let (backend_read, backend_write) = tokio::io::split(proxy_backend_end); + let proxy = tokio::spawn(async move { + session + .run_stateful_proxy_loop( + client_reader, + proxy_client_write, + backend_read, + backend_write, + state, + BackendId::from_index(0), + ) + .await + }); + + let mut responses = [0u8; b"111 first\r\n111 second\r\n".len()]; + tokio::time::timeout( + std::time::Duration::from_secs(1), + tokio::io::AsyncReadExt::read_exact(&mut client_end, &mut responses), + ) + .await + .expect("client did not receive both ordered replies") + .unwrap(); + assert_eq!(&responses, b"111 first\r\n111 second\r\n"); + + drop(client_end); + let _ = tokio::time::timeout(std::time::Duration::from_secs(1), proxy) + .await + .expect("stateful proxy did not stop") + .expect("stateful proxy task panicked"); + backend.await.unwrap(); + } + #[tokio::test] async fn test_client_disconnect_returns_metrics() { let session = test_session();