From afa1063e0debbe6454bf868e9f41c5b7cf353012 Mon Sep 17 00:00:00 2001 From: Shiv Bhosale Date: Sat, 16 May 2026 03:24:36 +0000 Subject: [PATCH] async/client: don't spawn per incoming frame ClientReader::handle_msg spawns a new task per frame. For a server-streaming RPC, the final DATA frame and the subsequent FLAG_REMOTE_CLOSED frame then race: if the close-frame task grabs the req_map lock first, it removes the stream from the map, and the preceding data-frame task finds nothing and silently drops the payload. The stream consumer sees Ok(None) (EOF) without ever observing the payload the server sent. The connection read loop already awaits handle_msg per frame, so processing inline preserves per-stream wire order. It also gives the per-stream mpsc natural back-pressure in place of the unbounded per-frame spawning. handle_err gets the same treatment for consistency. Add two regression examples: - async-stream-close-order: verifies DATA before CLOSE is not lost - async-data-order: verifies multi-frame streams maintain sequence Signed-off-by: Shiv Bhosale --- example/Cargo.toml | 10 +- example/async-data-order.rs | 167 ++++++++++++++++++++++++++++ example/async-stream-close-order.rs | 165 +++++++++++++++++++++++++++ src/asynchronous/client.rs | 32 +++--- tests/run-examples.rs | 32 ++++++ 5 files changed, 386 insertions(+), 20 deletions(-) create mode 100644 example/async-data-order.rs create mode 100644 example/async-stream-close-order.rs diff --git a/example/Cargo.toml b/example/Cargo.toml index 97c5cb0..7c8d15e 100644 --- a/example/Cargo.toml +++ b/example/Cargo.toml @@ -21,7 +21,7 @@ simple-logging = "2.0.2" nix = "0.23.0" ttrpc = { path = "../", features = ["async"] } ctrlc = { version = "3.0", features = ["termination"] } -tokio = { version = "1.0.1", features = ["signal", "time"] } +tokio = { version = "1.0.1", features = ["signal", "time", "rt-multi-thread", "macros"] } async-trait = "0.1.42" rand = "0.8.5" clap = { version = "4.5.40", features = ["derive"] } @@ -50,5 +50,13 @@ path = "./async-stream-server.rs" name = "async-stream-client" path = "./async-stream-client.rs" +[[example]] +name = "async-stream-close-order" +path = "./async-stream-close-order.rs" + +[[example]] +name = "async-data-order" +path = "./async-data-order.rs" + [build-dependencies] ttrpc-codegen = { path = "../ttrpc-codegen"} diff --git a/example/async-data-order.rs b/example/async-data-order.rs new file mode 100644 index 0000000..5e662df --- /dev/null +++ b/example/async-data-order.rs @@ -0,0 +1,167 @@ +// Regression test for per-stream DATA frame ordering on the async client. +// +// A server-streaming handler sends 10 DATA frames (seq 0..9) per stream. +// The client opens 200 concurrent streams on a multi-threaded runtime and +// asserts that every stream receives its DATA frames in sequence order. +// With the old spawn-per-frame design, 200 concurrent streams create enough +// task-list separation between same-stream DATA frames to trigger reordering. + +mod protocols; + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use async_trait::async_trait; +use protocols::asynchronous::{empty, streaming, streaming_ttrpc}; +use ttrpc::asynchronous::{Client, Server}; + +const SOCK: &str = "unix:///tmp/ttrpc-test-data-order"; +const ITERATIONS: usize = 100; +const FRAMES_PER_STREAM: u32 = 50; +const WORKER_THREADS: usize = 4; + +struct Svc; + +#[async_trait] +impl streaming_ttrpc::Streaming for Svc { + async fn echo( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _req: streaming::EchoPayload, + ) -> ::ttrpc::Result { + unimplemented!() + } + + async fn echo_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStream, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn sum_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStreamReceiver, + ) -> ::ttrpc::Result { + unimplemented!() + } + + async fn divide_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _sum: streaming::Sum, + _s: ::ttrpc::r#async::ServerStreamSender, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn echo_null( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStreamReceiver, + ) -> ::ttrpc::Result { + unimplemented!() + } + + async fn echo_null_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStream, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn echo_default_value( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _req: streaming::EchoPayload, + _s: ::ttrpc::r#async::ServerStreamSender, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn server_send_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _: empty::Empty, + s: ::ttrpc::r#async::ServerStreamSender, + ) -> ::ttrpc::Result<()> { + for seq in 0..FRAMES_PER_STREAM { + s.send(&streaming::EchoPayload { + seq, + msg: format!("{}", seq), + ..Default::default() + }) + .await + .unwrap(); + } + Ok(()) + } +} + +fn main() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(WORKER_THREADS) + .enable_all() + .build() + .unwrap() + .block_on(run()); +} + +async fn run() { + let path = SOCK.strip_prefix("unix://").unwrap(); + let _ = std::fs::remove_file(path); + + let service = streaming_ttrpc::create_streaming(Arc::new(Svc {})); + let mut server = Server::new().bind(SOCK).unwrap().register_service(service); + server.start().await.unwrap(); + + let c = Client::connect(SOCK).await.unwrap(); + let sc = streaming_ttrpc::StreamingClient::new(c); + + let out_of_order = Arc::new(AtomicUsize::new(0)); + let mut handles = Vec::with_capacity(ITERATIONS); + + for _ in 0..ITERATIONS { + let sc = sc.clone(); + let out_of_order = out_of_order.clone(); + handles.push(tokio::spawn(async move { + let ctx = ttrpc::context::with_timeout(10_000_000_000); + let mut stream: ttrpc::asynchronous::ClientStreamReceiver = sc + .server_send_stream(ctx, &empty::Empty::default()) + .await + .expect("failed to open stream"); + + let mut expected_seq: u32 = 0; + let mut saw_reorder = false; + while let Some(payload) = stream.recv().await.unwrap() { + if payload.seq != expected_seq { + saw_reorder = true; + } + expected_seq += 1; + } + if saw_reorder { + out_of_order.fetch_add(1, Ordering::Relaxed); + } + })); + } + + for h in handles { + let _ = h.await; + } + + server.shutdown().await.unwrap(); + + let bad = out_of_order.load(Ordering::Relaxed); + eprintln!( + "repeating the experiment {} times, we find out of order frames {}/{} times", + ITERATIONS, bad, ITERATIONS + ); + assert_eq!( + bad, 0, + "repeating the experiment {} times, we find out of order frames {}/{} times", + ITERATIONS, bad, ITERATIONS + ); +} diff --git a/example/async-stream-close-order.rs b/example/async-stream-close-order.rs new file mode 100644 index 0000000..4d58459 --- /dev/null +++ b/example/async-stream-close-order.rs @@ -0,0 +1,165 @@ +// Regression test for per-stream frame ordering on the async client. +// +// A server-streaming handler sends one payload and returns, producing a DATA +// frame immediately followed by a REMOTE_CLOSED frame on the wire. The client +// opens 200 concurrent streams in 10 batches (2000 total) on a multi-threaded +// runtime and asserts that every stream receives the payload before EOF. +// Concurrency creates scheduler contention that triggers the race far more +// reliably than sequential iterations. + +mod protocols; + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use async_trait::async_trait; +use protocols::asynchronous::{empty, streaming, streaming_ttrpc}; +use ttrpc::asynchronous::{Client, Server}; + +const SOCK: &str = "unix:///tmp/ttrpc-test-close-order"; +const ITERATIONS: usize = 2000; +const BATCH_SIZE: usize = 200; +const WORKER_THREADS: usize = 4; + +struct Svc; + +#[async_trait] +impl streaming_ttrpc::Streaming for Svc { + async fn echo( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _req: streaming::EchoPayload, + ) -> ::ttrpc::Result { + unimplemented!() + } + + async fn echo_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStream, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn sum_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStreamReceiver, + ) -> ::ttrpc::Result { + unimplemented!() + } + + async fn divide_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _sum: streaming::Sum, + _s: ::ttrpc::r#async::ServerStreamSender, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn echo_null( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStreamReceiver, + ) -> ::ttrpc::Result { + unimplemented!() + } + + async fn echo_null_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _s: ::ttrpc::r#async::ServerStream, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } + + async fn echo_default_value( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _req: streaming::EchoPayload, + s: ::ttrpc::r#async::ServerStreamSender, + ) -> ::ttrpc::Result<()> { + s.send(&streaming::EchoPayload { + seq: 1, + msg: "hello".into(), + ..Default::default() + }) + .await + .unwrap(); + Ok(()) + } + + async fn server_send_stream( + &self, + _ctx: &::ttrpc::r#async::TtrpcContext, + _req: empty::Empty, + _s: ::ttrpc::r#async::ServerStreamSender, + ) -> ::ttrpc::Result<()> { + unimplemented!() + } +} + +fn main() { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(WORKER_THREADS) + .enable_all() + .build() + .unwrap() + .block_on(run()); +} + +async fn run() { + let path = SOCK.strip_prefix("unix://").unwrap(); + let _ = std::fs::remove_file(path); + + let service = streaming_ttrpc::create_streaming(Arc::new(Svc {})); + let mut server = Server::new().bind(SOCK).unwrap().register_service(service); + server.start().await.unwrap(); + + let c = Client::connect(SOCK).await.unwrap(); + let sc = streaming_ttrpc::StreamingClient::new(c); + + let num_batches = ITERATIONS / BATCH_SIZE; + let eof_without_payload = Arc::new(AtomicUsize::new(0)); + + for _batch in 0..num_batches { + let mut handles = Vec::with_capacity(BATCH_SIZE); + for _ in 0..BATCH_SIZE { + let sc = sc.clone(); + let eof_without_payload = eof_without_payload.clone(); + handles.push(tokio::spawn(async move { + let ctx = ttrpc::context::with_timeout(10_000_000_000); + let mut stream = sc + .echo_default_value(ctx, &streaming::EchoPayload::default()) + .await + .expect("failed to open stream"); + + match stream.recv().await { + Ok(Some(_)) => {} + Ok(None) => { + eof_without_payload.fetch_add(1, Ordering::Relaxed); + } + Err(e) => panic!("unexpected error from recv: {:?}", e), + } + })); + } + + for h in handles { + let _ = h.await; + } + } + + server.shutdown().await.unwrap(); + + let lost = eof_without_payload.load(Ordering::Relaxed); + eprintln!( + "repeating the experiment {} times, we find data loss {}/{} times", + ITERATIONS, lost, ITERATIONS + ); + assert_eq!( + lost, 0, + "repeating the experiment {} times, we find data loss {}/{} times", + ITERATIONS, lost, ITERATIONS + ); +} diff --git a/src/asynchronous/client.rs b/src/asynchronous/client.rs index 524d643..0249dfe 100644 --- a/src/asynchronous/client.rs +++ b/src/asynchronous/client.rs @@ -232,7 +232,7 @@ impl WriterDelegate for ClientWriter { } async fn get_resp_tx( - req_map: Arc>>, + req_map: &Arc>>, header: &MessageHeader, ) -> Option { let resp_tx = match header.type_ { @@ -313,26 +313,20 @@ impl ReaderDelegate for ClientReader { async fn exit(&self) {} async fn handle_err(&self, header: MessageHeader, e: Error) { - let req_map = self.streams.clone(); - tokio::spawn(async move { - if let Some(resp_tx) = get_resp_tx(req_map, &header).await { - resp_tx - .send(Err(e)) - .await - .unwrap_or_else(|_e| error!("The request has returned")); - } - }); + if let Some(resp_tx) = get_resp_tx(&self.streams, &header).await { + resp_tx + .send(Err(e)) + .await + .unwrap_or_else(|_e| error!("The request has returned")); + } } async fn handle_msg(&self, msg: GenMessage) { - let req_map = self.streams.clone(); - tokio::spawn(async move { - if let Some(resp_tx) = get_resp_tx(req_map, &msg.header).await { - resp_tx - .send(Ok(msg)) - .await - .unwrap_or_else(|_e| error!("The request has returned")); - } - }); + if let Some(resp_tx) = get_resp_tx(&self.streams, &msg.header).await { + resp_tx + .send(Ok(msg)) + .await + .unwrap_or_else(|_e| error!("The request has returned")); + } } } diff --git a/tests/run-examples.rs b/tests/run-examples.rs index d2996c0..1df5d14 100644 --- a/tests/run-examples.rs +++ b/tests/run-examples.rs @@ -109,3 +109,35 @@ fn run_examples() -> Result<(), Box> { Ok(()) } + +#[test] +#[cfg(unix)] +fn stream_close_order() -> Result<(), Box> { + // Self-contained test: server + client in one process, multi-threaded + // runtime. Verifies that the final DATA frame of a server-streaming RPC + // is not dropped due to a race with the subsequent REMOTE_CLOSED frame. + let mut cmd = do_run_example("async-stream-close-order", &[]); + let mut child = cmd.spawn().unwrap(); + + let timeout = Duration::from_secs(120); + let start = std::time::Instant::now(); + loop { + if start.elapsed() > timeout { + child.kill().unwrap_or(()); + panic!("async-stream-close-order timed out"); + } + match child.try_wait() { + Ok(Some(status)) => { + wait_with_output("async-stream-close-order", child); + assert!( + status.success(), + "async-stream-close-order failed (data frames dropped)" + ); + break; + } + Ok(None) => continue, + Err(e) => panic!("Error waiting for async-stream-close-order: {:?}", e), + } + } + Ok(()) +}