gh-113538: Add asycio.Server.{close,abort}_clients (#114432)

These give applications the option of more forcefully terminating client
connections for asyncio servers. Useful when terminating a service and
there is limited time to wait for clients to finish up their work.

1 files changed, 88 insertions, 8 deletions

diff --git a/Lib/test/test_asyncio/test_server.py b/Lib/test/test_asyncio/test_server.py
index 918faac909b..0c55661bfb8 100644
--- a/Lib/test/test_asyncio/test_server.py
+++ b/Lib/test/test_asyncio/test_server.py
@@ -125,8 +125,12 @@ class SelectorStartServerTests(BaseStartServer, unittest.TestCase):
 class TestServer2(unittest.IsolatedAsyncioTestCase):
 
     async def test_wait_closed_basic(self):
-        async def serve(*args):
-            pass
+        async def serve(rd, wr):
+            try:
+                await rd.read()
+            finally:
+                wr.close()
+                await wr.wait_closed()
 
         srv = await asyncio.start_server(serve, socket_helper.HOSTv4, 0)
         self.addCleanup(srv.close)
@@ -137,7 +141,8 @@ class TestServer2(unittest.IsolatedAsyncioTestCase):
         self.assertFalse(task1.done())
 
         # active count != 0, not closed: should block
-        srv._attach()
+        addr = srv.sockets[0].getsockname()
+        (rd, wr) = await asyncio.open_connection(addr[0], addr[1])
         task2 = asyncio.create_task(srv.wait_closed())
         await asyncio.sleep(0)
         self.assertFalse(task1.done())
@@ -152,7 +157,8 @@ class TestServer2(unittest.IsolatedAsyncioTestCase):
         self.assertFalse(task2.done())
         self.assertFalse(task3.done())
 
-        srv._detach()
+        wr.close()
+        await wr.wait_closed()
         # active count == 0, closed: should unblock
         await task1
         await task2
@@ -161,8 +167,12 @@ class TestServer2(unittest.IsolatedAsyncioTestCase):
 
     async def test_wait_closed_race(self):
         # Test a regression in 3.12.0, should be fixed in 3.12.1
-        async def serve(*args):
-            pass
+        async def serve(rd, wr):
+            try:
+                await rd.read()
+            finally:
+                wr.close()
+                await wr.wait_closed()
 
         srv = await asyncio.start_server(serve, socket_helper.HOSTv4, 0)
         self.addCleanup(srv.close)
@@ -170,13 +180,83 @@ class TestServer2(unittest.IsolatedAsyncioTestCase):
         task = asyncio.create_task(srv.wait_closed())
         await asyncio.sleep(0)
         self.assertFalse(task.done())
-        srv._attach()
+        addr = srv.sockets[0].getsockname()
+        (rd, wr) = await asyncio.open_connection(addr[0], addr[1])
         loop = asyncio.get_running_loop()
         loop.call_soon(srv.close)
-        loop.call_soon(srv._detach)
+        loop.call_soon(wr.close)
         await srv.wait_closed()
 
+    async def test_close_clients(self):
+        async def serve(rd, wr):
+            try:
+                await rd.read()
+            finally:
+                wr.close()
+                await wr.wait_closed()
+
+        srv = await asyncio.start_server(serve, socket_helper.HOSTv4, 0)
+        self.addCleanup(srv.close)
+
+        addr = srv.sockets[0].getsockname()
+        (rd, wr) = await asyncio.open_connection(addr[0], addr[1])
+        self.addCleanup(wr.close)
+
+        task = asyncio.create_task(srv.wait_closed())
+        await asyncio.sleep(0)
+        self.assertFalse(task.done())
+
+        srv.close()
+        srv.close_clients()
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        self.assertTrue(task.done())
+
+    async def test_abort_clients(self):
+        async def serve(rd, wr):
+            nonlocal s_rd, s_wr
+            s_rd = rd
+            s_wr = wr
+            await wr.wait_closed()
+
+        s_rd = s_wr = None
+        srv = await asyncio.start_server(serve, socket_helper.HOSTv4, 0)
+        self.addCleanup(srv.close)
+
+        addr = srv.sockets[0].getsockname()
+        (c_rd, c_wr) = await asyncio.open_connection(addr[0], addr[1], limit=4096)
+        self.addCleanup(c_wr.close)
+
+        # Limit the socket buffers so we can reliably overfill them
+        s_sock = s_wr.get_extra_info('socket')
+        s_sock.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, 65536)
+        c_sock = c_wr.get_extra_info('socket')
+        c_sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 65536)
+
+        # Get the reader in to a paused state by sending more than twice
+        # the configured limit
+        s_wr.write(b'a' * 4096)
+        s_wr.write(b'a' * 4096)
+        s_wr.write(b'a' * 4096)
+        while c_wr.transport.is_reading():
+            await asyncio.sleep(0)
+
+        # Get the writer in a waiting state by sending data until the
+        # socket buffers are full on both server and client sockets and
+        # the kernel stops accepting more data
+        s_wr.write(b'a' * c_sock.getsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF))
+        s_wr.write(b'a' * s_sock.getsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF))
+        self.assertNotEqual(s_wr.transport.get_write_buffer_size(), 0)
+
+        task = asyncio.create_task(srv.wait_closed())
+        await asyncio.sleep(0)
+        self.assertFalse(task.done())
 
+        srv.close()
+        srv.abort_clients()
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        self.assertTrue(task.done())
 
 
 # Test the various corner cases of Unix server socket removal