Skip to content
12 changes: 12 additions & 0 deletions Lib/asyncio/proactor_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,18 @@ def _loop_writing(self, fut=None):
addr=addr)
except OSError as exc:
self._protocol.error_received(exc)
if self._buffer or self._closing:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kumaraditya303 this change doesn't make sense without the other changes from #156921

It would be easier if I could get that PR that fixes both issues reviewed in one go

# Reschedule the write loop so buffered data isn't stranded and
# a paused protocol is eventually resumed (gh-156698).
def resume_writing():
# a sendto() may have armed a write in the meantime;
# its own callback will drain the rest of the buffer.
if self._write_fut is None:
self._loop_writing()

self._loop.call_soon(resume_writing)
else:
self._maybe_resume_protocol()
except Exception as exc:
self._fatal_error(exc, 'Fatal write error on datagram transport')
else:
Expand Down
129 changes: 129 additions & 0 deletions Lib/test/test_asyncio/test_events.py
Original file line number Diff line number Diff line change
Expand Up @@ -1583,6 +1583,135 @@ def create_socket():
transport_1.close()
transport_2.close()

def _test_datagram_write_error_resumes_paused_protocol(self, first, second):
# See https://github.com/python/cpython/issues/156698: a
# datagram write error must not strand data left in the write
# buffer, nor leave a paused protocol paused forever.
loop = self.loop

class Protocol(asyncio.DatagramProtocol):
def connection_made(self, transport):
self.transport = transport
self.paused = False
self.resumed = False
self.errors = []
self.error_received_event = loop.create_future()

def pause_writing(self):
self.paused = True

def resume_writing(self):
self.resumed = True

def error_received(self, exc):
self.errors.append(exc)
if not self.error_received_event.done():
self.error_received_event.set_result(None)

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setblocking(False)
sock.bind(('127.0.0.1', 0))
transport, protocol = loop.run_until_complete(
loop.create_datagram_endpoint(Protocol, sock=sock))
addr = sock.getsockname()

# A high water mark of 0 makes pausing deterministic whenever
# anything is left in the write buffer.
transport.set_write_buffer_limits(0)

# The first sendto() may arm an in-flight write, so the second
# one can end up queued behind it; queuing is what trips
# pause_writing() at a high water mark of 0.
transport.sendto(first, addr)
transport.sendto(second, addr)

loop.run_until_complete(
asyncio.wait_for(protocol.error_received_event, 10))
self.assertTrue(protocol.errors)
self.assertIsInstance(protocol.errors[0], OSError)

# The write buffer must not be left stranded.
test_utils.run_until(
loop, lambda: transport.get_write_buffer_size() == 0)

# A protocol that got paused must eventually be resumed too --
# without requiring an unsolicited extra sendto() to un-stick it.
if protocol.paused:
test_utils.run_until(loop, lambda: protocol.resumed)

transport.close()
test_utils.run_briefly(loop)

def test_datagram_write_error_resumes_paused_protocol_in_flight(self):
# oversized datagram fails while in flight; a normal datagram
# queued right behind it must not be stranded.
oversized = b'\x00' * 70000
self._test_datagram_write_error_resumes_paused_protocol(
oversized, b'queued')

def test_datagram_write_error_resumes_paused_protocol_from_callback(self):
# oversized datagram fails once it reaches the front of the
# buffer; the protocol must not stay paused forever.
oversized = b'\x00' * 70000
self._test_datagram_write_error_resumes_paused_protocol(
b'ok', oversized)

def test_datagram_write_error_reentrant_sendto(self):
# See https://github.com/python/cpython/issues/156698: an
# error_received() callback that sends more data synchronously
# can itself arm a new write. The write-loop restart scheduled
# for the failed write must notice that and not try to start a
# second, conflicting one.
loop = self.loop
unhandled = []
loop.set_exception_handler(lambda loop, context: unhandled.append(context))

class Protocol(asyncio.DatagramProtocol):
def connection_made(self, transport):
self.transport = transport
self.sent_extra = False
self.errors = []
self.done = loop.create_future()

def datagram_received(self, data, addr):
if not self.done.done():
self.done.set_result(None)

def error_received(self, exc):
self.errors.append(exc)
if not self.sent_extra:
# Reentrantly kicks off another write while the
# failing one is still unwinding on the stack.
self.sent_extra = True
self.transport.sendto(b'extra', self.addr)

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.setblocking(False)
sock.bind(('127.0.0.1', 0))
transport, protocol = loop.run_until_complete(
loop.create_datagram_endpoint(Protocol, sock=sock))
protocol.addr = addr = sock.getsockname()

oversized = b'\x00' * 70000
transport.sendto(oversized, addr)
transport.sendto(b'queued', addr)

# The 'extra' datagram sent from error_received() is delivered
# back to the same socket; waiting for it proves the write loop
# kept running instead of wedging or crashing.
loop.run_until_complete(asyncio.wait_for(protocol.done, 10))

test_utils.run_until(
loop, lambda: transport.get_write_buffer_size() == 0)

transport.close()
test_utils.run_briefly(loop)

self.assertTrue(protocol.errors)
self.assertFalse(
unhandled,
f'unhandled exception in the write loop: {unhandled}')

def test_datagram_recvfrom_connection_reset_recovers(self):
# gh-127057: a UDP socket that sent a datagram to an address that
# wasn't listening can raise ConnectionResetError on a later
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
Fix :class:`asyncio.ProactorEventLoop` UDP transports so that a write
error no longer strands a paused protocol: the write loop is now
rescheduled when data remains buffered, and the protocol is resumed
when the buffer has drained.
Loading