Skip to content

fix: bound stateful upstream request window #123

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mjc merged 1 commit into from
Sep 14, 2026
+137 −9
Merged

fix: bound stateful upstream request window #123

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 51 additions & 5 deletions benches/end_to_end_proxy.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,17 +148,17 @@ struct BenchProxy {
}

impl BenchProxy {
async fn start(body_len: usize, cache: Option<Cache>) -> Self {
async fn start(body_len: usize, cache: Option<Cache>, 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();
Expand All @@ -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(),
}
}
Expand All @@ -189,6 +189,19 @@ impl BenchProxy {
)
.await
}

async fn article_window_roundtrip(&mut self) -> usize {
self.stream
.write_all(b"ARTICLE <bench@example.com>\r\nARTICLE <bench@example.com>\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<u8>) -> usize {
Expand Down Expand Up @@ -219,7 +232,7 @@ async fn read_exact_response_into(

fn bench_roundtrip(bencher: Bencher, body_len: usize, cache: Option<Cache>, 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);
Expand All @@ -232,6 +245,25 @@ fn bench_roundtrip(bencher: Bencher, body_len: usize, cache: Option<Cache>, 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};
Expand Down Expand Up @@ -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);
}
}
90 changes: 86 additions & 4 deletions src/session/handlers/stateful.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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());
}
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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();
Expand Down
Loading