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:
@@ -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