ProjectHentai/nonebot2
6
0
Fork 0
mirror of https://github.com/nonebot/nonebot2.git synced 2025-02-08 03:46:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

⚡ improve aiohttp main loop performance

Browse Source
This commit is contained in:
yanyongyu 2021-07-21 23:47:18 +08:00
parent cb68e0ee1b
commit 1f0be1e1e3
1 changed files with 14 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
nonebot/drivers/aiohttp.py
View File

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