fix(client): detect HTTP2 connection closures sooner
This commit is contained in:
		| @@ -1,7 +1,7 @@ | |||||||
| use bytes::IntoBuf; | use bytes::IntoBuf; | ||||||
| use futures::{Async, Future, Poll, Stream}; | use futures::{Async, Future, Poll, Stream}; | ||||||
| use futures::future::{self, Either}; | use futures::future::{self, Either}; | ||||||
| use futures::sync::mpsc; | use futures::sync::{mpsc, oneshot}; | ||||||
| use h2::client::{Builder, Handshake, SendRequest}; | use h2::client::{Builder, Handshake, SendRequest}; | ||||||
| use tokio_io::{AsyncRead, AsyncWrite}; | use tokio_io::{AsyncRead, AsyncWrite}; | ||||||
|  |  | ||||||
| @@ -18,6 +18,10 @@ type ClientRx<B> = ::client::dispatch::Receiver<Request<B>, Response<Body>>; | |||||||
| /// other handles to it have been dropped, so that it can shutdown. | /// other handles to it have been dropped, so that it can shutdown. | ||||||
| type ConnDropRef = mpsc::Sender<Never>; | type ConnDropRef = mpsc::Sender<Never>; | ||||||
|  |  | ||||||
|  | /// A oneshot channel watches the `Connection` task, and when it completes, | ||||||
|  | /// the "dispatch" task will be notified and can shutdown sooner. | ||||||
|  | type ConnEof = oneshot::Receiver<Never>; | ||||||
|  |  | ||||||
| pub(crate) struct Client<T, B> | pub(crate) struct Client<T, B> | ||||||
| where | where | ||||||
|     B: Payload, |     B: Payload, | ||||||
| @@ -29,7 +33,7 @@ where | |||||||
|  |  | ||||||
| enum State<T, B> where B: IntoBuf { | enum State<T, B> where B: IntoBuf { | ||||||
|     Handshaking(Handshake<T, B>), |     Handshaking(Handshake<T, B>), | ||||||
|     Ready(SendRequest<B>, ConnDropRef), |     Ready(SendRequest<B>, ConnDropRef, ConnEof), | ||||||
| } | } | ||||||
|  |  | ||||||
| impl<T, B> Client<T, B> | impl<T, B> Client<T, B> | ||||||
| @@ -66,6 +70,7 @@ where | |||||||
|                     // in h2 where dropping all SendRequests won't notify a |                     // in h2 where dropping all SendRequests won't notify a | ||||||
|                     // parked Connection. |                     // parked Connection. | ||||||
|                     let (tx, rx) = mpsc::channel(0); |                     let (tx, rx) = mpsc::channel(0); | ||||||
|  |                     let (cancel_tx, cancel_rx) = oneshot::channel(); | ||||||
|                     let rx = rx.into_future() |                     let rx = rx.into_future() | ||||||
|                         .map(|(msg, _)| match msg { |                         .map(|(msg, _)| match msg { | ||||||
|                             Some(never) => match never {}, |                             Some(never) => match never {}, | ||||||
| @@ -73,7 +78,10 @@ where | |||||||
|                         }) |                         }) | ||||||
|                         .map_err(|_| -> Never { unreachable!("mpsc cannot error") }); |                         .map_err(|_| -> Never { unreachable!("mpsc cannot error") }); | ||||||
|                     let fut = conn |                     let fut = conn | ||||||
|                         .inspect(|_| trace!("connection complete")) |                         .inspect(move |_| { | ||||||
|  |                             drop(cancel_tx); | ||||||
|  |                             trace!("connection complete") | ||||||
|  |                         }) | ||||||
|                         .map_err(|e| debug!("connection error: {}", e)) |                         .map_err(|e| debug!("connection error: {}", e)) | ||||||
|                         .select2(rx) |                         .select2(rx) | ||||||
|                         .then(|res| match res { |                         .then(|res| match res { | ||||||
| @@ -92,10 +100,21 @@ where | |||||||
|                             Err(Either::B((never, _))) => match never {}, |                             Err(Either::B((never, _))) => match never {}, | ||||||
|                         }); |                         }); | ||||||
|                     self.executor.execute(fut)?; |                     self.executor.execute(fut)?; | ||||||
|                     State::Ready(request_tx, tx) |                     State::Ready(request_tx, tx, cancel_rx) | ||||||
|                 }, |                 }, | ||||||
|                 State::Ready(ref mut tx, ref conn_dropper) => { |                 State::Ready(ref mut tx, ref conn_dropper, ref mut cancel_rx) => { | ||||||
|                     try_ready!(tx.poll_ready().map_err(::Error::new_h2)); |                     match tx.poll_ready() { | ||||||
|  |                         Ok(Async::Ready(())) => (), | ||||||
|  |                         Ok(Async::NotReady) => return Ok(Async::NotReady), | ||||||
|  |                         Err(err) => { | ||||||
|  |                             return if err.reason() == Some(::h2::Reason::NO_ERROR) { | ||||||
|  |                                 trace!("connection gracefully shutdown"); | ||||||
|  |                                 Ok(Async::Ready(Dispatched::Shutdown)) | ||||||
|  |                             } else { | ||||||
|  |                                 Err(::Error::new_h2(err)) | ||||||
|  |                             }; | ||||||
|  |                         } | ||||||
|  |                     } | ||||||
|                     match self.rx.poll() { |                     match self.rx.poll() { | ||||||
|                         Ok(Async::Ready(Some((req, cb)))) => { |                         Ok(Async::Ready(Some((req, cb)))) => { | ||||||
|                             // check that future hasn't been canceled already |                             // check that future hasn't been canceled already | ||||||
| @@ -157,7 +176,16 @@ where | |||||||
|                             continue; |                             continue; | ||||||
|                         }, |                         }, | ||||||
|  |  | ||||||
|                         Ok(Async::NotReady) => return Ok(Async::NotReady), |                         Ok(Async::NotReady) => { | ||||||
|  |                             match cancel_rx.poll() { | ||||||
|  |                                 Ok(Async::Ready(never)) => match never {}, | ||||||
|  |                                 Ok(Async::NotReady) => return Ok(Async::NotReady), | ||||||
|  |                                 Err(_conn_is_eof) => { | ||||||
|  |                                     trace!("connection task is closed, closing dispatch task"); | ||||||
|  |                                     return Ok(Async::Ready(Dispatched::Shutdown)); | ||||||
|  |                                 } | ||||||
|  |                             } | ||||||
|  |                         }, | ||||||
|  |  | ||||||
|                         Ok(Async::Ready(None)) => { |                         Ok(Async::Ready(None)) => { | ||||||
|                             trace!("client::dispatch::Sender dropped"); |                             trace!("client::dispatch::Sender dropped"); | ||||||
|   | |||||||
| @@ -2095,6 +2095,57 @@ mod conn { | |||||||
|         assert_eq!(vec, b"bar=foo"); |         assert_eq!(vec, b"bar=foo"); | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |  | ||||||
|  |     #[test] | ||||||
|  |     fn http2_detect_conn_eof() { | ||||||
|  |         use futures::future; | ||||||
|  |         use hyper::{Response, Server}; | ||||||
|  |         use hyper::service::service_fn_ok; | ||||||
|  |         use tokio::timer::Delay; | ||||||
|  |  | ||||||
|  |         let _ = pretty_env_logger::try_init(); | ||||||
|  |  | ||||||
|  |         let mut rt = Runtime::new().unwrap(); | ||||||
|  |  | ||||||
|  |         let server = Server::bind(&([127, 0, 0, 1], 0).into()) | ||||||
|  |             .http2_only(true) | ||||||
|  |             .serve(|| service_fn_ok(|_req| { | ||||||
|  |                 Response::new(Body::empty()) | ||||||
|  |             })); | ||||||
|  |         let addr = server.local_addr(); | ||||||
|  |         let (shdn_tx, shdn_rx) = oneshot::channel(); | ||||||
|  |         rt.spawn(server.with_graceful_shutdown(shdn_rx).map_err(|e| panic!("server error: {:?}", e))); | ||||||
|  |  | ||||||
|  |         let io = rt.block_on(tcp_connect(&addr)).expect("tcp connect"); | ||||||
|  |         let (mut client, conn) = rt.block_on( | ||||||
|  |             conn::Builder::new().http2_only(true).handshake::<_, Body>(io) | ||||||
|  |         ).expect("http handshake"); | ||||||
|  |         rt.spawn(conn.map_err(|e| panic!("client conn error: {:?}", e))); | ||||||
|  |  | ||||||
|  |  | ||||||
|  |         // Sanity check that client is ready | ||||||
|  |         rt.block_on(future::poll_fn(|| client.poll_ready())).expect("client poll ready sanity"); | ||||||
|  |  | ||||||
|  |         let req = Request::builder() | ||||||
|  |             .uri(format!("http://{}/", addr)) | ||||||
|  |             .body(Body::empty()) | ||||||
|  |             .expect("request builder"); | ||||||
|  |  | ||||||
|  |         rt.block_on(client.send_request(req)).expect("req1 send"); | ||||||
|  |  | ||||||
|  |         // Sanity check that client is STILL ready | ||||||
|  |         rt.block_on(future::poll_fn(|| client.poll_ready())).expect("client poll ready after"); | ||||||
|  |  | ||||||
|  |         // Trigger the server shutdown... | ||||||
|  |         let _ = shdn_tx.send(()); | ||||||
|  |  | ||||||
|  |         // Allow time for graceful shutdown roundtrips... | ||||||
|  |         rt.block_on(Delay::new(::std::time::Instant::now() + Duration::from_millis(100)).map_err(|e| panic!("delay error: {:?}", e))).expect("delay"); | ||||||
|  |  | ||||||
|  |         // After graceful shutdown roundtrips, the client should be closed... | ||||||
|  |         rt.block_on(future::poll_fn(|| client.poll_ready())).expect_err("client should be closed"); | ||||||
|  |     } | ||||||
|  |  | ||||||
|     struct DebugStream { |     struct DebugStream { | ||||||
|         tcp: TcpStream, |         tcp: TcpStream, | ||||||
|         shutdown_called: bool, |         shutdown_called: bool, | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user