@@ -1454,6 +1454,264 @@ def create_socket():
14541454 transport_1 .close ()
14551455 transport_2 .close ()
14561456
1457+ def test_datagram_write_error_resumes_paused_protocol (self ):
1458+ # See https://github.com/python/cpython/issues/156698: a
1459+ # datagram write error must not strand data left in the write
1460+ # buffer, nor leave a paused protocol paused forever.
1461+ loop = self .loop
1462+
1463+ class Protocol (asyncio .DatagramProtocol ):
1464+ def connection_made (self , transport ):
1465+ self .transport = transport
1466+ self .paused = False
1467+ self .resumed = False
1468+ self .errors = []
1469+ self .error_received_event = loop .create_future ()
1470+
1471+ def pause_writing (self ):
1472+ self .paused = True
1473+
1474+ def resume_writing (self ):
1475+ self .resumed = True
1476+
1477+ def error_received (self , exc ):
1478+ self .errors .append (exc )
1479+ if not self .error_received_event .done ():
1480+ self .error_received_event .set_result (None )
1481+
1482+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1483+ sock .setblocking (False )
1484+ sock .bind (('127.0.0.1' , 0 ))
1485+ transport , protocol = loop .run_until_complete (
1486+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1487+ addr = sock .getsockname ()
1488+
1489+ # A high water mark of 0 makes pausing deterministic whenever
1490+ # anything is left in the write buffer.
1491+ transport .set_write_buffer_limits (0 )
1492+
1493+ # The oversized datagram fails while it is in flight, and the
1494+ # normal datagram behind it is left queued -- queuing is also
1495+ # what trips pause_writing() at a high water mark of 0.
1496+ transport .sendto (b'\x00 ' * 70000 , addr )
1497+ transport .sendto (b'queued' , addr )
1498+
1499+ loop .run_until_complete (
1500+ asyncio .wait_for (protocol .error_received_event ,
1501+ support .SHORT_TIMEOUT ))
1502+ self .assertTrue (protocol .errors )
1503+ self .assertIsInstance (protocol .errors [0 ], OSError )
1504+
1505+ # The write buffer must not be left stranded.
1506+ test_utils .run_until (
1507+ loop , lambda : transport .get_write_buffer_size () == 0 )
1508+
1509+ # A protocol that got paused must eventually be resumed too --
1510+ # without requiring an unsolicited extra sendto() to un-stick it.
1511+ if protocol .paused :
1512+ test_utils .run_until (loop , lambda : protocol .resumed )
1513+
1514+ transport .close ()
1515+ test_utils .run_briefly (loop )
1516+
1517+ def test_datagram_write_error_reentrant_sendto (self ):
1518+ # See https://github.com/python/cpython/issues/156698: an
1519+ # error_received() callback that sends more data synchronously
1520+ # can itself schedule a new write. The write-loop restart scheduled
1521+ # for the failed write must notice that and not try to start a
1522+ # second, conflicting one.
1523+ loop = self .loop
1524+ unhandled = []
1525+ loop .set_exception_handler (lambda loop , context : unhandled .append (context ))
1526+
1527+ class Protocol (asyncio .DatagramProtocol ):
1528+ def connection_made (self , transport ):
1529+ self .transport = transport
1530+ self .sent_extra = False
1531+ self .errors = []
1532+ self .done = loop .create_future ()
1533+
1534+ def datagram_received (self , data , addr ):
1535+ if not self .done .done ():
1536+ self .done .set_result (None )
1537+
1538+ def error_received (self , exc ):
1539+ self .errors .append (exc )
1540+ if not self .sent_extra :
1541+ # Reentrantly kicks off another write while the
1542+ # failing one is still unwinding on the stack.
1543+ self .sent_extra = True
1544+ self .transport .sendto (b'extra' , self .addr )
1545+
1546+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1547+ sock .setblocking (False )
1548+ sock .bind (('127.0.0.1' , 0 ))
1549+ transport , protocol = loop .run_until_complete (
1550+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1551+ protocol .addr = addr = sock .getsockname ()
1552+
1553+ oversized = b'\x00 ' * 70000
1554+ transport .sendto (oversized , addr )
1555+ transport .sendto (b'queued' , addr )
1556+
1557+ # The 'extra' datagram sent from error_received() is delivered
1558+ # back to the same socket; waiting for it proves the write loop
1559+ # kept running instead of wedging or crashing.
1560+ loop .run_until_complete (
1561+ asyncio .wait_for (protocol .done , support .SHORT_TIMEOUT ))
1562+
1563+ test_utils .run_until (
1564+ loop , lambda : transport .get_write_buffer_size () == 0 )
1565+
1566+ transport .close ()
1567+ test_utils .run_briefly (loop )
1568+
1569+ self .assertTrue (protocol .errors )
1570+ self .assertFalse (
1571+ unhandled ,
1572+ f'unhandled exception in the write loop: { unhandled } )
1573+
1574+ def test_datagram_close_flushes_queued_data (self ):
1575+ # See https://github.com/python/cpython/issues/156920: _conn_lost
1576+ # used to mean "close() was requested" rather than "no more data
1577+ # will be sent". Since add_done_callback() always defers an
1578+ # already-completed write's callback with call_soon(), a sendto()
1579+ # immediately followed by close() -- with no await in between --
1580+ # leaves a write genuinely outstanding at close() time on every
1581+ # platform, not just a slow one. Closing must let that write (and
1582+ # anything queued behind it) drain and still call connection_lost(),
1583+ # instead of tripping the "no more data will be sent" guard before
1584+ # the drain has actually happened and hanging forever.
1585+ loop = self .loop
1586+
1587+ class Receiver (asyncio .DatagramProtocol ):
1588+ def connection_made (self , transport ):
1589+ self .received = []
1590+
1591+ def datagram_received (self , data , addr ):
1592+ self .received .append (data )
1593+
1594+ recv_sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1595+ recv_sock .setblocking (False )
1596+ recv_sock .bind (('127.0.0.1' , 0 ))
1597+ recv_transport , receiver = loop .run_until_complete (
1598+ loop .create_datagram_endpoint (Receiver , sock = recv_sock ))
1599+ addr = recv_sock .getsockname ()
1600+
1601+ class Protocol (asyncio .DatagramProtocol ):
1602+ def connection_made (self , transport ):
1603+ self .lost = loop .create_future ()
1604+
1605+ def connection_lost (self , exc ):
1606+ if not self .lost .done ():
1607+ self .lost .set_result (exc )
1608+
1609+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1610+ sock .setblocking (False )
1611+ sock .bind (('127.0.0.1' , 0 ))
1612+ transport , protocol = loop .run_until_complete (
1613+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1614+
1615+ # 'first' is still in flight (its completion callback hasn't run
1616+ # yet) and 'second' is queued behind it when close() is called.
1617+ transport .sendto (b'first' , addr )
1618+ transport .sendto (b'second' , addr )
1619+ transport .close ()
1620+
1621+ loop .run_until_complete (
1622+ asyncio .wait_for (protocol .lost , support .SHORT_TIMEOUT ))
1623+
1624+ test_utils .run_until (
1625+ loop , lambda : len (receiver .received ) >= 2 )
1626+ self .assertEqual (sorted (receiver .received ), [b'first' , b'second' ])
1627+
1628+ recv_transport .close ()
1629+ test_utils .run_briefly (loop )
1630+
1631+ def test_datagram_close_during_write_error_calls_connection_lost (self ):
1632+ # See https://github.com/python/cpython/issues/156920: if the
1633+ # write that's outstanding when close() is called goes on to fail
1634+ # (rather than succeed), the failure handler used to only re-schedule
1635+ # the write loop when data was still queued behind it. If that
1636+ # failing write was the last thing in the buffer, nothing re-scheduled
1637+ # the loop, so the close() in progress never got to call
1638+ # connection_lost() -- it hung forever instead of finishing once
1639+ # the buffer was actually empty.
1640+ loop = self .loop
1641+
1642+ class Protocol (asyncio .DatagramProtocol ):
1643+ def connection_made (self , transport ):
1644+ self .lost = loop .create_future ()
1645+ self .errors = []
1646+
1647+ def error_received (self , exc ):
1648+ self .errors .append (exc )
1649+
1650+ def connection_lost (self , exc ):
1651+ if not self .lost .done ():
1652+ self .lost .set_result (exc )
1653+
1654+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1655+ sock .setblocking (False )
1656+ sock .bind (('127.0.0.1' , 0 ))
1657+ transport , protocol = loop .run_until_complete (
1658+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1659+ addr = sock .getsockname ()
1660+
1661+ # 'ok' is still in flight when close() is called; 'oversized' is
1662+ # queued behind it and fails once it reaches the front of the
1663+ # buffer, leaving the buffer empty right as the error is handled.
1664+ oversized = b'\x00 ' * 70000
1665+ transport .sendto (b'ok' , addr )
1666+ transport .sendto (oversized , addr )
1667+ transport .close ()
1668+
1669+ loop .run_until_complete (
1670+ asyncio .wait_for (protocol .lost , support .SHORT_TIMEOUT ))
1671+ self .assertTrue (protocol .errors )
1672+
1673+ def test_datagram_write_error_close_from_callback (self ):
1674+ # See https://github.com/python/cpython/issues/156920: an
1675+ # error_received() callback that closes the transport must still
1676+ # result in connection_lost() being called eventually, instead of
1677+ # leaving the transport hanging forever. Two failing writes are
1678+ # used so that the first failure's error_received() call closes
1679+ # the transport while the second is still queued (close() defers
1680+ # to the write loop), and the second failure then empties the
1681+ # buffer with self._closing already True and no write in flight
1682+ # -- exercising the `self._closing` half of the
1683+ # `if self._buffer or self._closing:` condition in _loop_writing.
1684+ loop = self .loop
1685+
1686+ class Protocol (asyncio .DatagramProtocol ):
1687+ def connection_made (self , transport ):
1688+ self .transport = transport
1689+ self .errors = []
1690+ self .lost = loop .create_future ()
1691+
1692+ def error_received (self , exc ):
1693+ self .errors .append (exc )
1694+ self .transport .close ()
1695+
1696+ def connection_lost (self , exc ):
1697+ if not self .lost .done ():
1698+ self .lost .set_result (exc )
1699+
1700+ sock = socket .socket (socket .AF_INET , socket .SOCK_DGRAM )
1701+ sock .setblocking (False )
1702+ sock .bind (('127.0.0.1' , 0 ))
1703+ transport , protocol = loop .run_until_complete (
1704+ loop .create_datagram_endpoint (Protocol , sock = sock ))
1705+ addr = sock .getsockname ()
1706+
1707+ oversized = b'\x00 ' * 70000
1708+ transport .sendto (oversized , addr )
1709+ transport .sendto (oversized , addr )
1710+
1711+ loop .run_until_complete (
1712+ asyncio .wait_for (protocol .lost , support .SHORT_TIMEOUT ))
1713+ self .assertEqual (len (protocol .errors ), 2 )
1714+
14571715 def test_datagram_recvfrom_connection_reset_recovers (self ):
14581716 # gh-127057: a UDP socket that sent a datagram to an address that
14591717 # wasn't listening can raise ConnectionResetError on a later
0 commit comments