talksik/tonic
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add example of detecting client drops in bidirectional streams on server side (#931)

Browse Source
This commit is contained in:
Sylwester Rąpała
2022-03-03 17:50:43 +01:00
committed by GitHub
parent e30bb7ede7
commit 01e5be5080
3 changed files with 174 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/Cargo.toml
+15 -12
View File
@@ -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"] }
examples/src/streaming/client.rs
+66 -9
View File
@@ -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<dyn std::error::Error>> {
let mut client = EchoClient::connect("http://[::1]:50051").await.unwrap();
fn echo_requests_iter() -> impl Stream<Item = EchoRequest> {
tokio_stream::iter(1..usize::MAX).map(|i| EchoRequest {
message: format!("msg {:02}", i),
})
}
async fn streaming_echo(client: &mut EchoClient<Channel>, 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<Channel>, 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<Channel>, 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<dyn std::error::Error>> {
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(())
}
examples/src/streaming/server.rs
+93 -22
View File
@@ -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<T> = Result<Response<T>, Status>;
type ResponseStream = Pin<Box<dyn Stream<Item = Result<EchoResponse, Status>> + 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::<std::io::Error>() {
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::<h2::Error>() {
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<EchoRequest>,
) -> EchoResult<Self::ServerStreamingEchoStream> {
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<EchoResponse, Status>;
fn poll_next(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Option<Self::Item>> {
// 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<Streaming<EchoRequest>>,
req: Request<Streaming<EchoRequest>>,
) -> EchoResult<Self::BidirectionalStreamingEchoStream> {
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
))
}
}