tnot/konabot
1
0
Fork 0
forked from mttu-developers/konabot
Code Pull Requests Activity

调整 notify 的强制退出

Browse Source
This commit is contained in:
钟晓帕
2025-10-13 22:16:50 +08:00
parent 0ec66988fa
commit fc5b11c5e8
1 changed files with 32 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
konabot/plugins/simple_notify/__init__.py
View File

@ -1,8 +1,10 @@
import asyncio import asyncio as asynkio
import datetime import datetime
import functools
from pathlib import Path from pathlib import Path
from typing import Any, Literal, cast from typing import Any, Literal, cast
import signal
import nonebot import nonebot
import ptimeparse import ptimeparse
from loguru import logger from loguru import logger
@ -24,7 +26,9 @@ evt = on_message()
(Path(__file__).parent.parent.parent.parent / "data").mkdir(exist_ok=True) (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_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): class Notify(BaseModel):
@ -111,7 +115,11 @@ def create_notify_task(notify: Notify, fail2remove: bool = True):
async def mission(): async def mission():
begin_time = datetime.datetime.now() begin_time = datetime.datetime.now()
if begin_time < notify.notify_time: 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: else:
logger.warning( logger.warning(
f"期望在 {notify.notify_time} 在平台 {notify.platform} {notify.target_env}" 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() DATA_FILE_LOCK.release()
else: else:
pass pass
return asyncio.create_task(mission()) return asynkio.create_task(mission())
@evt.handle() @evt.handle()
@ -214,11 +222,11 @@ async def _():
DELTA = 2 DELTA = 2
logger.info(f"第一次探测到 Bot 连接，等待 {DELTA} 秒后开始通知") logger.info(f"第一次探测到 Bot 连接，等待 {DELTA} 秒后开始通知")
await asyncio.sleep(DELTA) await asynkio.sleep(DELTA)
await DATA_FILE_LOCK.acquire() await DATA_FILE_LOCK.acquire()
tasks: set[asyncio.Task[Any]] = set() # tasks: set[asynkio.Task[Any]] = set()
cfg = load_notify_config() cfg = load_notify_config()
if cfg.version == 1: if cfg.version == 1:
logger.info("将配置文件的版本升级为 2") logger.info("将配置文件的版本升级为 2")
@ -227,11 +235,26 @@ async def _():
counter = 0 counter = 0
for notify in [*cfg.notifies]: for notify in [*cfg.notifies]:
task = create_notify_task(notify, fail2remove=False) task = create_notify_task(notify, fail2remove=False)
tasks.add(task) ASYNK_TASKS.add(task)
task.add_done_callback(lambda self: tasks.remove(self)) task.add_done_callback(lambda self: ASYNK_TASKS.remove(self))
counter += 1 counter += 1
logger.info(f"成功创建了 {counter} 条代办事项") logger.info(f"成功创建了 {counter} 条代办事项")
save_notify_config(cfg) save_notify_config(cfg)
DATA_FILE_LOCK.release() 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)