fix(client): properly close idle connections after timeout
Additionally fixes if there were idle connections when a `Client` is dropped. Only fixes with the no-proto dispatcher, as changing internals for the tokio-proto dispatcher would be much harder, and it will replace it very soon. Closes #1397
This commit is contained in:
		| @@ -249,8 +249,12 @@ where C: Connect, | ||||
|                     let pool_key = Rc::new(domain.to_string()); | ||||
|                     self.connector.connect(url) | ||||
|                         .map(move |io| { | ||||
|                             let (tx, rx) = mpsc::channel(1); | ||||
|                             let pooled = pool.pooled(pool_key, RefCell::new(tx)); | ||||
|                             let (tx, rx) = mpsc::channel(0); | ||||
|                             let tx = HyperClient { | ||||
|                                 tx: RefCell::new(tx), | ||||
|                                 should_close: true, | ||||
|                             }; | ||||
|                             let pooled = pool.pooled(pool_key, tx); | ||||
|                             let conn = proto::Conn::<_, _, proto::ClientTransaction, _>::new(io, pooled.clone()); | ||||
|                             let dispatch = proto::dispatch::Dispatcher::new(proto::dispatch::Client::new(rx), conn); | ||||
|                             handle.spawn(dispatch.map_err(|err| error!("no_proto error: {}", err))); | ||||
| @@ -269,9 +273,10 @@ where C: Connect, | ||||
|                         e.into() | ||||
|                     }); | ||||
|  | ||||
|                 let resp = race.and_then(move |client| { | ||||
|                 let resp = race.and_then(move |mut client| { | ||||
|                     let (callback, rx) = oneshot::channel(); | ||||
|                     client.borrow_mut().start_send((head, body, callback)).unwrap(); | ||||
|                     client.tx.borrow_mut().start_send(proto::dispatch::ClientMsg::Request(head, body, callback)).unwrap(); | ||||
|                     client.should_close = false; | ||||
|                     rx.then(|res| { | ||||
|                         match res { | ||||
|                             Ok(Ok(res)) => Ok(res), | ||||
| @@ -309,7 +314,29 @@ impl<C, B> fmt::Debug for Client<C, B> { | ||||
| } | ||||
|  | ||||
| type ProtoClient<B> = ClientProxy<Message<RequestHead, B>, Message<proto::ResponseHead, TokioBody>, ::Error>; | ||||
| type HyperClient<B> = RefCell<::futures::sync::mpsc::Sender<(RequestHead, Option<B>, ::futures::sync::oneshot::Sender<::Result<::Response>>)>>; | ||||
|  | ||||
| struct HyperClient<B> { | ||||
|     tx: RefCell<::futures::sync::mpsc::Sender<proto::dispatch::ClientMsg<B>>>, | ||||
|     should_close: bool, | ||||
| } | ||||
|  | ||||
| impl<B> Clone for HyperClient<B> { | ||||
|     fn clone(&self) -> HyperClient<B> { | ||||
|         HyperClient { | ||||
|             tx: self.tx.clone(), | ||||
|             should_close: self.should_close, | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| impl<B> Drop for HyperClient<B> { | ||||
|     fn drop(&mut self) { | ||||
|         if self.should_close { | ||||
|             self.should_close = false; | ||||
|             let _ = self.tx.borrow_mut().try_send(proto::dispatch::ClientMsg::Close); | ||||
|         } | ||||
|     } | ||||
| } | ||||
|  | ||||
| enum Dispatch<B> { | ||||
|     Proto(Pool<ProtoClient<B>>), | ||||
|   | ||||
| @@ -511,6 +511,11 @@ where I: AsyncRead + AsyncWrite, | ||||
|  | ||||
|     } | ||||
|  | ||||
|     pub fn close_and_shutdown(&mut self) -> Poll<(), io::Error> { | ||||
|         try_ready!(self.flush()); | ||||
|         self.shutdown() | ||||
|     } | ||||
|  | ||||
|     pub fn shutdown(&mut self) -> Poll<(), io::Error> { | ||||
|         match self.io.io_mut().shutdown() { | ||||
|             Ok(Async::NotReady) => Ok(Async::NotReady), | ||||
| @@ -625,8 +630,7 @@ where I: AsyncRead + AsyncWrite, | ||||
|  | ||||
|     #[inline] | ||||
|     fn close(&mut self) -> Poll<(), Self::SinkError> { | ||||
|         try_ready!(self.poll_complete()); | ||||
|         self.shutdown() | ||||
|         self.close_and_shutdown() | ||||
|     } | ||||
| } | ||||
|  | ||||
|   | ||||
| @@ -13,6 +13,7 @@ pub struct Dispatcher<D, Bs, I, B, T, K> { | ||||
|     dispatch: D, | ||||
|     body_tx: Option<super::body::BodySender>, | ||||
|     body_rx: Option<Bs>, | ||||
|     is_closing: bool, | ||||
| } | ||||
|  | ||||
| pub trait Dispatch { | ||||
| @@ -34,7 +35,12 @@ pub struct Client<B> { | ||||
|     rx: ClientRx<B>, | ||||
| } | ||||
|  | ||||
| type ClientRx<B> = mpsc::Receiver<(RequestHead, Option<B>, oneshot::Sender<::Result<::Response>>)>; | ||||
| pub enum ClientMsg<B> { | ||||
|     Request(RequestHead, Option<B>, oneshot::Sender<::Result<::Response>>), | ||||
|     Close, | ||||
| } | ||||
|  | ||||
| type ClientRx<B> = mpsc::Receiver<ClientMsg<B>>; | ||||
|  | ||||
| impl<D, Bs, I, B, T, K> Dispatcher<D, Bs, I, B, T, K> | ||||
| where | ||||
| @@ -51,6 +57,7 @@ where | ||||
|             dispatch: dispatch, | ||||
|             body_tx: None, | ||||
|             body_rx: None, | ||||
|             is_closing: false, | ||||
|         } | ||||
|     } | ||||
|  | ||||
| @@ -60,7 +67,9 @@ where | ||||
|  | ||||
|     fn poll_read(&mut self) -> Poll<(), ::Error> { | ||||
|         loop { | ||||
|             if self.conn.can_read_head() { | ||||
|             if self.is_closing { | ||||
|                 return Ok(Async::Ready(())); | ||||
|             } else if self.conn.can_read_head() { | ||||
|                 match self.conn.read_head() { | ||||
|                     Ok(Async::Ready(Some((head, has_body)))) => { | ||||
|                         let body = if has_body { | ||||
| @@ -149,12 +158,16 @@ where | ||||
|  | ||||
|     fn poll_write(&mut self) -> Poll<(), ::Error> { | ||||
|         loop { | ||||
|             if self.body_rx.is_none() && self.dispatch.should_poll() { | ||||
|             if self.is_closing { | ||||
|                 return Ok(Async::Ready(())); | ||||
|             } else if self.body_rx.is_none() && self.dispatch.should_poll() { | ||||
|                 if let Some((head, body)) = try_ready!(self.dispatch.poll_msg()) { | ||||
|                     self.conn.write_head(head, body.is_some()); | ||||
|                     self.body_rx = body; | ||||
|                 } else { | ||||
|                     self.conn.close_write(); | ||||
|                     self.is_closing = true; | ||||
|                     //self.conn.close_read(); | ||||
|                     //self.conn.close_write(); | ||||
|                     return Ok(Async::Ready(())); | ||||
|                 } | ||||
|             } else if self.conn.has_queued_body() { | ||||
| @@ -190,6 +203,16 @@ where | ||||
|         }) | ||||
|     } | ||||
|  | ||||
|     fn poll_close(&mut self) -> Poll<(), ::Error> { | ||||
|         debug_assert!(self.is_closing); | ||||
|  | ||||
|         try_ready!(self.conn.close_and_shutdown()); | ||||
|         self.conn.close_read(); | ||||
|         self.conn.close_write(); | ||||
|         self.is_closing = false; | ||||
|         Ok(Async::Ready(())) | ||||
|     } | ||||
|  | ||||
|     fn is_done(&self) -> bool { | ||||
|         let read_done = self.conn.is_read_closed(); | ||||
|  | ||||
| @@ -224,6 +247,10 @@ where | ||||
|         self.poll_write()?; | ||||
|         self.poll_flush()?; | ||||
|  | ||||
|         if self.is_closing { | ||||
|             self.poll_close()?; | ||||
|         } | ||||
|  | ||||
|         if self.is_done() { | ||||
|             try_ready!(self.conn.shutdown()); | ||||
|             trace!("Dispatch::poll done"); | ||||
| @@ -285,6 +312,7 @@ where | ||||
|  | ||||
| // ===== impl Client ===== | ||||
|  | ||||
|  | ||||
| impl<B> Client<B> { | ||||
|     pub fn new(rx: ClientRx<B>) -> Client<B> { | ||||
|         Client { | ||||
| @@ -305,11 +333,13 @@ where | ||||
|  | ||||
|     fn poll_msg(&mut self) -> Poll<Option<(Self::PollItem, Option<Self::PollBody>)>, ::Error> { | ||||
|         match self.rx.poll() { | ||||
|             Ok(Async::Ready(Some((head, body, cb)))) => { | ||||
|             Ok(Async::Ready(Some(ClientMsg::Request(head, body, cb)))) => { | ||||
|                 self.callback = Some(cb); | ||||
|                 Ok(Async::Ready(Some((head, body)))) | ||||
|             }, | ||||
|             Ok(Async::Ready(Some(ClientMsg::Close))) | | ||||
|             Ok(Async::Ready(None)) => { | ||||
|                 trace!("client tx closed"); | ||||
|                 // user has dropped sender handle | ||||
|                 Ok(Async::Ready(None)) | ||||
|             }, | ||||
|   | ||||
		Reference in New Issue
	
	Block a user