diff --git a/src/core/Client.py b/src/core/Client.py index b2390c9..364081a 100644 --- a/src/core/Client.py +++ b/src/core/Client.py @@ -754,11 +754,13 @@ class Client: try: if not self.__writer.is_closing(): self.__writer.close() + await self.__writer.wait_closed() except Exception as e: self.log.debug(f"Error while closing writer: {e}") try: _, down_w = self._down_sock if down_w and not down_w.is_closing(): down_w.close() + await down_w.wait_closed() except Exception as e: self.log.debug(f"Error while closing download writer: {e}") diff --git a/src/core/tcp_server.py b/src/core/tcp_server.py index cdf95a3..8f48cb1 100644 --- a/src/core/tcp_server.py +++ b/src/core/tcp_server.py @@ -22,6 +22,8 @@ class TCPServer: self.host = host self.port = port self.run = False + self._connections = set() + self.server = None self.rl = RateLimiter(50, 10, 300) console.add_command("rl", self.rl.parse_console, None, "RateLimiter menu", {"rl": {"info": None, "unban": None, "ban": None, "help": None}}) @@ -135,7 +137,8 @@ class TCPServer: return False, None async def handle_client(self, reader, writer): - while True: + while self.run: + self._connections.add(writer) try: ip = writer.get_extra_info('peername')[0] if self.rl.is_banned(ip): @@ -163,12 +166,12 @@ class TCPServer: self.log.debug("Starting TCP server.") self.run = True try: - server = await asyncio.start_server(self.handle_client, self.host, self.port, + self.server = await asyncio.start_server(self.handle_client, self.host, self.port, backlog=int(config.Game["players"] * 2.3)) - self.log.debug(f"TCP server started on {server.sockets[0].getsockname()!r}") + self.log.debug(f"TCP server started on {self.server.sockets[0].getsockname()!r}") while True: - async with server: - await server.serve_forever() + async with self.server: + await self.server.serve_forever() except OSError as e: self.log.error(i18n.core_bind_failed.format(e)) raise e @@ -181,5 +184,14 @@ class TCPServer: self.run = False self.Core.run = False - def stop(self): + async def stop(self): self.log.debug("Stopping TCP server") + try: + self.server.close() + for conn in self._connections: + conn.close() + await conn.wait_closed() + await self.server.wait_closed() + except Exception as e: + self.log.exception(e) + self.log.debug("Stopped") diff --git a/src/core/tcp_server.pyi b/src/core/tcp_server.pyi index c223d34..dd88c52 100644 --- a/src/core/tcp_server.pyi +++ b/src/core/tcp_server.pyi @@ -15,11 +15,13 @@ from modules import RateLimiter class TCPServer: def __init__(self, core: Core, host, port): + self.server = await asyncio.start_server(self.handle_client, "", 0, backlog=int(config.Game["players"] * 2.3)) self.log = utils.get_logger("TCPServer") self.loop = asyncio.get_event_loop() self.Core = core self.host = host self.port = port + self._connections = set() self.run = False self.rl = RateLimiter(50, 10, 15)