Flush needs to register listener too

baloo · Dec 29, 2019 · 8ccdfd2 · 8ccdfd2
1 parent 24b4f1b
commit 8ccdfd2
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/src/utils/io/write.rs b/src/utils/io/write.rs
@@ -128,9 +128,16 @@ where
 
     fn poll_flush(
         mut self: Pin<&mut Self>,
-        _cx: &mut task::Context<'_>,
+        cx: &mut task::Context<'_>,
     ) -> Poll<Result<(), Self::Error>> {
-        unsafe { self.get_unchecked_mut() }.flush()
+        let writer = unsafe { self.get_unchecked_mut() };
+        match writer.flush() {
+            Poll::Pending => {
+                writer.listen(cx.waker().clone());
+                Poll::Pending
+            }
+            Poll::Ready(out) => Poll::Ready(out),
+        }
     }
 
     fn poll_close(