@@ -1583,6 +1583,264 @@ def create_socket():
15831583 transport_1 .close ()
15841584 transport_2 .close ()
15851585
1586+ def test_datagram_write_error_resumes_paused_protocol (self ):
1587+ # See https://github.com/python/cpython/issues/156698: a
1588+ # datagram write error must not strand data left in the write
1589+ # buffer, nor leave a paused protocol paused forever.
1590+ loop = self .loop
1591+
1592+ class Protocol (asyncio .DatagramProtocol ):
1593+ def connection_made (self , transport ):
1594+ self .transport = transport
1595+ self .paused = False
1596+ self .resumed = False
1597+ self .errors = []
1598+ self .error_received_event = loop .create_future ()
1599+
1600+ def pause_writing (self ):
1601+ self .paused = True
1602+
1603+ def resume_writing (self ):
1604+ self .resumed = True
1605+
1606+ def error_received (self , exc ):
1607+ self .errors .append (exc )
1608+ if not self .error_received_event .done ():
1609+ self .error_received_event .set_result (None )
1610+
1611+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1612+ sock .setblocking (False )
1613+ sock .bind (('127.0.0.1' , 0 ))
1614+ transport , protocol = loop .run_until_complete (
1615+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1616+ addr = sock .getsockname ()
1617+
1618+ # A high water mark of 0 makes pausing deterministic whenever
1619+ # anything is left in the write buffer.
1620+ transport .set_write_buffer_limits (0 )
1621+
1622+ # The oversized datagram fails while it is in flight, and the
1623+ # normal datagram behind it is left queued -- queuing is also
1624+ # what trips pause_writing() at a high water mark of 0.
1625+ transport .sendto (b'\x00 ' * 70000 , addr )
1626+ transport .sendto (b'queued' , addr )
1627+
1628+ loop .run_until_complete (
1629+ asyncio .wait_for (protocol .error_received_event ,
1630+ support .SHORT_TIMEOUT ))
1631+ self .assertTrue (protocol .errors )
1632+ self .assertIsInstance (protocol .errors [0 ], OSError )
1633+
1634+ # The write buffer must not be left stranded.
1635+ test_utils .run_until (
1636+ loop , lambda : transport .get_write_buffer_size () == 0 )
1637+
1638+ # A protocol that got paused must eventually be resumed too --
1639+ # without requiring an unsolicited extra sendto() to un-stick it.
1640+ if protocol .paused :
1641+ test_utils .run_until (loop , lambda : protocol .resumed )
1642+
1643+ transport .close ()
1644+ test_utils .run_briefly (loop )
1645+
1646+ def test_datagram_write_error_reentrant_sendto (self ):
1647+ # See https://github.com/python/cpython/issues/156698: an
1648+ # error_received() callback that sends more data synchronously
1649+ # can itself schedule a new write. The write-loop restart scheduled
1650+ # for the failed write must notice that and not try to start a
1651+ # second, conflicting one.
1652+ loop = self .loop
1653+ unhandled = []
1654+ loop .set_exception_handler (lambda loop , context : unhandled .append (context ))
1655+
1656+ class Protocol (asyncio .DatagramProtocol ):
1657+ def connection_made (self , transport ):
1658+ self .transport = transport
1659+ self .sent_extra = False
1660+ self .errors = []
1661+ self .done = loop .create_future ()
1662+
1663+ def datagram_received (self , data , addr ):
1664+ if not self .done .done ():
1665+ self .done .set_result (None )
1666+
1667+ def error_received (self , exc ):
1668+ self .errors .append (exc )
1669+ if not self .sent_extra :
1670+ # Reentrantly kicks off another write while the
1671+ # failing one is still unwinding on the stack.
1672+ self .sent_extra = True
1673+ self .transport .sendto (b'extra' , self .addr )
1674+
1675+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1676+ sock .setblocking (False )
1677+ sock .bind (('127.0.0.1' , 0 ))
1678+ transport , protocol = loop .run_until_complete (
1679+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1680+ protocol .addr = addr = sock .getsockname ()
1681+
1682+ oversized = b'\x00 ' * 70000
1683+ transport .sendto (oversized , addr )
1684+ transport .sendto (b'queued' , addr )
1685+
1686+ # The 'extra' datagram sent from error_received() is delivered
1687+ # back to the same socket; waiting for it proves the write loop
1688+ # kept running instead of wedging or crashing.
1689+ loop .run_until_complete (
1690+ asyncio .wait_for (protocol .done , support .SHORT_TIMEOUT ))
1691+
1692+ test_utils .run_until (
1693+ loop , lambda : transport .get_write_buffer_size () == 0 )
1694+
1695+ transport .close ()
1696+ test_utils .run_briefly (loop )
1697+
1698+ self .assertTrue (protocol .errors )
1699+ self .assertFalse (
1700+ unhandled ,
1701+ f'unhandled exception in the write loop: { unhandled } )
1702+
1703+ def test_datagram_close_flushes_queued_data (self ):
1704+ # See https://github.com/python/cpython/issues/156920: _conn_lost
1705+ # used to mean "close() was requested" rather than "no more data
1706+ # will be sent". Since add_done_callback() always defers an
1707+ # already-completed write's callback with call_soon(), a sendto()
1708+ # immediately followed by close() -- with no await in between --
1709+ # leaves a write genuinely outstanding at close() time on every
1710+ # platform, not just a slow one. Closing must let that write (and
1711+ # anything queued behind it) drain and still call connection_lost(),
1712+ # instead of tripping the "no more data will be sent" guard before
1713+ # the drain has actually happened and hanging forever.
1714+ loop = self .loop
1715+
1716+ class Receiver (asyncio .DatagramProtocol ):
1717+ def connection_made (self , transport ):
1718+ self .received = []
1719+
1720+ def datagram_received (self , data , addr ):
1721+ self .received .append (data )
1722+
1723+ recv_sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1724+ recv_sock .setblocking (False )
1725+ recv_sock .bind (('127.0.0.1' , 0 ))
1726+ recv_transport , receiver = loop .run_until_complete (
1727+ loop .create_datagram_endpoint (Receiver , sock = recv_sock ))
1728+ addr = recv_sock .getsockname ()
1729+
1730+ class Protocol (asyncio .DatagramProtocol ):
1731+ def connection_made (self , transport ):
1732+ self .lost = loop .create_future ()
1733+
1734+ def connection_lost (self , exc ):
1735+ if not self .lost .done ():
1736+ self .lost .set_result (exc )
1737+
1738+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1739+ sock .setblocking (False )
1740+ sock .bind (('127.0.0.1' , 0 ))
1741+ transport , protocol = loop .run_until_complete (
1742+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1743+
1744+ # 'first' is still in flight (its completion callback hasn't run
1745+ # yet) and 'second' is queued behind it when close() is called.
1746+ transport .sendto (b'first' , addr )
1747+ transport .sendto (b'second' , addr )
1748+ transport .close ()
1749+
1750+ loop .run_until_complete (
1751+ asyncio .wait_for (protocol .lost , support .SHORT_TIMEOUT ))
1752+
1753+ test_utils .run_until (
1754+ loop , lambda : len (receiver .received ) >= 2 )
1755+ self .assertEqual (sorted (receiver .received ), [b'first' , b'second' ])
1756+
1757+ recv_transport .close ()
1758+ test_utils .run_briefly (loop )
1759+
1760+ def test_datagram_close_during_write_error_calls_connection_lost (self ):
1761+ # See https://github.com/python/cpython/issues/156920: if the
1762+ # write that's outstanding when close() is called goes on to fail
1763+ # (rather than succeed), the failure handler used to only re-schedule
1764+ # the write loop when data was still queued behind it. If that
1765+ # failing write was the last thing in the buffer, nothing re-scheduled
1766+ # the loop, so the close() in progress never got to call
1767+ # connection_lost() -- it hung forever instead of finishing once
1768+ # the buffer was actually empty.
1769+ loop = self .loop
1770+
1771+ class Protocol (asyncio .DatagramProtocol ):
1772+ def connection_made (self , transport ):
1773+ self .lost = loop .create_future ()
1774+ self .errors = []
1775+
1776+ def error_received (self , exc ):
1777+ self .errors .append (exc )
1778+
1779+ def connection_lost (self , exc ):
1780+ if not self .lost .done ():
1781+ self .lost .set_result (exc )
1782+
1783+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1784+ sock .setblocking (False )
1785+ sock .bind (('127.0.0.1' , 0 ))
1786+ transport , protocol = loop .run_until_complete (
1787+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1788+ addr = sock .getsockname ()
1789+
1790+ # 'ok' is still in flight when close() is called; 'oversized' is
1791+ # queued behind it and fails once it reaches the front of the
1792+ # buffer, leaving the buffer empty right as the error is handled.
1793+ oversized = b'\x00 ' * 70000
1794+ transport .sendto (b'ok' , addr )
1795+ transport .sendto (oversized , addr )
1796+ transport .close ()
1797+
1798+ loop .run_until_complete (
1799+ asyncio .wait_for (protocol .lost , support .SHORT_TIMEOUT ))
1800+ self .assertTrue (protocol .errors )
1801+
1802+ def test_datagram_write_error_close_from_callback (self ):
1803+ # See https://github.com/python/cpython/issues/156920: an
1804+ # error_received() callback that closes the transport must still
1805+ # result in connection_lost() being called eventually, instead of
1806+ # leaving the transport hanging forever. Two failing writes are
1807+ # used so that the first failure's error_received() call closes
1808+ # the transport while the second is still queued (close() defers
1809+ # to the write loop), and the second failure then empties the
1810+ # buffer with self._closing already True and no write in flight
1811+ # -- exercising the `self._closing` half of the
1812+ # `if self._buffer or self._closing:` condition in _loop_writing.
1813+ loop = self .loop
1814+
1815+ class Protocol (asyncio .DatagramProtocol ):
1816+ def connection_made (self , transport ):
1817+ self .transport = transport
1818+ self .errors = []
1819+ self .lost = loop .create_future ()
1820+
1821+ def error_received (self , exc ):
1822+ self .errors .append (exc )
1823+ self .transport .close ()
1824+
1825+ def connection_lost (self , exc ):
1826+ if not self .lost .done ():
1827+ self .lost .set_result (exc )
1828+
1829+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1830+ sock .setblocking (False )
1831+ sock .bind (('127.0.0.1' , 0 ))
1832+ transport , protocol = loop .run_until_complete (
1833+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1834+ addr = sock .getsockname ()
1835+
1836+ oversized = b'\x00 ' * 70000
1837+ transport .sendto (oversized , addr )
1838+ transport .sendto (oversized , addr )
1839+
1840+ loop .run_until_complete (
1841+ asyncio .wait_for (protocol .lost , support .SHORT_TIMEOUT ))
1842+ self .assertEqual (len (protocol .errors ), 2 )
1843+
15861844 def test_datagram_recvfrom_connection_reset_recovers (self ):
15871845 # gh-127057: a UDP socket that sent a datagram to an address that
15881846 # wasn't listening can raise ConnectionResetError on a later
0 commit comments