Add Server.close_clients() and Server.abort_clients()

tests/test_tcp.py

@@ -640,6 +640,68 @@ async def start_server():
         self.loop.run_until_complete(start_server())
         self.assertEqual(CNT, TOTAL_CNT)
 
+    def _test_server_close_clients(self, method_name):
+        if self.implementation == 'asyncio' and sys.version_info < (3, 13):
+            raise unittest.SkipTest(
+                'asyncio.Server.{}() requires Python 3.13'.format(method_name))
+
+        CNT = 3
+        handlers_started = 0
+        handlers_done = 0
+        all_started = all_done = None
+
+        async def handle_client(reader, writer):
+            nonlocal handlers_started, handlers_done
+            handlers_started += 1
+            if handlers_started == CNT:
+                all_started.set()
+            # Returns once the connection is closed by the server.
+            await reader.read()
+            handlers_done += 1
+            if handlers_done == CNT:
+                all_done.set()
+
+        async def test():
+            nonlocal all_started, all_done
+            all_started = asyncio.Event()
+            all_done = asyncio.Event()
+
+            srv = await asyncio.start_server(
+                handle_client, '127.0.0.1', 0, family=socket.AF_INET)
+            addr = srv.sockets[0].getsockname()
+
+            clients = []
+            for _ in range(CNT):
+                clients.append(await asyncio.open_connection(*addr))
+            await asyncio.wait_for(all_started.wait(), 5)
+
+            srv.close()
+            await asyncio.sleep(0.1)
+            # Closing the server leaves the client connections open.
+            self.assertEqual(handlers_done, 0)
+
+            getattr(srv, method_name)()
+
+            for reader, writer in clients:
+                try:
+                    data = await asyncio.wait_for(reader.read(), 5)
+                except ConnectionResetError:
+                    pass
+                else:
+                    self.assertEqual(data, b'')
+                writer.close()
+
+            await asyncio.wait_for(srv.wait_closed(), 5)
+            await asyncio.wait_for(all_done.wait(), 5)
+
+        self.loop.run_until_complete(test())
+
+    def test_server_close_clients(self):
+        self._test_server_close_clients('close_clients')
+
+    def test_server_abort_clients(self):
+        self._test_server_close_clients('abort_clients')
+
     def test_tcp_handle_exception_in_connection_made(self):
         # Test that if connection_made raises an exception,
         # 'create_connection' still returns.

uvloop/handles/basetransport.pyx

@@ -174,12 +174,12 @@ cdef class UVBaseTransport(UVSocketHandle):
 
             server = self._server
             if server is not None:
-                (<Server>server)._detach()
+                (<Server>server)._detach(self)
                 self._server = None
 
     cdef inline _set_server(self, Server server):
         self._server = server
-        (<Server>server)._attach()
+        (<Server>server)._attach(self)
 
     cdef inline _set_waiter(self, object waiter):
         if waiter is not None and not isfuture(waiter):

uvloop/server.pxd

@@ -3,6 +3,7 @@ cdef class Server:
         list _servers
         list _waiters
         int _active_count
+        object _clients
         Loop _loop
         bint _serving
         object _serving_forever_fut
@@ -12,8 +13,8 @@ cdef class Server:
     cdef _start_serving(self)
     cdef _wakeup(self)
 
-    cdef _attach(self)
-    cdef _detach(self)
+    cdef _attach(self, UVBaseTransport transport)
+    cdef _detach(self, UVBaseTransport transport)
 
     cdef _ref(self)
     cdef _unref(self)

uvloop/server.pyx

@@ -7,6 +7,7 @@ cdef class Server:
         self._servers = []
         self._waiters = []
         self._active_count = 0
+        self._clients = weakref_WeakSet()
         self._serving_forever_fut = None
 
     cdef _add_server(self, UVStreamServer srv):
@@ -29,13 +30,15 @@ cdef class Server:
             if not waiter.done():
                 waiter.set_result(waiter)
 
-    cdef _attach(self):
+    cdef _attach(self, UVBaseTransport transport):
         assert self._servers is not None
         self._active_count += 1
+        self._clients.add(transport)
 
-    cdef _detach(self):
+    cdef _detach(self, UVBaseTransport transport):
         assert self._active_count > 0
         self._active_count -= 1
+        self._clients.discard(transport)
         if self._active_count == 0 and self._servers is None:
             self._wakeup()
 
@@ -93,6 +96,14 @@ cdef class Server:
         finally:
             self._unref()
 
+    def close_clients(self):
+        for transport in list(self._clients):
+            transport.close()
+
+    def abort_clients(self):
+        for transport in list(self._clients):
+            transport.abort()
+
     def is_serving(self):
         return self._serving