gh-156920: fix ProactorEventLoop datagram transports drop buffered datagrams on close() and never call connection_lost()

Author: graingert

Lib/asyncio/proactor_events.py

@@ -105,8 +105,9 @@ def close(self):
         if self._closing:
             return
         self._closing = True
-        self._conn_lost += 1
         if not self._buffer and self._write_fut is None:
+            # Nothing left to flush: no more data will be sent.
+            self._conn_lost += 1
             self._loop.call_soon(self._call_connection_lost, None)
         if self._read_fut is not None:
             self._read_fut.cancel()
@@ -386,6 +387,7 @@ def _loop_writing(self, f=None, data=None):
                 self._buffer = None
             if not data:
                 if self._closing:
+                    self._conn_lost += 1
                     self._loop.call_soon(self._call_connection_lost, None)
                 if self._eof_written:
                     self._sock.shutdown(socket.SHUT_WR)
@@ -480,6 +482,11 @@ def get_write_buffer_size(self):
     def abort(self):
         self._force_close(None)
 
+    def _force_close(self, exc):
+        # The base class drops the buffer; the size is tracked separately.
+        self._buffer_size = 0
+        super()._force_close(exc)
+
     def sendto(self, data, addr=None):
         if not isinstance(data, (bytes, bytearray, memoryview)):
             raise TypeError('data argument must be bytes-like object (%r)',
@@ -509,6 +516,8 @@ def sendto(self, data, addr=None):
     def _loop_writing(self, fut=None):
         try:
             if self._conn_lost:
+                # No more data will be sent: either everything buffered has
+                # already been flushed, or _force_close() dropped it.
                 return
 
             assert fut is self._write_fut
@@ -517,9 +526,10 @@ def _loop_writing(self, fut=None):
                 # We are in a _loop_writing() done callback, get the result
                 fut.result()
 
-            if not self._buffer or (self._conn_lost and self._address):
-                # The connection has been closed
+            if not self._buffer:
+                # Everything buffered has been sent
                 if self._closing:
+                    self._conn_lost += 1
                     self._loop.call_soon(self._call_connection_lost, None)
                 return
 
@@ -534,17 +544,26 @@ def _loop_writing(self, fut=None):
                                                               addr=addr)
         except OSError as exc:
             self._protocol.error_received(exc)
-            if self._buffer:
-                # Re-arm 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.
+            # error_received() is arbitrary protocol code: it may have sent
+            # (arming a write of its own, directly or via call_soon()),
+            # closed, or aborted the transport.
+            if self._buffer or self._closing:
+                # Either data is still queued, or a close() is waiting on
+                # the write loop to drain it and call connection_lost().
+                # This write failed, so there is no completion callback
+                # pending to re-enter the loop -- schedule one (gh-156698).
+                def write_next():
+                    # error_received() may have armed a write of its own,
+                    # directly or with call_soon(); its completion callback
+                    # will drain the rest of the buffer.
                     if self._write_fut is None:
                         self._loop_writing()
 
-                self._loop.call_soon(resume_writing)
+                self._loop.call_soon(write_next)
             else:
+                # Nothing left to write, so a paused protocol has to be
+                # resumed here: the next entry into _loop_writing() returns
+                # early on an empty buffer without doing it.
                 self._maybe_resume_protocol()
         except Exception as exc:
             self._fatal_error(exc, 'Fatal write error on datagram transport')

Lib/test/test_asyncio/test_events.py

@@ -1712,6 +1712,103 @@ def error_received(self, exc):
             unhandled,
             f'unhandled exception in the write loop: {unhandled}')
 
+    def test_datagram_close_flushes_queued_data(self):
+        # See https://github.com/python/cpython/issues/156920: _conn_lost
+        # used to mean "close() was requested" rather than "no more data
+        # will be sent". Since add_done_callback() always defers an
+        # already-completed write's callback with call_soon(), a sendto()
+        # immediately followed by close() -- with no await in between --
+        # leaves a write genuinely outstanding at close() time on every
+        # platform, not just a slow one. Closing must let that write (and
+        # anything queued behind it) drain and still call connection_lost(),
+        # instead of tripping the "no more data will be sent" guard before
+        # the drain has actually happened and hanging forever.
+        loop = self.loop
+
+        class Receiver(asyncio.DatagramProtocol):
+            def connection_made(self, transport):
+                self.received = []
+
+            def datagram_received(self, data, addr):
+                self.received.append(data)
+
+        recv_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        recv_sock.setblocking(False)
+        recv_sock.bind(('127.0.0.1', 0))
+        recv_transport, receiver = loop.run_until_complete(
+            loop.create_datagram_endpoint(Receiver, sock=recv_sock))
+        addr = recv_sock.getsockname()
+
+        class Protocol(asyncio.DatagramProtocol):
+            def connection_made(self, transport):
+                self.lost = loop.create_future()
+
+            def connection_lost(self, exc):
+                if not self.lost.done():
+                    self.lost.set_result(exc)
+
+        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))
+
+        # 'first' is still in flight (its completion callback hasn't run
+        # yet) and 'second' is queued behind it when close() is called.
+        transport.sendto(b'first', addr)
+        transport.sendto(b'second', addr)
+        transport.close()
+
+        loop.run_until_complete(asyncio.wait_for(protocol.lost, 10))
+
+        test_utils.run_until(
+            loop, lambda: len(receiver.received) >= 2)
+        self.assertEqual(sorted(receiver.received), [b'first', b'second'])
+
+        recv_transport.close()
+        test_utils.run_briefly(loop)
+
+    def test_datagram_close_during_write_error_calls_connection_lost(self):
+        # See https://github.com/python/cpython/issues/156920: if the
+        # write that's outstanding when close() is called goes on to fail
+        # (rather than succeed), the failure handler used to only re-arm
+        # the write loop when data was still queued behind it. If that
+        # failing write was the last thing in the buffer, nothing re-armed
+        # the loop, so the close() in progress never got to call
+        # connection_lost() -- it hung forever instead of finishing once
+        # the buffer was actually empty.
+        loop = self.loop
+
+        class Protocol(asyncio.DatagramProtocol):
+            def connection_made(self, transport):
+                self.lost = loop.create_future()
+                self.errors = []
+
+            def error_received(self, exc):
+                self.errors.append(exc)
+
+            def connection_lost(self, exc):
+                if not self.lost.done():
+                    self.lost.set_result(exc)
+
+        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()
+
+        # 'ok' is still in flight when close() is called; 'oversized' is
+        # queued behind it and fails once it reaches the front of the
+        # buffer, leaving the buffer empty right as the error is handled.
+        oversized = b'\x00' * 70000
+        transport.sendto(b'ok', addr)
+        transport.sendto(oversized, addr)
+        transport.close()
+
+        loop.run_until_complete(asyncio.wait_for(protocol.lost, 10))
+        self.assertTrue(protocol.errors)
+
     def test_internal_fds(self):
         loop = self.create_event_loop()
         if not isinstance(loop, selector_events.BaseSelectorEventLoop):

Misc/NEWS.d/next/Library/2026-09-04-09-18-05.gh-issue-156920.lONsKT.rst

@@ -0,0 +1,5 @@
+Fix :mod:`asyncio` on Windows: closing a :class:`~asyncio.DatagramTransport`
+under :class:`~asyncio.ProactorEventLoop` while datagrams were still queued,
+or while an in-flight write failed right as ``close()`` was draining the
+buffer, could strand the queued data and never call ``connection_lost()``,
+hanging the close indefinitely.