bpo-36801: Fix waiting in StreamWriter.drain for closing SSL transport (GH-13098)
https://bugs.python.org/issue36801
(cherry picked from commit 1cc0ee7d9f6a2817918fafd24c18d8bb093a85d3)
Co-authored-by: Andrew Svetlov <andrew.svetlov@gmail.com>
diff --git a/Lib/asyncio/streams.py b/Lib/asyncio/streams.py
index d6531f8..50badd6 100644
--- a/Lib/asyncio/streams.py
+++ b/Lib/asyncio/streams.py
@@ -208,6 +208,9 @@
self._drain_waiter = waiter
await waiter
+ def _get_close_waiter(self, stream):
+ raise NotImplementedError
+
class StreamReaderProtocol(FlowControlMixin, protocols.Protocol):
"""Helper class to adapt between Protocol and StreamReader.
@@ -265,6 +268,9 @@
return False
return True
+ def _get_close_waiter(self, stream):
+ return self._closed
+
def __del__(self):
# Prevent reports about unhandled exceptions.
# Better than self._closed._log_traceback = False hack
@@ -320,7 +326,7 @@
return self._transport.is_closing()
async def wait_closed(self):
- await self._protocol._closed
+ await self._protocol._get_close_waiter(self)
def get_extra_info(self, name, default=None):
return self._transport.get_extra_info(name, default)
@@ -338,13 +344,12 @@
if exc is not None:
raise exc
if self._transport.is_closing():
- # Yield to the event loop so connection_lost() may be
- # called. Without this, _drain_helper() would return
- # immediately, and code that calls
- # write(...); await drain()
- # in a loop would never call connection_lost(), so it
- # would not see an error when the socket is closed.
- await sleep(0, loop=self._loop)
+ # Wait for protocol.connection_lost() call
+ # Raise connection closing error if any,
+ # ConnectionResetError otherwise
+ fut = self._protocol._get_close_waiter(self)
+ await fut
+ raise ConnectionResetError('Connection lost')
await self._protocol._drain_helper()