调整 notify 的强制退出

konabot/plugins/simple_notify/__init__.py

@@ -1,8 +1,10 @@
-import asyncio
+import asyncio as asynkio
 import datetime
+import functools
 from pathlib import Path
 from typing import Any, Literal, cast
 
+import signal
 import nonebot
 import ptimeparse
 from loguru import logger
@@ -24,7 +26,9 @@ evt = on_message()
 
 (Path(__file__).parent.parent.parent.parent / "data").mkdir(exist_ok=True)
 DATA_FILE_PATH = Path(__file__).parent.parent.parent.parent / "data" / "notify.json"
-DATA_FILE_LOCK = asyncio.Lock()
+DATA_FILE_LOCK = asynkio.Lock()
+
+ASYNK_TASKS: set[asynkio.Task[Any]] = set()
 
 
 class Notify(BaseModel):
@@ -111,7 +115,11 @@ def create_notify_task(notify: Notify, fail2remove: bool = True):
     async def mission():
         begin_time = datetime.datetime.now()
         if begin_time < notify.notify_time:
-            await asyncio.sleep((notify.notify_time - begin_time).total_seconds())
+            try:
+                await asynkio.sleep((notify.notify_time - begin_time).total_seconds())
+            except asynkio.CancelledError:
+                logger.debug("代办提醒被信号中止，任务退出")
+                return
         else:
             logger.warning(
                 f"期望在 {notify.notify_time} 在平台 {notify.platform} {notify.target_env}"
@@ -128,7 +136,7 @@ def create_notify_task(notify: Notify, fail2remove: bool = True):
             DATA_FILE_LOCK.release()
         else:
             pass
-    return asyncio.create_task(mission())
+    return asynkio.create_task(mission())
 
 
 @evt.handle()
@@ -214,11 +222,11 @@ async def _():
 
     DELTA = 2
     logger.info(f"第一次探测到 Bot 连接，等待 {DELTA} 秒后开始通知")
-    await asyncio.sleep(DELTA)
+    await asynkio.sleep(DELTA)
 
     await DATA_FILE_LOCK.acquire()
 
-    tasks: set[asyncio.Task[Any]] = set()
+    # tasks: set[asynkio.Task[Any]] = set()
     cfg = load_notify_config()
     if cfg.version == 1:
         logger.info("将配置文件的版本升级为 2")
@@ -227,11 +235,26 @@ async def _():
         counter = 0
         for notify in [*cfg.notifies]:
             task = create_notify_task(notify, fail2remove=False)
-            tasks.add(task)
-            task.add_done_callback(lambda self: tasks.remove(self))
+            ASYNK_TASKS.add(task)
+            task.add_done_callback(lambda self: ASYNK_TASKS.remove(self))
             counter += 1
         logger.info(f"成功创建了 {counter} 条代办事项")
     save_notify_config(cfg)
     DATA_FILE_LOCK.release()
 
-    await asyncio.gather(*tasks)
+    loop = asynkio.get_running_loop()
+
+    # 解决 asynk task 没有被 cancel 的问题
+    async def shutdown(sig: signal.Signals):
+        logger.info(f"收到 {sig.name} 指令，正在关闭所有的东西")
+        for task in ASYNK_TASKS:
+            task.cancel()
+        await asynkio.gather(*ASYNK_TASKS, return_exceptions=True)
+        logger.info("所有的代办提醒 Task 都已经退出了")
+
+    for sig in (signal.SIGINT, signal.SIGTERM):
+        loop.add_signal_handler(sig, functools.partial(
+            asynkio.create_task, shutdown(sig)
+        ))
+
+    await asynkio.gather(*ASYNK_TASKS)