ElapsingDreams/LiteyukiBot
1
0
Fork 0
forked from bot/app
Code Pull Requests Activity

🐛 Ctrl+C 无法终止Channel接收线程的问题

Browse Source
This commit is contained in:
远野千束 2024-08-16 23:43:43 +08:00
parent 0417805e46
commit 85a13251a5
6 changed files with 123 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
liteyuki/bot/__init__.py
View File

@ -1,6 +1,8 @@
import asyncio
import atexit
import os
import platform
import signal
import sys
import threading
import time
@ -38,11 +40,16 @@ class LiteyukiBot:
self.loop = asyncio.new_event_loop()
asyncio.set_event_loop(self.loop)
self.stop_event = threading.Event()
self.call_restart_count = 0
load_plugins("liteyuki/plugins") # 加载轻雪插件
signal.signal(signal.SIGINT, self._handle_exit)
signal.signal(signal.SIGTERM, self._handle_exit)
atexit.register(self.process_manager.terminate_all) # 注册退出时的函数
def run(self):
"""
启动逻辑
@ -60,12 +67,24 @@ class LiteyukiBot:
"""
try:
while not self.stop_event.is_set():
time.sleep(1)
time.sleep(0.5)
except KeyboardInterrupt:
logger.info("Liteyuki is stopping...")
self.stop()
finally:
self.lifespan.after_shutdown()
def _handle_exit(self, signum, frame):
"""
信号处理
Args:
signum:
frame:
Returns:
"""
logger.info("Received signal, stopping all processes.")
self.stop()
sys.exit(0)
def restart(self, delay: int = 0):
"""
@ -116,16 +135,12 @@ class LiteyukiBot:
Returns:
"""
self.init_config()
self.init_logger()
def init_logger(self):
# 修改nonebot的日志配置
init_log(config=self.config)
def init_config(self):
pass
def stop(self):
"""
停止轻雪

4
liteyuki/comm/channel.py
View File

@ -93,10 +93,10 @@ class Channel:
装饰器，装饰一个函数在接收到数据后执行
"""
if (not self.is_sub_receive_loop_running) and not IS_MAIN_PROCESS:
threading.Thread(target=self._start_sub_receive_loop).start()
threading.Thread(target=self._start_sub_receive_loop, daemon=True).start()
if (not self.is_main_receive_loop_running) and IS_MAIN_PROCESS:
threading.Thread(target=self._start_main_receive_loop).start()
threading.Thread(target=self._start_main_receive_loop, daemon=True).start()
def decorator(func: ON_RECEIVE_FUNC) -> ON_RECEIVE_FUNC:
async def wrapper(data: Any) -> Any:

74
liteyuki/comm/storage.py
View File

@ -102,6 +102,77 @@ class KeyValueStore:
return self.active_chan.receive()
class KeyValueStoreNoLock:
def __init__(self):
self._store = {}
self.active_chan = Channel(_id="shared_memory-active")
self.passive_chan = Channel(_id="shared_memory-passive")
def set(self, key: str, value: any) -> None:
"""
设置键值对
Args:
key: 键
value: 值
"""
if IS_MAIN_PROCESS:
self._store[key] = value
else:
# 向主进程发送请求拿取
self.passive_chan.send(("set", key, value))
def get(self, key: str, default: Optional[any] = None) -> any:
"""
获取键值对
Args:
key: 键
default: 默认值
Returns:
any: 值
"""
if IS_MAIN_PROCESS:
return self._store.get(key, default)
else:
self.passive_chan.send(("get", key, default))
return self.active_chan.receive()
def delete(self, key: str, ignore_key_error: bool = True) -> None:
"""
删除键值对
Args:
key: 键
ignore_key_error: 是否忽略键不存在的错误
Returns:
"""
if IS_MAIN_PROCESS:
if key in self._store:
try:
del self._store[key]
del _locks[key]
except KeyError as e:
if not ignore_key_error:
raise e
else:
# 向主进程发送请求删除
self.passive_chan.send(("delete", key))
def get_all(self) -> dict[str, any]:
"""
获取所有键值对
Returns:
dict[str, any]: 键值对
"""
if IS_MAIN_PROCESS:
return self._store
else:
self.passive_chan.send(("get_all",))
return self.active_chan.receive()
class GlobalKeyValueStore:
_instance = None
_lock = threading.Lock()
@ -142,7 +213,8 @@ if IS_MAIN_PROCESS:
@shared_memory.passive_chan.on_receive(lambda d: d[0] == "get_all")
def on_get_all(d):
shared_memory.active_chan.send(shared_memory.get_all())
if d[0] == "get_all":
shared_memory.active_chan.send(shared_memory.get_all())
else:
shared_memory = None

56
liteyuki/core/manager.py
View File

@ -11,6 +11,7 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
import atexit
import signal
import sys
import threading
from multiprocessing import Process
from typing import Any, Callable, TYPE_CHECKING, TypeAlias
@ -63,10 +64,6 @@ class ProcessManager:
self.targets: dict[str, tuple[callable, tuple, dict]] = {}
self.processes: dict[str, Process] = {}
atexit.register(self.terminate_all)
signal.signal(signal.SIGINT, self._handle_exit)
signal.signal(signal.SIGTERM, self._handle_exit)
def start(self, name: str):
"""
开启后自动监控进程，并添加到进程字典中
@ -82,34 +79,31 @@ class ProcessManager:
def _start_process():
process = Process(target=self.targets[name][0], args=self.targets[name][1],
kwargs=self.targets[name][2])
kwargs=self.targets[name][2], daemon=True)
self.processes[name] = process
process.start()
# 启动进程并监听信号
_start_process()
def _start_monitor():
while True:
data = chan_active.receive()
if data == 0:
# 停止
logger.info(f"Stopping process {name}")
self.bot.lifespan.before_process_shutdown()
self.terminate(name)
break
elif data == 1:
# 重启
logger.info(f"Restarting process {name}")
self.bot.lifespan.before_process_shutdown()
self.bot.lifespan.before_process_restart()
self.terminate(name)
_start_process()
continue
else:
logger.warning("Unknown data received, ignored.")
_start_monitor()
while True:
data = chan_active.receive()
if data == 0:
# 停止
logger.info(f"Stopping process {name}")
self.bot.lifespan.before_process_shutdown()
self.terminate(name)
break
elif data == 1:
# 重启
logger.info(f"Restarting process {name}")
self.bot.lifespan.before_process_shutdown()
self.bot.lifespan.before_process_restart()
self.terminate(name)
_start_process()
continue
else:
logger.warning("Unknown data received, ignored.")
def start_all(self):
"""
@ -118,11 +112,6 @@ class ProcessManager:
for name in self.targets:
threading.Thread(target=self.start, args=(name,), daemon=True).start()
def _handle_exit(self, signum, frame):
logger.info("Received signal, stopping all processes.")
self.terminate_all()
exit(0)
def add_target(self, name: str, target: TARGET_FUNC, args: tuple = (), kwargs=None):
"""
添加进程
@ -159,8 +148,9 @@ class ProcessManager:
Returns:
"""
if name not in self.targets:
raise logger.warning(f"Process {name} not found.")
if name not in self.processes:
logger.warning(f"Process {name} not found.")
return
process = self.processes[name]
process.terminate()
process.join(TIMEOUT)

1
src/liteyuki_plugins/nonebot_launcher/__init__.py
View File

@ -10,7 +10,6 @@ Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
"""
import nonebot
from liteyuki.utils import IS_MAIN_PROCESS
from liteyuki.plugin import PluginMetadata
from .nb_utils import adapter_manager, driver_manager

10
src/utils/base/__init__.py
View File

@ -1,7 +1,7 @@
import threading
from nonebot import logger
from liteyuki.comm.channel import get_channel
from liteyuki.comm.channel import active_channel
def reload(delay: float = 0.0, receiver: str = "nonebot"):
@ -14,13 +14,9 @@ def reload(delay: float = 0.0, receiver: str = "nonebot"):
Returns:
"""
chan = get_channel(receiver + "-active")
if chan is None:
logger.error(f"Channel {receiver}-active not found, cannot reload.")
return
if delay > 0:
threading.Timer(delay, chan.send, args=(1,)).start()
threading.Timer(delay, active_channel.send, args=(1,)).start()
return
else:
chan.send(1)
active_channel.send(1)