recover: #1517 Kill all pending accepts when TCP listener is closed (#2224)

1 files changed, 19 insertions, 2 deletions

diff --git a/cli/tokio_util.rs b/cli/tokio_util.rs
index e1f8587c3..a57cbbd2e 100644
--- a/cli/tokio_util.rs
+++ b/cli/tokio_util.rs
@@ -78,14 +78,31 @@ pub fn accept(r: Resource) -> Accept {
 pub struct Accept {
   state: AcceptState,
 }
-
 impl Future for Accept {
   type Item = (TcpStream, SocketAddr);
   type Error = io::Error;
 
   fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
     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"),
     };