modify: amqtt 打包支持
This commit is contained in:
+4
-1
@@ -87,9 +87,12 @@ base = [
|
||||
"modelvault",
|
||||
"volcenginesdkarkruntime",
|
||||
"chongming_jwt",
|
||||
"chongming-database",
|
||||
"chongming_database",
|
||||
"chongming_config",
|
||||
"chongming_password",
|
||||
"aiosqlite",
|
||||
"sqlmodel",
|
||||
"amqtt"
|
||||
]
|
||||
linux = [
|
||||
"gunicorn.glogging",
|
||||
|
||||
@@ -5,6 +5,8 @@ import multiprocessing
|
||||
import asyncio
|
||||
import socket
|
||||
import sys
|
||||
import fcntl
|
||||
import tempfile
|
||||
|
||||
import boto3
|
||||
from botocore.exceptions import ClientError
|
||||
@@ -32,6 +34,46 @@ def is_port_in_use(port: int, host: str = "127.0.0.1") -> bool:
|
||||
s.settimeout(0.5)
|
||||
return s.connect_ex((host, port)) == 0
|
||||
|
||||
|
||||
def _run_mqtt_broker(host: str, port: int):
|
||||
"""子进程入口：启动 amqtt Broker"""
|
||||
import asyncio as _asyncio
|
||||
from amqtt.broker import Broker
|
||||
|
||||
async def _serve():
|
||||
cfg = {
|
||||
"listeners": {
|
||||
"default": {
|
||||
"type": "tcp", "bind": f"{host}:{port}", "max_connections": 500,
|
||||
},
|
||||
},
|
||||
"timeout_disconnect_delay": 2,
|
||||
"plugins": {
|
||||
"amqtt.plugins.authentication.AnonymousAuthPlugin": {"allow_anonymous": True},
|
||||
"amqtt.plugins.sys.broker.BrokerSysPlugin": {"sys_interval": 0},
|
||||
},
|
||||
}
|
||||
broker = Broker(cfg)
|
||||
try:
|
||||
await broker.start()
|
||||
except Exception:
|
||||
# 端口已被其他 worker 的 Broker 占用，静默退出
|
||||
return
|
||||
print(f"[MQTT Broker] 已启动 @ {host}:{port}")
|
||||
try:
|
||||
while True:
|
||||
await _asyncio.sleep(1)
|
||||
except (KeyboardInterrupt, _asyncio.CancelledError):
|
||||
pass
|
||||
finally:
|
||||
await broker.shutdown()
|
||||
|
||||
try:
|
||||
_asyncio.run(_serve())
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
|
||||
def _run_nanio():
|
||||
"""子进程入口：直接调用 nanio.cli.main()"""
|
||||
# 设置环境变量（从主进程继承，但重新设置确保有效）
|
||||
@@ -183,44 +225,100 @@ async def register_routes(app: FastAPI, manager: WorkerManager):
|
||||
register_dynamic_routes(app, manager, prefix=prefix)
|
||||
# 可选：等待路由注册生效（已在 register_dynamic_routes 内部完成日志）
|
||||
|
||||
def setup_mqtt_broker(app: FastAPI, host: str, port: int):
|
||||
"""启动 MQTT Broker 子进程（文件锁确保多 worker 仅启动一个）"""
|
||||
# 使用文件锁确保多个 gunicorn worker 中只有一个启动 Broker
|
||||
lock_path = os.path.join(tempfile.gettempdir(), "chongming_mqtt_broker.lock")
|
||||
lock_fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o644)
|
||||
try:
|
||||
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except BlockingIOError:
|
||||
# 其他 worker 已持有锁，说明 Broker 正在/已经启动
|
||||
os.close(lock_fd)
|
||||
# 等待端口就绪
|
||||
for _ in range(50):
|
||||
if is_port_in_use(port, host):
|
||||
logger.info(f"MQTT Broker 已在 {host}:{port} 运行，跳过")
|
||||
return
|
||||
time.sleep(0.1)
|
||||
raise RuntimeError("MQTT Broker 端口未就绪（已被其他 worker 锁定）")
|
||||
|
||||
# 持有锁，负责启动 Broker
|
||||
if is_port_in_use(port, host):
|
||||
logger.info(f"MQTT Broker 已在 {host}:{port} 运行，跳过")
|
||||
return
|
||||
proc = multiprocessing.Process(
|
||||
target=_run_mqtt_broker, args=(host, port), daemon=False
|
||||
)
|
||||
proc.start()
|
||||
app.state.mqtt_broker_process = proc
|
||||
app.state._mqtt_lock_fd = lock_fd # 保持锁直到进程退出
|
||||
logger.info(f"MQTT Broker 子进程已启动 (PID: {proc.pid})")
|
||||
for _ in range(50):
|
||||
if is_port_in_use(port, host):
|
||||
logger.info(f"MQTT Broker 已监听 {port} 端口")
|
||||
return
|
||||
if not proc.is_alive():
|
||||
raise RuntimeError("MQTT Broker 进程异常退出")
|
||||
time.sleep(0.1)
|
||||
raise RuntimeError("MQTT Broker 端口未就绪")
|
||||
|
||||
|
||||
async def shutdown_mqtt_broker(app: FastAPI):
|
||||
"""终止 MQTT Broker 子进程"""
|
||||
proc = getattr(app.state, "mqtt_broker_process", None)
|
||||
if proc and proc.is_alive():
|
||||
proc.terminate()
|
||||
proc.join(timeout=5)
|
||||
if proc.is_alive():
|
||||
proc.kill()
|
||||
logger.info("MQTT Broker 子进程已终止")
|
||||
|
||||
|
||||
async def setup_mqtt(app: FastAPI, worker_manager: WorkerManager) -> MQTTBridge | None:
|
||||
"""
|
||||
启动 MQTT Bridge（内嵌 Broker + 客户端桥接）。
|
||||
启动 MQTT Bridge（客户端模式，Broker 由子进程管理）。
|
||||
|
||||
从 config.toml 的环境段（如 [development.mqtt]）读取配置，
|
||||
默认监听 localhost:1883。若 enabled 为 false，则跳过启动。
|
||||
从 config.toml 的环境段读取配置，若 enabled 为 false 则跳过。
|
||||
使用文件锁确保多个 gunicorn worker 中只有一个启动 Bridge。
|
||||
"""
|
||||
env_config = config.get(env, {})
|
||||
mqtt_config = env_config.get("mqtt", {})
|
||||
if not mqtt_config.get("enabled", True):
|
||||
logger.info("MQTT Bridge 已禁用（mqtt.enabled = false）")
|
||||
logger.info("MQTT 已禁用")
|
||||
return None
|
||||
|
||||
host = mqtt_config.get("host", "localhost")
|
||||
port = mqtt_config.get("port", 1883)
|
||||
broker_config = mqtt_config.get("broker", {})
|
||||
|
||||
bridge = MQTTBridge(
|
||||
worker_manager=worker_manager,
|
||||
host=host,
|
||||
port=port,
|
||||
broker_config=broker_config,
|
||||
)
|
||||
# 1. 启动 Broker 子进程（文件锁已确保只有一个）
|
||||
setup_mqtt_broker(app, host, port)
|
||||
|
||||
# 2. 文件锁确保只有一个 worker 启动 MQTT Bridge
|
||||
lock_path = os.path.join(tempfile.gettempdir(), "chongming_mqtt_bridge.lock")
|
||||
lock_fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o644)
|
||||
try:
|
||||
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except BlockingIOError:
|
||||
os.close(lock_fd)
|
||||
logger.info("MQTT Bridge 已由其他 worker 启动，跳过")
|
||||
return None
|
||||
|
||||
# 持有锁，启动 Bridge
|
||||
bridge = MQTTBridge(worker_manager=worker_manager, host=host, port=port)
|
||||
app.state.mqtt_bridge = bridge
|
||||
app.state._mqtt_bridge_lock_fd = lock_fd
|
||||
await bridge.start()
|
||||
logger.info(f"MQTT Bridge 已启动 (host={host}, port={port})")
|
||||
logger.info(f"MQTT Bridge 已启动 ({host}:{port})")
|
||||
return bridge
|
||||
|
||||
|
||||
async def shutdown_mqtt(app: FastAPI):
|
||||
"""停止 MQTT Bridge。"""
|
||||
"""停止 MQTT Bridge 和 Broker 子进程"""
|
||||
bridge: MQTTBridge | None = getattr(app.state, "mqtt_bridge", None)
|
||||
if bridge is None:
|
||||
return
|
||||
|
||||
if bridge.is_running:
|
||||
if bridge and bridge.is_running:
|
||||
await bridge.stop()
|
||||
await shutdown_mqtt_broker(app)
|
||||
|
||||
|
||||
async def shutdown_nanio(app: FastAPI):
|
||||
|
||||
@@ -1,292 +1,88 @@
|
||||
"""
|
||||
MQTT Bridge — 将 MQTT 消息桥接到 Worker 插件
|
||||
|
||||
架构:
|
||||
MQTT 客户端订阅 chongming/+/+ 主题
|
||||
→ 解析 topic: chongming/<worker_type>/<operation>
|
||||
→ 调用 WorkerManager.call(worker_type, subject, payload)
|
||||
→ 将响应发布到 chongming/<worker_type>/<operation>/response
|
||||
|
||||
同时内嵌一个 amqtt Broker，无需外部 MQTT 服务器。
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from amqtt.broker import Broker
|
||||
"""MQTT Bridge — 纯客户端模式，Broker 由子进程管理"""
|
||||
import asyncio, json
|
||||
from typing import Any, Optional
|
||||
from amqtt.client import MQTTClient
|
||||
from amqtt.mqtt.constants import QOS_1
|
||||
|
||||
from .logger import get_logger
|
||||
|
||||
logger = get_logger("mqtt")
|
||||
|
||||
# MQTT 主题前缀
|
||||
MQTT_TOPIC_PREFIX = "chongming"
|
||||
# 订阅通配符：匹配 chongming/<worker_type>/<operation>
|
||||
MQTT_SUBSCRIBE_TOPIC = f"{MQTT_TOPIC_PREFIX}/+/+"
|
||||
|
||||
|
||||
class MQTTBridge:
|
||||
"""
|
||||
MQTT 桥接器
|
||||
|
||||
内嵌 amqtt Broker + 客户端，将 MQTT 消息转发到 WorkerManager。
|
||||
|
||||
用法:
|
||||
bridge = MQTTBridge(worker_manager, host="localhost", port=1883)
|
||||
await bridge.start()
|
||||
# ... 应用运行 ...
|
||||
await bridge.stop()
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
worker_manager,
|
||||
host: str = "localhost",
|
||||
port: int = 1883,
|
||||
broker_config: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self._worker_manager = worker_manager
|
||||
self._host = host
|
||||
self._port = port
|
||||
self._broker_config = broker_config or {}
|
||||
|
||||
self._broker: Optional[Broker] = None
|
||||
def __init__(self, worker_manager, host="localhost", port=1883):
|
||||
self._wm = worker_manager; self._host = host; self._port = port
|
||||
self._client: Optional[MQTTClient] = None
|
||||
self._broker_task: Optional[asyncio.Task] = None
|
||||
self._listen_task: Optional[asyncio.Task] = None
|
||||
self._lt: Optional[asyncio.Task] = None
|
||||
self._running = False
|
||||
|
||||
# ============================================================
|
||||
# 生命周期
|
||||
# ============================================================
|
||||
|
||||
async def start(self) -> None:
|
||||
"""启动 MQTT Broker 和客户端桥接"""
|
||||
if self._running:
|
||||
logger.warning("MQTT Bridge 已在运行中")
|
||||
return
|
||||
|
||||
# 1. 启动内嵌 Broker
|
||||
await self._start_broker()
|
||||
|
||||
# 2. 启动客户端桥接
|
||||
await self._start_client()
|
||||
|
||||
async def start(self):
|
||||
if self._running: return
|
||||
try: await self._sc()
|
||||
except Exception as e:
|
||||
logger.warning(f"MQTT 不可用: {e}"); return
|
||||
self._running = True
|
||||
logger.info(f"MQTT Bridge 已启动 (broker: {self._host}:{self._port})")
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""停止客户端和 Broker"""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
logger.info("正在关闭 MQTT Bridge...")
|
||||
|
||||
# 1. 停止客户端
|
||||
await self._stop_client()
|
||||
|
||||
# 2. 停止 Broker
|
||||
await self._stop_broker()
|
||||
logger.info(f"MQTT Bridge 已连接 {self._host}:{self._port}")
|
||||
|
||||
async def stop(self):
|
||||
if not self._running: return
|
||||
self._running = False
|
||||
logger.info("MQTT Bridge 已关闭")
|
||||
await self._sc_stop()
|
||||
|
||||
# ============================================================
|
||||
# Broker（内嵌 MQTT 服务器）
|
||||
# ============================================================
|
||||
|
||||
async def _start_broker(self) -> None:
|
||||
"""启动内嵌 amqtt Broker"""
|
||||
config = {
|
||||
"listeners": {
|
||||
"default": {
|
||||
"type": "tcp",
|
||||
"bind": f"{self._host}:{self._port}",
|
||||
"max_connections": 500,
|
||||
},
|
||||
},
|
||||
"timeout_disconnect_delay": 2,
|
||||
"sys_interval": 0, # 禁用 sys 主题周期性发布，避免 None > 0 异常
|
||||
"auth": {
|
||||
"allow-anonymous": True,
|
||||
},
|
||||
}
|
||||
# 合并用户自定义配置
|
||||
config.update(self._broker_config)
|
||||
|
||||
self._broker = Broker(config)
|
||||
|
||||
async def _run_broker():
|
||||
try:
|
||||
await self._broker.start() # type: ignore
|
||||
logger.info(f"amqtt Broker 已启动 @ {self._host}:{self._port}")
|
||||
# 保持 broker 运行直到被取消
|
||||
while True:
|
||||
await asyncio.sleep(1)
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(f"Broker 运行异常: {e}")
|
||||
|
||||
self._broker_task = asyncio.create_task(_run_broker())
|
||||
# 等待 broker 就绪
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
async def _stop_broker(self) -> None:
|
||||
"""停止 Broker"""
|
||||
if self._broker_task and not self._broker_task.done():
|
||||
self._broker_task.cancel()
|
||||
try:
|
||||
await self._broker_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
if self._broker:
|
||||
try:
|
||||
await self._broker.shutdown()
|
||||
except Exception as e:
|
||||
logger.error(f"Broker 关闭异常: {e}")
|
||||
self._broker = None
|
||||
|
||||
logger.info("amqtt Broker 已停止")
|
||||
|
||||
# ============================================================
|
||||
# 客户端桥接
|
||||
# ============================================================
|
||||
|
||||
async def _start_client(self) -> None:
|
||||
"""启动 MQTT 客户端并订阅 worker 主题"""
|
||||
self._client = MQTTClient()
|
||||
async def _sc(self):
|
||||
self._client = MQTTClient(config={"auto_reconnect": False})
|
||||
await self._client.connect(f"mqtt://{self._host}:{self._port}/")
|
||||
await self._client.subscribe([(MQTT_SUBSCRIBE_TOPIC, QOS_1)])
|
||||
logger.info(f"MQTT 客户端已订阅: {MQTT_SUBSCRIBE_TOPIC}")
|
||||
|
||||
self._listen_task = asyncio.create_task(self._message_loop())
|
||||
|
||||
async def _stop_client(self) -> None:
|
||||
"""断开 MQTT 客户端"""
|
||||
if self._listen_task and not self._listen_task.done():
|
||||
self._listen_task.cancel()
|
||||
try:
|
||||
await self._listen_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info(f"已订阅 {MQTT_SUBSCRIBE_TOPIC}")
|
||||
self._lt = asyncio.create_task(self._loop())
|
||||
|
||||
async def _sc_stop(self):
|
||||
# 先断开连接，让 deliver_message() 解锁
|
||||
if self._client:
|
||||
try:
|
||||
await self._client.unsubscribe([MQTT_SUBSCRIBE_TOPIC])
|
||||
await self._client.disconnect()
|
||||
except Exception as e:
|
||||
logger.error(f"MQTT 客户端断开异常: {e}")
|
||||
except Exception: pass
|
||||
self._client = None
|
||||
# 再取消消息循环
|
||||
if self._lt and not self._lt.done():
|
||||
self._lt.cancel()
|
||||
try: await self._lt
|
||||
except asyncio.CancelledError: pass
|
||||
|
||||
logger.info("MQTT 客户端已断开")
|
||||
|
||||
# ============================================================
|
||||
# 消息处理
|
||||
# ============================================================
|
||||
|
||||
async def _message_loop(self) -> None:
|
||||
"""消息循环：接收 MQTT 消息并转发到 WorkerManager"""
|
||||
while True:
|
||||
async def _loop(self):
|
||||
while self._running:
|
||||
try:
|
||||
message = await self._client.deliver_message() # type: ignore
|
||||
await self._handle_message(message)
|
||||
except asyncio.CancelledError:
|
||||
m = await self._client.deliver_message() # type: ignore
|
||||
if m is None: break
|
||||
await self._handle(m)
|
||||
except asyncio.CancelledError: break
|
||||
except Exception as e:
|
||||
if self._running:
|
||||
logger.error(f"消息异常: {e}")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"MQTT 消息处理异常: {e}")
|
||||
|
||||
async def _handle_message(self, message) -> None:
|
||||
"""
|
||||
处理收到的 MQTT 消息
|
||||
|
||||
解析 topic: chongming/<worker_type>/<operation>
|
||||
转发到 WorkerManager.call()
|
||||
发布响应到 chongming/<worker_type>/<operation>/response
|
||||
"""
|
||||
topic: str = message.topic # type: ignore
|
||||
payload_bytes: bytes = message.data # type: ignore
|
||||
|
||||
# 解析主题
|
||||
parts = topic.split("/")
|
||||
if len(parts) < 3:
|
||||
logger.warning(f"无效的 MQTT 主题: {topic}")
|
||||
return
|
||||
|
||||
worker_type = parts[1]
|
||||
operation = parts[2]
|
||||
subject = f"{worker_type}.{operation}"
|
||||
response_topic = f"{MQTT_TOPIC_PREFIX}/{worker_type}/{operation}/response"
|
||||
|
||||
# 解析载荷
|
||||
async def _handle(self, m):
|
||||
t: str = m.topic; p = m.data # type: ignore
|
||||
parts = t.split("/")
|
||||
if len(parts) < 3: return
|
||||
wt, op = parts[1], parts[2]
|
||||
subj = f"{wt}.{op}"
|
||||
rt = f"{MQTT_TOPIC_PREFIX}/{wt}/{op}/response"
|
||||
try: pl = json.loads(p.decode("utf-8"))
|
||||
except json.JSONDecodeError: return
|
||||
pl.setdefault("_method", "GET")
|
||||
logger.info(f"MQTT -> {subj}")
|
||||
try:
|
||||
payload_str = payload_bytes.decode("utf-8") if payload_bytes else "{}"
|
||||
payload = json.loads(payload_str)
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"无效的 JSON 载荷: {payload_bytes}")
|
||||
await self._publish_response(
|
||||
response_topic,
|
||||
{"error": "Invalid JSON payload"},
|
||||
)
|
||||
return
|
||||
r = await self._wm.call(worker_type=wt, subject=subj, payload=pl,
|
||||
timeout=float(pl.pop("_timeout",30)), cacheable=bool(pl.pop("_cacheable",False)))
|
||||
if r is not None: await self._pub(rt, r)
|
||||
except RuntimeError as e: await self._pub(rt, {"error":str(e)})
|
||||
except Exception as e: await self._pub(rt, {"error":str(e)})
|
||||
|
||||
# 确保 _method 字段存在
|
||||
if "_method" not in payload:
|
||||
payload["_method"] = "GET"
|
||||
|
||||
logger.info(
|
||||
f"MQTT → Worker: {topic} [{payload.get('_method', 'GET')}] "
|
||||
f"-> {worker_type}/{subject}"
|
||||
)
|
||||
|
||||
# 调用 WorkerManager
|
||||
async def _pub(self, topic, data):
|
||||
try:
|
||||
result = await self._worker_manager.call(
|
||||
worker_type=worker_type,
|
||||
subject=subject,
|
||||
payload=payload,
|
||||
timeout=float(payload.pop("_timeout", 30.0)),
|
||||
cacheable=bool(payload.pop("_cacheable", False)),
|
||||
)
|
||||
# result 为 None 时不发布响应（如心跳上报等无需返回的场景）
|
||||
if result is not None:
|
||||
await self._publish_response(response_topic, result)
|
||||
logger.debug(f"Worker → MQTT: {response_topic} success")
|
||||
else:
|
||||
logger.debug(f"Worker → MQTT: {topic} (无响应，已跳过)")
|
||||
except RuntimeError as e:
|
||||
error_resp = {"error": str(e), "worker_type": worker_type, "subject": subject}
|
||||
await self._publish_response(response_topic, error_resp)
|
||||
logger.warning(f"Worker 调用失败: {e}")
|
||||
except Exception as e:
|
||||
error_resp = {"error": str(e), "worker_type": worker_type, "subject": subject}
|
||||
await self._publish_response(response_topic, error_resp)
|
||||
logger.error(f"MQTT 桥接异常: {e}")
|
||||
|
||||
async def _publish_response(self, topic: str, data: Any) -> None:
|
||||
"""将响应发布到 MQTT 主题"""
|
||||
try:
|
||||
payload = json.dumps(data, ensure_ascii=False, default=str)
|
||||
await self._client.publish(topic, payload.encode("utf-8"), QOS_1) # type: ignore
|
||||
except Exception as e:
|
||||
logger.error(f"MQTT 发布失败 [{topic}]: {e}")
|
||||
|
||||
# ============================================================
|
||||
# 属性
|
||||
# ============================================================
|
||||
pd = json.dumps(data, ensure_ascii=False, default=str)
|
||||
await self._client.publish(topic, pd.encode("utf-8"), QOS_1) # type: ignore
|
||||
except Exception as e: logger.error(f"发布失败 [{topic}]: {e}")
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self._running
|
||||
|
||||
@property
|
||||
def broker_host(self) -> str:
|
||||
return self._host
|
||||
|
||||
@property
|
||||
def broker_port(self) -> int:
|
||||
return self._port
|
||||
def is_running(self): return self._running
|
||||
|
||||
@@ -265,6 +265,8 @@ def _stage_pyinstaller(config: Config) -> tuple[str, str]:
|
||||
|
||||
cmd.append("--collect-all")
|
||||
cmd.append("nanio")
|
||||
cmd.append("--collect-all")
|
||||
cmd.append("amqtt")
|
||||
|
||||
cmd.append("main.py")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user