From 431e455642cee8f7258814d6bb9b426a99d1256b Mon Sep 17 00:00:00 2001 From: "Kevin (Kun) \"Kassimo\" Qian" Date: Tue, 15 Jan 2019 17:36:51 -0800 Subject: Kill all pending accepts when TCP listener is closed (#1517) --- src/tokio_util.rs | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) (limited to 'src/tokio_util.rs') diff --git a/src/tokio_util.rs b/src/tokio_util.rs index 32542aa43..2eb0211db 100644 --- a/src/tokio_util.rs +++ b/src/tokio_util.rs @@ -63,7 +63,25 @@ impl Future for Accept { fn poll(&mut self) -> Poll { let (stream, addr) = match self.state { - AcceptState::Pending(ref mut r) => try_ready!(r.poll_accept()), + // Similar to try_ready!, but also track/untrack accept task + // in TcpListener resource. + // In this way, when the listener is closed, the task can be + // notified to error out (instead of stuck forever). + AcceptState::Pending(ref mut r) => match r.poll_accept() { + Ok(futures::prelude::Async::Ready(t)) => { + r.untrack_task(); + t + } + Ok(futures::prelude::Async::NotReady) => { + // Would error out if another accept task is being tracked. + r.track_task()?; + return Ok(futures::prelude::Async::NotReady); + } + Err(e) => { + r.untrack_task(); + return Err(From::from(e)); + } + }, AcceptState::Empty => panic!("poll Accept after it's done"), }; -- cgit v1.2.3