From 32104d16c123f53c285c3fd885a993aba6fca8d3 Mon Sep 17 00:00:00 2001 From: ybai08 <271969568+ybai08@users.noreply.github.com> Date: Sun, 18 Oct 2026 00:00:02 -0500 Subject: [PATCH] Add Server.close_clients() and Server.abort_clients() Python 3.13 added close_clients() and abort_clients() to asyncio.Server, to close (or abort) the connections accepted by a server, but uvloop.loop.Server did not have them and raised AttributeError. Keep track of the server's client transports in a WeakSet (as asyncio does) next to the existing active connection count, and implement both methods on top of it. Fixes #683. --- tests/test_tcp.py | 62 ++++++++++++++++++++++++++++++++ uvloop/handles/basetransport.pyx | 4 +-- uvloop/server.pxd | 5 +-- uvloop/server.pyx | 15 ++++++-- 4 files changed, 80 insertions(+), 6 deletions(-) diff --git a/tests/test_tcp.py b/tests/test_tcp.py index b933166b7..9a3a3f1b8 100644 --- a/tests/test_tcp.py +++ b/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. diff --git a/uvloop/handles/basetransport.pyx b/uvloop/handles/basetransport.pyx index 28b307947..f134ca506 100644 --- a/uvloop/handles/basetransport.pyx +++ b/uvloop/handles/basetransport.pyx @@ -174,12 +174,12 @@ cdef class UVBaseTransport(UVSocketHandle): server = self._server if server is not None: - (server)._detach() + (server)._detach(self) self._server = None cdef inline _set_server(self, Server server): self._server = server - (server)._attach() + (server)._attach(self) cdef inline _set_waiter(self, object waiter): if waiter is not None and not isfuture(waiter): diff --git a/uvloop/server.pxd b/uvloop/server.pxd index ef10f8131..e1087f8d4 100644 --- a/uvloop/server.pxd +++ b/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) diff --git a/uvloop/server.pyx b/uvloop/server.pyx index 845bcfda4..f7225e7ea 100644 --- a/uvloop/server.pyx +++ b/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