test: add regression test for stream close-frame ordering

Exercises the fix in the previous commit. A server-streaming handler
sends one payload and returns (DATA frame immediately followed by
REMOTE_CLOSED). The client calls it 1000 times on a multi_thread
runtime and asserts that every iteration receives the payload before
EOF. Without the handle_msg fix this fails with ~1% of iterations
silently dropping the payload.

Signed-off-by: Shiv Bhosale <shvbsle@amazon.com>

Author: shvbsle

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,9 @@ 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"
+
 [build-dependencies]
 ttrpc-codegen = { path = "../ttrpc-codegen"}

example/async-stream-close-order.rs

@@ -0,0 +1,137 @@
+// 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
+// calls this 1000 times on a multi-threaded runtime and asserts that every
+// iteration receives the payload before EOF.
+
+mod protocols;
+
+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 = 1000;
+
+struct Svc;
+
+#[async_trait]
+impl streaming_ttrpc::Streaming for Svc {
+    async fn echo(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _req: streaming::EchoPayload,
+    ) -> ::ttrpc::Result<streaming::EchoPayload> {
+        unimplemented!()
+    }
+
+    async fn echo_stream(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _s: ::ttrpc::r#async::ServerStream<streaming::EchoPayload, streaming::EchoPayload>,
+    ) -> ::ttrpc::Result<()> {
+        unimplemented!()
+    }
+
+    async fn sum_stream(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _s: ::ttrpc::r#async::ServerStreamReceiver<streaming::Part>,
+    ) -> ::ttrpc::Result<streaming::Sum> {
+        unimplemented!()
+    }
+
+    async fn divide_stream(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _sum: streaming::Sum,
+        _s: ::ttrpc::r#async::ServerStreamSender<streaming::Part>,
+    ) -> ::ttrpc::Result<()> {
+        unimplemented!()
+    }
+
+    async fn echo_null(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _s: ::ttrpc::r#async::ServerStreamReceiver<streaming::EchoPayload>,
+    ) -> ::ttrpc::Result<empty::Empty> {
+        unimplemented!()
+    }
+
+    async fn echo_null_stream(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _s: ::ttrpc::r#async::ServerStream<empty::Empty, streaming::EchoPayload>,
+    ) -> ::ttrpc::Result<()> {
+        unimplemented!()
+    }
+
+    async fn echo_default_value(
+        &self,
+        _ctx: &::ttrpc::r#async::TtrpcContext,
+        _req: streaming::EchoPayload,
+        s: ::ttrpc::r#async::ServerStreamSender<streaming::EchoPayload>,
+    ) -> ::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<streaming::EchoPayload>,
+    ) -> ::ttrpc::Result<()> {
+        unimplemented!()
+    }
+}
+
+#[tokio::main(flavor = "multi_thread")]
+async fn main() {
+    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 mut got_payload = 0usize;
+    let mut eof_without_payload = 0usize;
+
+    for _ in 0..ITERATIONS {
+        let ctx = ttrpc::context::with_timeout(5_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(_)) => got_payload += 1,
+            Ok(None) => eof_without_payload += 1,
+            Err(e) => panic!("unexpected error from recv: {:?}", e),
+        }
+    }
+
+    server.shutdown().await.unwrap();
+
+    eprintln!(
+        "iterations={ITERATIONS} got_payload={got_payload} eof_without_payload={eof_without_payload}"
+    );
+    assert_eq!(
+        eof_without_payload, 0,
+        "server sent a payload on every iteration but the client saw EOF \
+         {eof_without_payload}/{ITERATIONS} times (data frames dropped)"
+    );
+}

tests/run-examples.rs

@@ -109,3 +109,35 @@ fn run_examples() -> Result<(), Box<dyn std::error::Error>> {
 
     Ok(())
 }
+
+#[test]
+#[cfg(unix)]
+fn stream_close_order() -> Result<(), Box<dyn std::error::Error>> {
+    // 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(())
+}