From 3ee018e3ddbf2dba58ee3ae0bdb3f7323ac914e5 Mon Sep 17 00:00:00 2001 From: chakcy <947105045@qq.com> Date: Fri, 24 Jul 2026 18:53:36 +0800 Subject: [PATCH] =?UTF-8?q?modify:=20amqtt=20=E6=89=93=E5=8C=85=E6=94=AF?= =?UTF-8?q?=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config.toml | 5 +- src/chongming/app/core/lifespan_setup.py | 132 ++++++++-- src/chongming/app/core/mqtt_bridge.py | 318 ++++------------------- src/chongming/scripts/build.py | 2 + 4 files changed, 178 insertions(+), 279 deletions(-) diff --git a/config.toml b/config.toml index 8e602df..d1c965c 100644 --- a/config.toml +++ b/config.toml @@ -87,9 +87,12 @@ base = [ "modelvault", "volcenginesdkarkruntime", "chongming_jwt", - "chongming-database", + "chongming_database", + "chongming_config", + "chongming_password", "aiosqlite", "sqlmodel", + "amqtt" ] linux = [ "gunicorn.glogging", diff --git a/src/chongming/app/core/lifespan_setup.py b/src/chongming/app/core/lifespan_setup.py index 1942b43..2cab302 100644 --- a/src/chongming/app/core/lifespan_setup.py +++ b/src/chongming/app/core/lifespan_setup.py @@ -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): diff --git a/src/chongming/app/core/mqtt_bridge.py b/src/chongming/app/core/mqtt_bridge.py index 35113ca..12b78fd 100644 --- a/src/chongming/app/core/mqtt_bridge.py +++ b/src/chongming/app/core/mqtt_bridge.py @@ -1,292 +1,88 @@ -""" -MQTT Bridge — 将 MQTT 消息桥接到 Worker 插件 - -架构: - MQTT 客户端订阅 chongming/+/+ 主题 - → 解析 topic: chongming// - → 调用 WorkerManager.call(worker_type, subject, payload) - → 将响应发布到 chongming///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// 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: - break + 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: - logger.error(f"MQTT 消息处理异常: {e}") + if self._running: + logger.error(f"消息异常: {e}") + break - async def _handle_message(self, message) -> None: - """ - 处理收到的 MQTT 消息 - - 解析 topic: chongming// - 转发到 WorkerManager.call() - 发布响应到 chongming///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 diff --git a/src/chongming/scripts/build.py b/src/chongming/scripts/build.py index 17e7fd3..7a20071 100644 --- a/src/chongming/scripts/build.py +++ b/src/chongming/scripts/build.py @@ -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")