From 01e5be508051eebf19c233d48b57797a17331383 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sylwester=20R=C4=85pa=C5=82a?= Date: Thu, 3 Mar 2022 17:50:43 +0100 Subject: [PATCH] Add example of detecting client drops in bidirectional streams on server side (#931) --- examples/Cargo.toml | 27 ++++---- examples/src/streaming/client.rs | 75 +++++++++++++++++--- examples/src/streaming/server.rs | 115 +++++++++++++++++++++++++------ 3 files changed, 174 insertions(+), 43 deletions(-) diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 3857d69..6a61e10 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -180,37 +180,40 @@ path = "src/streaming/server.rs" [dependencies] async-stream = "0.3" -futures = {version = "0.3", default-features = false, features = ["alloc"]} +futures = { version = "0.3", default-features = false, features = ["alloc"] } prost = "0.9" -tokio = {version = "1.0", features = ["rt-multi-thread", "time", "fs", "macros", "net"]} -tokio-stream = {version = "0.1", features = ["net"]} -tonic = {path = "../tonic", features = ["tls", "compression"]} -tower = {version = "0.4"} +tokio = { version = "1.0", features = [ "rt-multi-thread", "time", "fs", "macros", "net",] } +tokio-stream = { version = "0.1", features = ["net"] } +tonic = { path = "../tonic", features = ["tls", "compression"] } +tower = { version = "0.4" } # Required for routeguide rand = "0.8" -serde = {version = "1.0", features = ["derive"]} +serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" # Tracing tracing = "0.1.16" tracing-attributes = "0.1" tracing-futures = "0.2" -tracing-subscriber = {version = "0.3", features = ["tracing-log"]} +tracing-subscriber = { version = "0.3", features = ["tracing-log"] } # Required for wellknown types prost-types = "0.9" # Hyper example http = "0.2" http-body = "0.4.2" -hyper = {version = "0.14", features = ["full"]} +hyper = { version = "0.14", features = ["full"] } pin-project = "1.0" warp = "0.3" # Health example -tonic-health = {path = "../tonic-health"} +tonic-health = { path = "../tonic-health" } # Reflection example listenfd = "0.3" -tonic-reflection = {path = "../tonic-reflection"} +tonic-reflection = { path = "../tonic-reflection" } # grpc-web example bytes = "1" -tonic-web = {path = "../tonic-web"} +tonic-web = { path = "../tonic-web" } +# streaming example +h2 = "0.3" + [build-dependencies] -tonic-build = {path = "../tonic-build", features = ["prost", "compression"]} +tonic-build = { path = "../tonic-build", features = ["prost", "compression"] } diff --git a/examples/src/streaming/client.rs b/examples/src/streaming/client.rs index 67a876a..20e7117 100644 --- a/examples/src/streaming/client.rs +++ b/examples/src/streaming/client.rs @@ -2,28 +2,85 @@ pub mod pb { tonic::include_proto!("grpc.examples.echo"); } +use futures::stream::Stream; +use std::time::Duration; +use tokio_stream::StreamExt; +use tonic::transport::Channel; + use pb::{echo_client::EchoClient, EchoRequest}; -#[tokio::main] -async fn main() -> Result<(), Box> { - let mut client = EchoClient::connect("http://[::1]:50051").await.unwrap(); +fn echo_requests_iter() -> impl Stream { + tokio_stream::iter(1..usize::MAX).map(|i| EchoRequest { + message: format!("msg {:02}", i), + }) +} +async fn streaming_echo(client: &mut EchoClient, num: usize) { let stream = client .server_streaming_echo(EchoRequest { message: "foo".into(), }) .await + .unwrap() + .into_inner(); + + // stream is infinite - take just 5 elements and then disconnect + let mut stream = stream.take(num); + while let Some(item) = stream.next().await { + println!("\trecived: {}", item.unwrap().message); + } + // stream is droped here and the disconnect info is send to server +} + +async fn bidirectional_streaming_echo(client: &mut EchoClient, num: usize) { + let in_stream = echo_requests_iter().take(num); + + let response = client + .bidirectional_streaming_echo(in_stream) + .await .unwrap(); - println!("Connected...now sleeping for 2 seconds..."); + let mut resp_stream = response.into_inner(); - tokio::time::sleep(std::time::Duration::from_secs(2)).await; + while let Some(recived) = resp_stream.next().await { + let recived = recived.unwrap(); + println!("\trecived message: `{}`", recived.message); + } +} - // Disconnect - drop(stream); - drop(client); +async fn bidirectional_streaming_echo_throttle(client: &mut EchoClient, dur: Duration) { + let in_stream = echo_requests_iter().throttle(dur); - println!("Disconnected..."); + let response = client + .bidirectional_streaming_echo(in_stream) + .await + .unwrap(); + + let mut resp_stream = response.into_inner(); + + while let Some(recived) = resp_stream.next().await { + let recived = recived.unwrap(); + println!("\trecived message: `{}`", recived.message); + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let mut client = EchoClient::connect("http://[::1]:50051").await.unwrap(); + + println!("Streaming echo:"); + streaming_echo(&mut client, 5).await; + tokio::time::sleep(Duration::from_secs(1)).await; //do not mess server println functions + + // Echo stream that sends 17 requests then gracefull end that conection + println!("\r\nBidirectional stream echo:"); + bidirectional_streaming_echo(&mut client, 17).await; + + // Echo stream that sends up to `usize::MAX` requets. One request each 2s. + // Exiting client with CTRL+C demostrate how to distinguise broken pipe from + //gracefull client disconnection (above example) on the server side. + println!("\r\nBidirectional stream echo (kill client with CTLR+C):"); + bidirectional_streaming_echo_throttle(&mut client, Duration::from_secs(2)).await; Ok(()) } diff --git a/examples/src/streaming/server.rs b/examples/src/streaming/server.rs index 9c07387..8646260 100644 --- a/examples/src/streaming/server.rs +++ b/examples/src/streaming/server.rs @@ -3,10 +3,9 @@ pub mod pb { } use futures::Stream; -use std::net::ToSocketAddrs; -use std::pin::Pin; -use std::task::{Context, Poll}; -use tokio::sync::oneshot; +use std::{error::Error, io::ErrorKind, net::ToSocketAddrs, pin::Pin, time::Duration}; +use tokio::sync::mpsc; +use tokio_stream::{wrappers::ReceiverStream, StreamExt}; use tonic::{transport::Server, Request, Response, Status, Streaming}; use pb::{EchoRequest, EchoResponse}; @@ -14,6 +13,29 @@ use pb::{EchoRequest, EchoResponse}; type EchoResult = Result, Status>; type ResponseStream = Pin> + Send>>; +fn match_for_io_error(err_status: &Status) -> Option<&std::io::Error> { + let mut err: &(dyn Error + 'static) = err_status; + + loop { + if let Some(io_err) = err.downcast_ref::() { + return Some(io_err); + } + + // h2::Error do not expose std::io::Error with `source()` + // https://github.com/hyperium/h2/pull/462 + if let Some(h2_err) = err.downcast_ref::() { + if let Some(io_err) = h2_err.get_io() { + return Some(io_err); + } + } + + err = match err.source() { + Some(err) => err, + None => return None, + }; + } +} + #[derive(Debug)] pub struct EchoServer {} @@ -29,28 +51,36 @@ impl pb::echo_server::Echo for EchoServer { &self, req: Request, ) -> EchoResult { - println!("Client connected from: {:?}", req.remote_addr()); + println!("EchoServer::server_streaming_echo"); + println!("\tclient connected from: {:?}", req.remote_addr()); - let (tx, rx) = oneshot::channel::<()>(); + // creating infinite stream with requested message + let repeat = std::iter::repeat(EchoResponse { + message: req.into_inner().message, + }); + let mut stream = Box::pin(tokio_stream::iter(repeat).throttle(Duration::from_millis(200))); + // spawn and channel are required if you want handle "disconnect" functionality + // the `out_stream` will not be polled after client disconnect + let (tx, rx) = mpsc::channel(128); tokio::spawn(async move { - let _ = rx.await; - println!("The rx resolved therefore the client disconnected!"); + while let Some(item) = stream.next().await { + match tx.send(Result::<_, Status>::Ok(item)).await { + Ok(_) => { + // item (server response) was queued to be send to client + } + Err(_item) => { + // output_stream was build from rx and both are dropped + break; + } + } + } + println!("\tclient disconnected"); }); - struct ClientDisconnect(oneshot::Sender<()>); - - impl Stream for ClientDisconnect { - type Item = Result; - - fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { - // A stream that never resolves to anything.... - Poll::Pending - } - } - + let output_stream = ReceiverStream::new(rx); Ok(Response::new( - Box::pin(ClientDisconnect(tx)) as Self::ServerStreamingEchoStream + Box::pin(output_stream) as Self::ServerStreamingEchoStream )) } @@ -65,9 +95,50 @@ impl pb::echo_server::Echo for EchoServer { async fn bidirectional_streaming_echo( &self, - _: Request>, + req: Request>, ) -> EchoResult { - Err(Status::unimplemented("not implemented")) + println!("EchoServer::bidirectional_streaming_echo"); + + let mut in_stream = req.into_inner(); + let (tx, rx) = mpsc::channel(128); + + // this spawn here is required if you want to handle connection error. + // If we just map `in_stream` and write it back as `out_stream` the `out_stream` + // will be drooped when connection error occurs and error will never be propagated + // to mapped version of `in_stream`. + tokio::spawn(async move { + while let Some(result) = in_stream.next().await { + match result { + Ok(v) => tx + .send(Ok(EchoResponse { message: v.message })) + .await + .expect("working rx"), + Err(err) => { + if let Some(io_err) = match_for_io_error(&err) { + if io_err.kind() == ErrorKind::BrokenPipe { + // here you can handle special case when client + // disconnected in unexpected way + eprintln!("\tclient disconnected: broken pipe"); + break; + } + } + + match tx.send(Err(err)).await { + Ok(_) => (), + Err(_err) => break, // response was droped + } + } + } + } + println!("\tstream ended"); + }); + + // echo just write the same data that was received + let out_stream = ReceiverStream::new(rx); + + Ok(Response::new( + Box::pin(out_stream) as Self::BidirectionalStreamingEchoStream + )) } }