diff options
Diffstat (limited to 'Lib/asyncio/proactor_events.py')
-rw-r--r-- | Lib/asyncio/proactor_events.py | 52 |
1 files changed, 41 insertions, 11 deletions
diff --git a/Lib/asyncio/proactor_events.py b/Lib/asyncio/proactor_events.py index 877dfb0..337ed0f 100644 --- a/Lib/asyncio/proactor_events.py +++ b/Lib/asyncio/proactor_events.py @@ -30,7 +30,7 @@ class _ProactorBasePipeTransport(transports._FlowControlMixin, super().__init__(extra, loop) self._set_extra(sock) self._sock = sock - self._protocol = protocol + self.set_protocol(protocol) self._server = server self._buffer = None # None or bytearray. self._read_fut = None @@ -159,16 +159,26 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport, def __init__(self, loop, sock, protocol, waiter=None, extra=None, server=None): + self._loop_reading_cb = None + self._paused = True super().__init__(loop, sock, protocol, waiter, extra, server) - self._paused = False + self._reschedule_on_resume = False + self._loop.call_soon(self._loop_reading) + self._paused = False - if protocols._is_buffered_protocol(protocol): - self._loop_reading = self._loop_reading__get_buffer + def set_protocol(self, protocol): + if isinstance(protocol, protocols.BufferedProtocol): + self._loop_reading_cb = self._loop_reading__get_buffer else: - self._loop_reading = self._loop_reading__data_received + self._loop_reading_cb = self._loop_reading__data_received - self._loop.call_soon(self._loop_reading) + super().set_protocol(protocol) + + if self.is_reading(): + # reset reading callback / buffers / self._read_fut + self.pause_reading() + self.resume_reading() def is_reading(self): return not self._paused and not self._closing @@ -179,6 +189,13 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport, self._paused = True if self._read_fut is not None and not self._read_fut.done(): + # TODO: This is an ugly hack to cancel the current read future + # *and* avoid potential race conditions, as read cancellation + # goes through `future.cancel()` and `loop.call_soon()`. + # We then use this special attribute in the reader callback to + # exit *immediately* without doing any cleanup/rescheduling. + self._read_fut.__asyncio_cancelled_on_pause__ = True + self._read_fut.cancel() self._read_fut = None self._reschedule_on_resume = True @@ -210,7 +227,14 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport, if not keep_open: self.close() - def _loop_reading__data_received(self, fut=None): + def _loop_reading(self, fut=None): + self._loop_reading_cb(fut) + + def _loop_reading__data_received(self, fut): + if (fut is not None and + getattr(fut, '__asyncio_cancelled_on_pause__', False)): + return + if self._paused: self._reschedule_on_resume = True return @@ -253,14 +277,18 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport, if not self._closing: raise else: - self._read_fut.add_done_callback(self._loop_reading) + self._read_fut.add_done_callback(self._loop_reading__data_received) finally: if data: self._protocol.data_received(data) elif data == b'': self._loop_reading__on_eof() - def _loop_reading__get_buffer(self, fut=None): + def _loop_reading__get_buffer(self, fut): + if (fut is not None and + getattr(fut, '__asyncio_cancelled_on_pause__', False)): + return + if self._paused: self._reschedule_on_resume = True return @@ -310,7 +338,9 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport, return try: - buf = self._protocol.get_buffer() + buf = self._protocol.get_buffer(-1) + if not len(buf): + raise RuntimeError('get_buffer() returned an empty buffer') except Exception as exc: self._fatal_error( exc, 'Fatal error: protocol.get_buffer() call failed.') @@ -319,7 +349,7 @@ class _ProactorReadPipeTransport(_ProactorBasePipeTransport, try: # schedule a new read self._read_fut = self._loop._proactor.recv_into(self._sock, buf) - self._read_fut.add_done_callback(self._loop_reading) + self._read_fut.add_done_callback(self._loop_reading__get_buffer) except ConnectionAbortedError as exc: if not self._closing: self._fatal_error(exc, 'Fatal read error on pipe transport') |