mirror of
https://github.com/LiteyukiStudio/LiteyukiBot.git
synced 2024-11-22 21:27:43 +08:00
273 lines
7.7 KiB
Python
273 lines
7.7 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Copyright (C) 2020-2024 LiteyukiStudio. All Rights Reserved
|
||
|
||
@Time : 2024/7/27 上午11:12
|
||
@Author : snowykami
|
||
@Email : snowykami@outlook.com
|
||
@File : manager.py
|
||
@Software: PyCharm
|
||
"""
|
||
import asyncio
|
||
import multiprocessing
|
||
import threading
|
||
from multiprocessing import Process
|
||
from typing import Any, Callable, TYPE_CHECKING, TypeAlias
|
||
|
||
from croterline.context import Context
|
||
from croterline.process import SubProcess, ProcessFuncType
|
||
|
||
from liteyuki.log import logger
|
||
from liteyuki.utils import IS_MAIN_PROCESS
|
||
|
||
if TYPE_CHECKING:
|
||
from liteyuki.bot.lifespan import Lifespan
|
||
from liteyuki.comm.storage import KeyValueStore
|
||
|
||
from liteyuki.comm import Channel
|
||
|
||
if IS_MAIN_PROCESS:
|
||
from liteyuki.comm.channel import get_channel, publish_channel, get_channels
|
||
from liteyuki.comm.storage import shared_memory
|
||
from liteyuki.comm.channel import (
|
||
channel_deliver_active_channel,
|
||
channel_deliver_passive_channel,
|
||
)
|
||
else:
|
||
from liteyuki.comm import channel
|
||
from liteyuki.comm import storage
|
||
|
||
TARGET_FUNC: TypeAlias = Callable[..., Any]
|
||
TIMEOUT = 10
|
||
|
||
__all__ = ["ProcessManager", "sub_process_manager"]
|
||
multiprocessing.set_start_method("spawn", force=True)
|
||
|
||
|
||
class ChannelDeliver:
|
||
def __init__(
|
||
self,
|
||
active: Channel[Any],
|
||
passive: Channel[Any],
|
||
channel_deliver_active: Channel[Channel[Any]],
|
||
channel_deliver_passive: Channel[tuple[str, dict]],
|
||
publish: Channel[tuple[str, Any]],
|
||
):
|
||
self.active = active
|
||
self.passive = passive
|
||
self.channel_deliver_active = channel_deliver_active
|
||
self.channel_deliver_passive = channel_deliver_passive
|
||
self.publish = publish
|
||
|
||
|
||
# 函数处理一些跨进程通道的
|
||
def _delivery_channel_wrapper(
|
||
func: TARGET_FUNC, cd: ChannelDeliver, sm: "KeyValueStore", *args, **kwargs
|
||
):
|
||
"""
|
||
子进程入口函数
|
||
处理一些操作
|
||
"""
|
||
# 给子进程设置通道
|
||
if IS_MAIN_PROCESS:
|
||
raise RuntimeError("Function should only be called in a sub process.")
|
||
|
||
channel.active_channel = cd.active # 子进程主动通道
|
||
channel.passive_channel = cd.passive # 子进程被动通道
|
||
channel.channel_deliver_active_channel = (
|
||
cd.channel_deliver_active
|
||
) # 子进程通道传递主动通道
|
||
channel.channel_deliver_passive_channel = (
|
||
cd.channel_deliver_passive
|
||
) # 子进程通道传递被动通道
|
||
channel.publish_channel = cd.publish # 子进程发布通道
|
||
|
||
# 给子进程创建共享内存实例
|
||
|
||
storage.shared_memory = sm
|
||
|
||
func(*args, **kwargs)
|
||
|
||
|
||
class ProcessManager:
|
||
"""
|
||
进程管理器
|
||
"""
|
||
|
||
def __init__(self, lifespan: "Lifespan"):
|
||
self.lifespan = lifespan
|
||
self.targets: dict[str, tuple[Callable, tuple, dict]] = {}
|
||
self.processes: dict[str, Process] = {}
|
||
|
||
def _run_process(self, name: str):
|
||
"""
|
||
开启后自动监控进程，并添加到进程字典中，会阻塞，请创建task
|
||
Args:
|
||
name:
|
||
Returns:
|
||
"""
|
||
if name not in self.targets:
|
||
raise KeyError(f"Process {name} not found.")
|
||
|
||
chan_active = get_channel(f"{name}-active")
|
||
|
||
def _start_process():
|
||
process = Process(
|
||
target=self.targets[name][0],
|
||
args=self.targets[name][1],
|
||
kwargs=self.targets[name][2],
|
||
daemon=True,
|
||
)
|
||
self.processes[name] = process
|
||
process.start()
|
||
|
||
# 启动进程并监听信号
|
||
_start_process()
|
||
while True:
|
||
data = chan_active.receive()
|
||
if data == 0:
|
||
# 停止
|
||
logger.info(f"Stopping process {name}")
|
||
self.terminate(name)
|
||
break
|
||
elif data == 1:
|
||
# 重启
|
||
logger.info(f"Restarting process {name}")
|
||
self.terminate(name)
|
||
_start_process()
|
||
continue
|
||
else:
|
||
logger.warning("Unknown data received, ignored.")
|
||
|
||
def start_all(self):
|
||
"""
|
||
对外启动方法，启动所有进程，创建asyncio task
|
||
"""
|
||
# [asyncio.create_task(self._run_process(name)) for name in self.targets]
|
||
|
||
for name in self.targets:
|
||
logger.debug(f"Starting process {name}")
|
||
threading.Thread(
|
||
target=self._run_process, args=(name,), daemon=True
|
||
).start()
|
||
|
||
def add_target(self, name: str, target: TARGET_FUNC, args: tuple = (), kwargs=None):
|
||
"""
|
||
添加进程
|
||
Args:
|
||
name: 进程名，用于获取和唯一标识
|
||
target: 进程函数
|
||
args: 进程函数参数
|
||
kwargs: 进程函数关键字参数，通常会默认传入chan_active和chan_passive
|
||
"""
|
||
if kwargs is None:
|
||
kwargs = {}
|
||
chan_active: Channel = Channel(name=f"{name}-active")
|
||
chan_passive: Channel = Channel(name=f"{name}-passive")
|
||
|
||
channel_deliver = ChannelDeliver(
|
||
active=chan_active,
|
||
passive=chan_passive,
|
||
channel_deliver_active=channel_deliver_active_channel,
|
||
channel_deliver_passive=channel_deliver_passive_channel,
|
||
publish=publish_channel,
|
||
)
|
||
|
||
self.targets[name] = (
|
||
_delivery_channel_wrapper,
|
||
(target, channel_deliver, shared_memory, *args),
|
||
kwargs,
|
||
)
|
||
# 主进程通道
|
||
|
||
def join_all(self):
|
||
for name, process in self.targets:
|
||
process.join()
|
||
|
||
def terminate(self, name: str):
|
||
"""
|
||
终止进程并从进程字典中删除
|
||
Args:
|
||
name:
|
||
|
||
Returns:
|
||
|
||
"""
|
||
if name not in self.processes:
|
||
logger.warning(f"Process {name} not found.")
|
||
return
|
||
process = self.processes[name]
|
||
process.terminate()
|
||
process.join(TIMEOUT)
|
||
if process.is_alive():
|
||
process.kill()
|
||
logger.success(f"Process {name} terminated.")
|
||
|
||
def terminate_all(self):
|
||
for name in self.targets:
|
||
self.terminate(name)
|
||
|
||
def is_process_alive(self, name: str) -> bool:
|
||
"""
|
||
检查进程是否存活
|
||
Args:
|
||
name:
|
||
|
||
Returns:
|
||
|
||
"""
|
||
if name not in self.targets:
|
||
logger.warning(f"Process {name} not found.")
|
||
return self.processes[name].is_alive()
|
||
|
||
|
||
# new version
|
||
|
||
|
||
class _SubProcessManager:
|
||
|
||
def __init__(self):
|
||
self.processes: dict[str, SubProcess] = {}
|
||
|
||
def new_process(
|
||
self, name: str, *args, **kwargs
|
||
) -> Callable[[ProcessFuncType], None]:
|
||
def decorator(func: ProcessFuncType):
|
||
self.processes[name] = SubProcess(name, func, *args, **kwargs)
|
||
|
||
return decorator
|
||
|
||
def add(self, name: str, func: ProcessFuncType, *args, **kwargs):
|
||
"""
|
||
添加子进程
|
||
Args:
|
||
func: 子进程函数
|
||
name: 子进程名称
|
||
args: 子进程函数参数
|
||
kwargs: 子进程函数关键字参数
|
||
Returns:
|
||
"""
|
||
self.processes[name] = SubProcess(name, func, *args, **kwargs)
|
||
|
||
def start(self, name: str):
|
||
"""
|
||
启动指定子进程
|
||
Args:
|
||
name: 子进程名称
|
||
Returns:
|
||
"""
|
||
if name not in self.processes:
|
||
raise KeyError(f"Process {name} not found.")
|
||
self.processes[name].start()
|
||
|
||
def start_all(self):
|
||
"""
|
||
启动所有子进程
|
||
"""
|
||
for name, process in self.processes.items():
|
||
process.start()
|
||
logger.debug(f"Starting process {name}")
|
||
|
||
|
||
sub_process_manager = _SubProcessManager()
|