Issue #25441: asyncio: Raise error from drain() when socket is closed. (Merge 3.4->3.5)

1 files changed, 9 insertions, 0 deletions

diff --git a/Lib/asyncio/streams.py b/Lib/asyncio/streams.py
index bb9fb31..fb786ed 100644
--- a/Lib/asyncio/streams.py
+++ b/Lib/asyncio/streams.py
@@ -301,6 +301,15 @@ class StreamWriter:
             exc = self._reader.exception()
             if exc is not None:
                 raise exc
+        if self._transport is not None:
+            if self._transport._closing:
+                # Yield to the event loop so connection_lost() may be
+                # called.  Without this, _drain_helper() would return
+                # immediately, and code that calls
+                #     write(...); yield from drain()
+                # in a loop would never call connection_lost(), so it
+                # would not see an error when the socket is closed.
+                yield
         yield from self._protocol._drain_helper()