diff --git a/nonebot/drivers/aiohttp.py b/nonebot/drivers/aiohttp.py index 7bff8cc1..1871eec1 100644 --- a/nonebot/drivers/aiohttp.py +++ b/nonebot/drivers/aiohttp.py @@ -55,7 +55,7 @@ class Driver(ForwardDriver): self.http_pollings: List[HTTPPollingSetup] = [] self.websockets: List[WebSocketSetup] = [] self.connections: List[asyncio.Task] = [] - self.should_exit: bool = False + self.should_exit: asyncio.Event = asyncio.Event() self.force_exit: bool = False @property @@ -114,7 +114,7 @@ class Driver(ForwardDriver): async def serve(self): self.install_signal_handlers() await self.startup() - if self.should_exit: + if self.should_exit.is_set(): return await self.main_loop() await self.shutdown() @@ -138,8 +138,7 @@ class Driver(ForwardDriver): "Ignored!") async def main_loop(self): - while not self.should_exit: - await asyncio.sleep(0.1) + await self.should_exit.wait() async def shutdown(self): # run shutdown @@ -193,10 +192,10 @@ class Driver(ForwardDriver): signal.signal(sig, self.handle_exit) def handle_exit(self, sig, frame): - if self.should_exit: + if self.should_exit.is_set(): self.force_exit = True else: - self.should_exit = True + self.should_exit.set() async def _http_loop(self, setup: HTTPPollingSetup): url = URL(setup.url) @@ -231,7 +230,7 @@ class Driver(ForwardDriver): async with aiohttp.ClientSession(headers=headers, timeout=timeout, version=version) as session: - while not self.should_exit: + while not self.should_exit.is_set(): logger.debug( f"Bot {setup.self_id} from adapter {setup.adapter} request {url}" ) @@ -253,7 +252,8 @@ class Driver(ForwardDriver): pass except Exception as e: logger.opt(colors=True, exception=e).error( - "Unexpected exception occurred while http polling") + "Unexpected exception occurred " + "while http polling") finally: self._bot_disconnect(bot) @@ -298,7 +298,7 @@ class Driver(ForwardDriver): BotClass = self._adapters[setup.adapter] bot = BotClass(setup.self_id, request) self._bot_connect(bot) - while not self.should_exit: + while not self.should_exit.is_set(): msg = await ws.receive() if msg.type == aiohttp.WSMsgType.text: asyncio.create_task( @@ -316,9 +316,10 @@ class Driver(ForwardDriver): "WebSocket connection closed by peer. " "Try to reconnect...") break - except aiohttp.WSServerHandshakeError as e: + except (aiohttp.ClientResponseError, + aiohttp.ClientConnectionError) as e: logger.opt(colors=True, exception=e).error( - f"Error while connecting to {url}" + f"Error while connecting to {url}. " "Try to reconnect...") finally: if bot: @@ -330,7 +331,8 @@ class Driver(ForwardDriver): pass except Exception as e: logger.opt(colors=True, exception=e).error( - "Unexpected exception occurred while websocket loop") + "Unexpected exception occurred " + "while websocket loop") @dataclass