fix(client): ensure idle connection is pooled before response body finishes
This commit is contained in:
		| @@ -128,6 +128,10 @@ impl<B> SendRequest<B> | ||||
|         self.dispatch.poll_ready(cx) | ||||
|     } | ||||
|  | ||||
|     pub(super) fn is_ready(&self) -> bool { | ||||
|         self.dispatch.is_ready() | ||||
|     } | ||||
|  | ||||
|     pub(super) fn is_closed(&self) -> bool { | ||||
|         self.dispatch.is_closed() | ||||
|     } | ||||
|   | ||||
| @@ -45,6 +45,10 @@ impl<T, U> Sender<T, U> { | ||||
|         } | ||||
|     } | ||||
|  | ||||
|     pub fn is_ready(&self) -> bool { | ||||
|         self.giver.is_wanting() | ||||
|     } | ||||
|  | ||||
|     pub fn is_closed(&self) -> bool { | ||||
|         self.giver.is_canceled() | ||||
|     } | ||||
|   | ||||
| @@ -7,6 +7,7 @@ use std::sync::Arc; | ||||
| use std::time::Duration; | ||||
|  | ||||
| use futures::{Async, Future, FutureExt, Never, Poll}; | ||||
| use futures::channel::oneshot; | ||||
| use futures::future; | ||||
| use futures::task; | ||||
| use http::{Method, Request, Response, Uri, Version}; | ||||
| @@ -235,7 +236,7 @@ where C: Connect<Error=io::Error> + Sync + 'static, | ||||
|                         ClientError::Normal(err) | ||||
|                     } | ||||
|                 }) | ||||
|                 .and_then(move |res| { | ||||
|                 .and_then(move |mut res| { | ||||
|                     future::lazy(move |cx| { | ||||
|                         // when pooled is dropped, it will try to insert back into the | ||||
|                         // pool. To delay that, spawn a future that completes once the | ||||
| @@ -245,14 +246,24 @@ where C: Connect<Error=io::Error> + Sync + 'static, | ||||
|                         // for a new request to start. | ||||
|                         // | ||||
|                         // It won't be ready if there is a body to stream. | ||||
|                         if let Ok(Async::Pending) = pooled.tx.poll_ready(cx) { | ||||
|                         if pooled.tx.is_ready() { | ||||
|                             drop(pooled); | ||||
|                         } else if !res.body().is_empty() { | ||||
|                             let (delayed_tx, delayed_rx) = oneshot::channel(); | ||||
|                             res.body_mut().delayed_eof(delayed_rx); | ||||
|                             // If the executor doesn't have room, oh well. Things will likely | ||||
|                             // be blowing up soon, but this specific task isn't required. | ||||
|                             execute(future::poll_fn(move |cx| { | ||||
|                                 pooled.tx.poll_ready(cx).or(Ok(Async::Ready(()))) | ||||
|                             }), cx).ok(); | ||||
|                             let fut = future::poll_fn(move |cx| { | ||||
|                                 pooled.tx.poll_ready(cx) | ||||
|                             }) | ||||
|                                 .then(move |_| { | ||||
|                                     // At this point, `pooled` is dropped, and had a chance | ||||
|                                     // to insert into the pool (if conn was idle) | ||||
|                                     drop(delayed_tx); | ||||
|                                     Ok(()) | ||||
|                                 }); | ||||
|                             execute(fut, cx).ok(); | ||||
|                         } | ||||
|  | ||||
|                         Ok(res) | ||||
|                     }) | ||||
|                 }); | ||||
|   | ||||
| @@ -3,7 +3,7 @@ use std::borrow::Cow; | ||||
| use std::fmt; | ||||
|  | ||||
| use bytes::Bytes; | ||||
| use futures::{Async, Future, Poll, Stream, StreamExt}; | ||||
| use futures::{Async, Future, Never, Poll, Stream, StreamExt}; | ||||
| use futures::task; | ||||
| use futures::channel::{mpsc, oneshot}; | ||||
| use http::HeaderMap; | ||||
| @@ -125,6 +125,15 @@ impl<E: Entity> Stream for EntityStream<E> { | ||||
| #[must_use = "streams do nothing unless polled"] | ||||
| pub struct Body { | ||||
|     kind: Kind, | ||||
|     /// Allow the client to pass a future to delay the `Body` from returning | ||||
|     /// EOF. This allows the `Client` to try to put the idle connection | ||||
|     /// back into the pool before the body is "finished". | ||||
|     /// | ||||
|     /// The reason for this is so that creating a new request after finishing | ||||
|     /// streaming the body of a response could sometimes result in creating | ||||
|     /// a brand new connection, since the pool didn't know about the idle | ||||
|     /// connection yet. | ||||
|     delayed_eof: Option<DelayEof>, | ||||
| } | ||||
|  | ||||
| enum Kind { | ||||
| @@ -137,6 +146,17 @@ enum Kind { | ||||
|     Empty, | ||||
| } | ||||
|  | ||||
| type DelayEofUntil = oneshot::Receiver<Never>; | ||||
|  | ||||
| enum DelayEof { | ||||
|     /// Initial state, stream hasn't seen EOF yet. | ||||
|     NotEof(DelayEofUntil), | ||||
|     /// Transitions to this state once we've seen `poll` try to | ||||
|     /// return EOF (`None`). This future is then polled, and | ||||
|     /// when it completes, the Body finally returns EOF (`None`). | ||||
|     Eof(DelayEofUntil), | ||||
| } | ||||
|  | ||||
| /// A sender half used with `Body::channel()`. | ||||
| #[derive(Debug)] | ||||
| pub struct Sender { | ||||
| @@ -253,6 +273,63 @@ impl Body { | ||||
|     fn new(kind: Kind) -> Body { | ||||
|         Body { | ||||
|             kind: kind, | ||||
|             delayed_eof: None, | ||||
|         } | ||||
|     } | ||||
|  | ||||
|     pub(crate) fn delayed_eof(&mut self, fut: DelayEofUntil) { | ||||
|         self.delayed_eof = Some(DelayEof::NotEof(fut)); | ||||
|     } | ||||
|  | ||||
|     fn poll_eof(&mut self, cx: &mut task::Context) -> Poll<Option<Chunk>, ::Error> { | ||||
|         match self.delayed_eof.take() { | ||||
|             Some(DelayEof::NotEof(mut delay)) => { | ||||
|                 match self.poll_inner(cx) { | ||||
|                     ok @ Ok(Async::Ready(Some(..))) | | ||||
|                     ok @ Ok(Async::Pending) => { | ||||
|                         self.delayed_eof = Some(DelayEof::NotEof(delay)); | ||||
|                         ok | ||||
|                     }, | ||||
|                     Ok(Async::Ready(None)) => match delay.poll(cx) { | ||||
|                         Ok(Async::Ready(never)) => match never {}, | ||||
|                         Ok(Async::Pending) => { | ||||
|                             self.delayed_eof = Some(DelayEof::Eof(delay)); | ||||
|                             Ok(Async::Pending) | ||||
|                         }, | ||||
|                         Err(_done) => { | ||||
|                             Ok(Async::Ready(None)) | ||||
|                         }, | ||||
|                     }, | ||||
|                     Err(e) => Err(e), | ||||
|                 } | ||||
|             }, | ||||
|             Some(DelayEof::Eof(mut delay)) => { | ||||
|                 match delay.poll(cx) { | ||||
|                     Ok(Async::Ready(never)) => match never {}, | ||||
|                     Ok(Async::Pending) => { | ||||
|                         self.delayed_eof = Some(DelayEof::Eof(delay)); | ||||
|                         Ok(Async::Pending) | ||||
|                     }, | ||||
|                     Err(_done) => { | ||||
|                         Ok(Async::Ready(None)) | ||||
|                     }, | ||||
|                 } | ||||
|             }, | ||||
|             None => self.poll_inner(cx), | ||||
|         } | ||||
|     } | ||||
|  | ||||
|     fn poll_inner(&mut self, cx: &mut task::Context) -> Poll<Option<Chunk>, ::Error> { | ||||
|         match self.kind { | ||||
|             Kind::Chan { ref mut rx, .. } => match rx.poll_next(cx).expect("mpsc cannot error") { | ||||
|                 Async::Ready(Some(Ok(chunk))) => Ok(Async::Ready(Some(chunk))), | ||||
|                 Async::Ready(Some(Err(err))) => Err(err), | ||||
|                 Async::Ready(None) => Ok(Async::Ready(None)), | ||||
|                 Async::Pending => Ok(Async::Pending), | ||||
|             }, | ||||
|             Kind::Wrapped(ref mut s) => s.poll_next(cx), | ||||
|             Kind::Once(ref mut val) => Ok(Async::Ready(val.take())), | ||||
|             Kind::Empty => Ok(Async::Ready(None)), | ||||
|         } | ||||
|     } | ||||
| } | ||||
| @@ -269,17 +346,7 @@ impl Entity for Body { | ||||
|     type Error = ::Error; | ||||
|  | ||||
|     fn poll_data(&mut self, cx: &mut task::Context) -> Poll<Option<Self::Data>, Self::Error> { | ||||
|         match self.kind { | ||||
|             Kind::Chan { ref mut rx, .. } => match rx.poll_next(cx).expect("mpsc cannot error") { | ||||
|                 Async::Ready(Some(Ok(chunk))) => Ok(Async::Ready(Some(chunk))), | ||||
|                 Async::Ready(Some(Err(err))) => Err(err), | ||||
|                 Async::Ready(None) => Ok(Async::Ready(None)), | ||||
|                 Async::Pending => Ok(Async::Pending), | ||||
|             }, | ||||
|             Kind::Wrapped(ref mut s) => s.poll_next(cx), | ||||
|             Kind::Once(ref mut val) => Ok(Async::Ready(val.take())), | ||||
|             Kind::Empty => Ok(Async::Ready(None)), | ||||
|         } | ||||
|         self.poll_eof(cx) | ||||
|     } | ||||
|  | ||||
|     fn is_end_stream(&self) -> bool { | ||||
| @@ -300,6 +367,7 @@ impl Entity for Body { | ||||
|             Kind::Empty => Some(0) | ||||
|         } | ||||
|     } | ||||
|      | ||||
| } | ||||
|  | ||||
| impl fmt::Debug for Body { | ||||
|   | ||||
		Reference in New Issue
	
	Block a user