From 255f7cf28f1f4f96bb02a24501aa48b801d0d31f Mon Sep 17 00:00:00 2001 From: chakcy <947105045@qq.com> Date: Thu, 27 Aug 2026 09:49:31 +0800 Subject: [PATCH] =?UTF-8?q?modify:=20worker=20=E7=AE=A1=E7=90=86=E7=AB=AF?= =?UTF-8?q?=E7=82=B9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 3 +- src/chongming_worker/plugin.py | 25 +- src/chongming_worker/runtime.py | 584 ++++++++++++++++++++++++++++++-- 3 files changed, 585 insertions(+), 27 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index fe12eb6..fdb3b1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = ["hatchling"] [project] name = "chongming-worker" -version = "0.3.0" +version = "0.4.0" description = "chongming worker" readme = "README.md" requires-python = ">=3.12" @@ -12,6 +12,7 @@ dependencies = [ "eclipse-zenoh>=1.9.0", "pydantic>=2.0", "tomli>=2.4.1", + "tomli-w>=1.0.0", "minio>=7.2.0", ] diff --git a/src/chongming_worker/plugin.py b/src/chongming_worker/plugin.py index 7b14dcf..c78d496 100644 --- a/src/chongming_worker/plugin.py +++ b/src/chongming_worker/plugin.py @@ -124,6 +124,27 @@ class Plugin: async def on_cleanup(self) -> None: """cleanup 时调用，清理自定义资源。""" + async def reload(self, config: dict) -> None: + """ + 配置热重载：以新配置重新初始化插件（重建客户端等），供动态改配置后即时生效。 + + 在内部互斥锁下执行：先清理旧资源（on_cleanup），再用新 config 重新 on_init。 + 由 runtime 的 config_set 指令在 index0 副本处调用。 + + Args: + config: 更新后的完整运行时配置（已合并新值） + """ + async with self._init_lock: + self.logger.info("[%s] 配置热重载开始", self.type) + new_config = dict(config or {}) + self.config = new_config + try: + await self.on_cleanup() + except Exception as e: # noqa: BLE001 - 清理失败不阻断重载 + self.logger.warning("[%s] 热重载清理旧资源异常: %s", self.type, e) + await self.on_init() + self.logger.info("[%s] 配置热重载完成", self.type) + # ------------------------------------------------------------------ # 对外统一接口（app lifespan 直接绑定这两个 bound method） # ------------------------------------------------------------------ @@ -160,7 +181,9 @@ class Plugin: 独立方法便于测试用桩替换（如 test_auth_race / verify_example_plugin）， 避免在测试中依赖真实 zenoh 会话。 """ - return _create_workers(session, config, count, worker_manager) + return _create_workers( + session, config, count, worker_manager, reload_cb=self.reload + ) async def cleanup(self, worker_item: str) -> None: """ diff --git a/src/chongming_worker/runtime.py b/src/chongming_worker/runtime.py index 153e370..0d1dbda 100644 --- a/src/chongming_worker/runtime.py +++ b/src/chongming_worker/runtime.py @@ -7,8 +7,12 @@ import os import asyncio import inspect import logging +import pathlib +import random +import tomllib import orjson -from typing import Any, Dict, List, Optional +import tomli_w +from typing import Any, Awaitable, Callable, Dict, List, Optional import zenoh @@ -23,6 +27,317 @@ logger = logging.getLogger("chongming_worker.runtime") _warned_legacy_payload = False +# ============================================================ +# 单副本管理协议（admin Queryable） +# ============================================================ +# 每个 worker 副本额外注册一个内部 Queryable： +# worker///__admin__ +# 网关（或运维工具）经该 key 下发指令，可精细化控制单个副本： +# - restart 平滑退出当前副本并以新 worker_id 重建（liveliness down/up） +# - stop 平滑退出当前副本且不自动重建 +# 无操作（查询状态）时返回副本实时状态。 +# 该 key 不来自业务 TOML 配置，不参与契约校验与路由注册。 + +ADMIN_CMD_RESTART = "restart" +ADMIN_CMD_STOP = "stop" +ADMIN_CMD_DRAIN = "drain" +ADMIN_CMD_CONFIG_SET = "config_set" +ADMIN_CMD_RESCALE = "rescale" + +# 优雅排空时，等待在途请求完成的最大时长（秒），超时则强制关闭该副本 +DRAIN_INFLIGHT_TIMEOUT = 30.0 + +# 进程级 in-flight 请求计数（key=worker_id）。单线程事件循环内读写，无需加锁。 +# 用于「优雅排空」：drain 摘除 liveliness/queryable 后，等待正在处理的请求完成再退出。 +_inflight: Dict[str, int] = {} + +# 进程级 Worker 池注册表（key=worker_type）。由 create_workers/WorkerPool 在同一进程内 +# 建立，供 rescale 管理指令按目标副本数扩缩当前进程内槽位（inline 与 external 共用）。 +_process_pools: Dict[str, "WorkerPool"] = {} + + +def get_process_pool(worker_type: str) -> Optional["WorkerPool"]: + """获取当前进程内某类型的 Worker 池（未创建返回 None）。""" + return _process_pools.get(worker_type) + + +def _register_process_pool(worker_type: str, pool: "WorkerPool") -> None: + _process_pools[worker_type] = pool + +# 禁止通过 config_set 写入的保留键： +# [workers.config] 若覆盖这些键会破坏插件加载/路由语义（属于"路由配置除外"范畴）。 +CONFIG_FORBIDDEN_KEYS = { + "type", "module", "factory", "cleanup", "count", "load_mode", + "mqtt_topics", "workers", "worker", "defaults", "registration", +} + + +def _plugin_config_path(worker_type: str) -> pathlib.Path: + """插件 config.toml 路径（与 load_plugin_config 的定位规则一致）。""" + return pathlib.Path("config/workers") / worker_type / "config.toml" + + +def _write_plugin_config_section(worker_type: str, updates: Dict[str, Any]) -> Dict[str, Any]: + """把 updates 合并写入插件 config.toml 的 [workers.config] 段（原子写回）。 + + 仅改动 [workers.config] 下的键，保留文件其它段与段落。 + 返回写入后的 [workers.config] 完整 dict（供回显）。 + + Raises: + ValueError: 包含保留键 / config.toml 缺失 / 无 [workers] 段 + """ + bad = [k for k in updates if k in CONFIG_FORBIDDEN_KEYS] + if bad: + raise ValueError(f"禁止修改保留配置键: {bad}（路由/加载控制配置除外，不在此功能范围）") + + path = _plugin_config_path(worker_type) + with open(path, "rb") as f: + data = tomllib.load(f) + + workers = data.get("workers") + if not isinstance(workers, dict): + raise ValueError(f"{path} 缺少 [workers] 节") + + cfg = workers.setdefault("config", {}) + if not isinstance(cfg, dict): + raise ValueError(f"{path} 的 [workers.config] 不是表，无法更新") + cfg.update(updates) + + # 原子写回：临时文件 + rename，避免写一半被读 + tmp = path.with_suffix(".toml.tmp") + with open(tmp, "w", encoding="utf-8") as f: + f.write(tomli_w.dumps(data)) + os.replace(tmp, path) + return dict(cfg) + + +def _write_plugin_count(worker_type: str, count: int) -> int: + """把某类型的副本目标数写入插件 config.toml 的 [workers].count（原子写回）。 + + 由 rescale 指令的处理副本在**其所在进程的工作目录**下执行，因此 inline（网关进程） + 与 external（独立进程/容器、配置目录已挂载）都能持久化到自己那侧的配置，重启后保持。 + 返回写回的计数值。 + + Raises: + ValueError: 文件缺失 / 无 [workers] 节 + """ + path = _plugin_config_path(worker_type) + with open(path, "rb") as f: + data = tomllib.load(f) + + workers = data.get("workers") + if not isinstance(workers, dict): + raise ValueError(f"{path} 缺少 [workers] 节") + + workers["count"] = int(count) + + tmp = path.with_suffix(".toml.tmp") + with open(tmp, "w", encoding="utf-8") as f: + f.write(tomli_w.dumps(data)) + os.replace(tmp, path) + return int(count) + + +def _merge_runtime_config(config: Dict[str, Any], updates: Dict[str, Any]) -> Dict[str, Any]: + """基于本轮 running config 合并 updates，得到重启后应使用的运行时配置。""" + allowed = {k: v for k, v in updates.items() if k not in CONFIG_FORBIDDEN_KEYS} + return {**config, **allowed} + + +class WorkerController: + """单个 worker 副本的跨重启控制对象。 + + 由 create_workers 为每个副本槽位（index）创建、贯穿该槽位的多次重启： + - worker_id = worker---，重启 epoch 递增、id 变化 + - exit_request 记录当前副本被要求退出时的原因 + （""=正常/被取消；"restart"=管理指令重建；"stop"=管理指令停止） + - current_worker_id 为当前运行副本的 id（供 admin 指令回显） + 单个 controller 串行驱动 create_worker 多次运行，实现单副本自愈重启。 + """ + + def __init__(self, worker_type: str, index: int): + self.worker_type = worker_type + self.index = index + self.epoch = 0 + self.exit_request: str = "" + self.current_worker_id: str = "" + # 配置变更后待应用的新运行时配置（由 config_set 设置，重启时生效） + self.pending_config: Optional[Dict[str, Any]] = None + self.pending_updates: Optional[Dict[str, Any]] = None + # 配置热重载回调（入口：config_set 且 index0 时触发，用于重建插件客户端，见 Plugin.reload） + self.reload_cb: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None + # 优雅排空标记：drain 指令置位后，副本先摘除 liveliness/queryable（隔绝新请求）， + # 再等待正在处理的请求完成（in-flight=0）才退出；管理循环据此不重建（数量-1）。 + self.draining: bool = False + # 副本主阻塞事件：create_worker 用它代替局部 exit_event，同样可被 admin 指令与 + # WorkerPool 触发的缩容/排空唤醒（单一复制均通过置位此事件实现优雅退出）。 + self.control_event: asyncio.Event = asyncio.Event() + + def make_worker_id(self) -> str: + """生成下一个副本 id（epoch 自增，保证重启后 id 变化）。""" + self.epoch += 1 + return f"worker-{self.index}-{self.epoch}-{os.getpid()}" + + +async def _handle_admin_command( + query: zenoh.Query, + controller: WorkerController, + worker_type: str, + worker_id: str, + exit_event: Optional[asyncio.Event] = None, + config: Optional[Dict[str, Any]] = None, +) -> None: + """处理对单个副本的 admin 指令（在 worker 事件循环内执行）。 + + 载荷（信封或裸 JSON，优先信封 data 字段）形如 {"cmd": "restart"}。 + + - restart：登记退出原因为 restart 并解除副本主阻塞，管理循环据此重建。 + - stop： 登记退出原因为 stop 并解除主阻塞，管理循环据此停止该副本。 + - config_set：把 data.config 写入插件 config.toml 的 [workers.config]（仅 index0 + 落盘避免同文件并发写，其它副本不落盘），并登记 restart 让本副本以新配置重建。 + 含保留键 / 文件缺失 / 写回失败 -> 返回 {status:"error", error}，不触发重启。 + - 其他/缺失：查询状态，不触发退出，返回副本实时信息。 + """ + cmd = "" + raw = {} + try: + payload = bytes(query.payload) if query.payload else b"" + raw = orjson.loads(payload) if payload else {} + raw = raw.get("data") if isinstance(raw, dict) and "data" in raw else raw + cmd = str((raw or {}).get("cmd", "") or "").lower() + except Exception: + cmd = "" + + applied_updates: Dict[str, Any] = {} + setup_error = "" + persist_note = "" + + if cmd == ADMIN_CMD_RESTART: + controller.exit_request = ADMIN_CMD_RESTART + logger.info("[Worker] %s 收到重启指令（将以新 id 重建）", worker_id) + elif cmd == ADMIN_CMD_STOP: + controller.exit_request = ADMIN_CMD_STOP + logger.info("[Worker] %s 收到停止指令", worker_id) + elif cmd == ADMIN_CMD_DRAIN: + # 优雅排空：登记 stop（数量-1、不重建），由主循环摘除 liveliness/queryable + # 后等待 in-flight 请求完成再退出。 + controller.exit_request = ADMIN_CMD_STOP + controller.draining = True + logger.info("[Worker] %s 收到优雅排空指令（将隔离新请求并等待在途请求完成）", worker_id) + elif cmd == ADMIN_CMD_CONFIG_SET: + updates = raw.get("config") + if not isinstance(updates, dict) or not updates: + setup_error = "config_set 需要 data.config 字段（非空对象）" + else: + try: + # 仅 index0 落盘写回，其余副本不落盘（同一进程多副本/共享文件避免并发写） + if controller.index == 0: + new_section = _write_plugin_config_section(worker_type, updates) + applied_updates = { + k: v for k, v in updates.items() + if k not in CONFIG_FORBIDDEN_KEYS + } + logger.info( + "[Worker] %s 已将配置写入 %s: %s", + worker_id, _plugin_config_path(worker_type), new_section, + ) + else: + applied_updates = { + k: v for k, v in updates.items() + if k not in CONFIG_FORBIDDEN_KEYS + } + # 热重载：index0 副本触发插件重新初始化（重建客户端等），使新配置立即生效 + if controller.index == 0 and controller.reload_cb is not None: + await controller.reload_cb( + _merge_runtime_config(config or {}, applied_updates) + ) + logger.info( + "[Worker] %s 配置热重载完成，将以新配置重启", worker_id, + ) + # 待应用配置：合并进本轮 running config，重启后以新配置运行 + controller.pending_updates = dict(applied_updates) + controller.pending_config = _merge_runtime_config( + config or {}, applied_updates + ) + controller.exit_request = ADMIN_CMD_RESTART + logger.info( + "[Worker] %s 收到配置变更指令，将以新配置重启（%d 个键）", + worker_id, len(applied_updates), + ) + except Exception as e: # noqa: BLE001 - 写回/校验失败按指令失败返回 + setup_error = f"配置写入失败: {e}" + logger.error("[Worker] %s %s", worker_id, setup_error) + elif cmd == ADMIN_CMD_RESCALE: + # 调整当前进程内该类型的副本槽位数（缩容走优雅排空）。作用于本进程 Worker 池， + # 不退出本副本自身；跨进程（external）指令经 Zenoh 到达其容器进程内的指定副本 → 其进程池扩缩。 + try: + target = int(raw.get("scale", -1)) + except (TypeError, ValueError): + target = -1 + if target < 0: + setup_error = "rescale 需要 data.scale（非负整数目标副本数）" + else: + pool = get_process_pool(worker_type) + if pool is None: + setup_error = f"当前进程无 {worker_type} 类型的 Worker 池，无法缩放" + else: + asyncio.ensure_future(pool.set_target(target)) + logger.info( + "[Worker] %s 收到缩放指令，目标副本数=%d（当前池=%d）", + worker_id, target, len(pool._slots), + ) + # 持久化目标数量到本进程侧配置（重启后保持），失败不影响本次内存缩放 + try: + _write_plugin_count(worker_type, target) + persist_note = f"count={target} 已持久化" + except Exception as e: # noqa: BLE001 + persist_note = f"count 持久化失败: {e}" + logger.warning("[Worker] %s %s", worker_id, persist_note) + else: + query.reply( + key_expr=query.key_expr, + payload=orjson.dumps({ + "worker_type": worker_type, + "worker_id": worker_id, + "index": controller.index, + "epoch": controller.epoch, + "status": "running", + }), + ) + return + + if setup_error: + query.reply( + key_expr=query.key_expr, + payload=orjson.dumps({ + "status": "error", + "cmd": cmd, + "worker_type": worker_type, + "worker_id": worker_id, + "error": setup_error, + }), + ) + return + + query.reply( + key_expr=query.key_expr, + payload=orjson.dumps({ + "status": "accepted", + "cmd": cmd, + "worker_type": worker_type, + "worker_id": worker_id, + "exit_request": controller.exit_request, + "persist": persist_note, + }), + ) + + # 解除该副本主阻塞（exit_event.wait），使其按 exit_request 收尾并交给管理循环重建/停止 + if exit_event is not None and cmd in ( + ADMIN_CMD_RESTART, ADMIN_CMD_STOP, ADMIN_CMD_DRAIN, ADMIN_CMD_CONFIG_SET, + ): + exit_event.set() + + async def stream_replies_to_query(query: zenoh.Query, result) -> None: """ 将生成器结果逐块回复给查询方（Zenoh 多回复流式）。 @@ -70,19 +385,25 @@ async def create_worker( worker_id: str, config: Optional[Dict[str, Any]] = None, worker_manager=None, + controller: Optional[WorkerController] = None, ): """ - 创建一个 Worker 协程 + 创建一个 Worker 协程（单次运行，配合 controller 可实现单副本自愈重启） 自动发现通过 @rpc_method 注册的端点，为每个 subject 声明 Zenoh Queryable， 收到请求后按信封 meta.method 字段分派到具体方法。 + 当传入 controller（由 create_workers 的管理循环创建）时，额外注册一个内部 + admin Queryable（worker///__admin__），支持对单个副本下发 + restart / stop 指令；退出原因通过 controller.exit_request 回传管理循环。 + Args: session: Zenoh 会话 worker_type: Worker 类型标识（匹配 subject 的第一部分） worker_id: Worker 唯一 ID config: 自定义配置参数 worker_manager: WorkerManager 实例（用于本地注册） + controller: WorkerController（单个副本多次运行共用，用于管理指令与重启控制） """ logger.info("[Worker] %s 启动 (type=%s)", worker_id, worker_type) @@ -107,14 +428,19 @@ async def create_worker( for registry_key, rpc in rpc_methods.items(): subject_methods.setdefault(rpc.subject, []).append(rpc) + loop = asyncio.get_running_loop() + # 副本退出事件：admin 指令（restart/stop/drain）、WorkerPool 缩容触发时置位， + # 解除本副本主阻塞。有 controller 时复用其 control_event，便于池内缩容触达。 + exit_event = controller.control_event if controller is not None else asyncio.Event() + # 5. 为每个 subject 创建回调函数并声明 Queryable - stop_event = asyncio.Event() def make_callback(subject: str, methods: List[RPCMethod]): """创建处理特定 subject 的回调函数""" - loop = asyncio.get_running_loop() async def _handle_query(query: zenoh.Query): matched = None # 方法匹配结果（异常分支判断流式错误格式用） + # in-flight 计数：供 drain 排在完成后按需退出 + _inflight[worker_id] = _inflight.get(worker_id, 0) + 1 try: # 解析请求负载（信封协议: {"meta": {...}, "data": {...}}） payload_bytes = bytes(query.payload) if query.payload else b"" @@ -185,6 +511,11 @@ async def create_worker( payload=orjson.dumps(error_response) ) logger.error("[Worker] %s 异常 %s: %s", worker_id, subject, e) + finally: + # in-flight 完成，递减计数（可能已 < 1 允许负数并 clamps） + _inflight[worker_id] = _inflight.get(worker_id, 0) - 1 + if _inflight[worker_id] <= 0: + _inflight.pop(worker_id, None) def callback(query: zenoh.Query): # 当前 create_worker 运行在事件循环中，可直接 create_task @@ -192,7 +523,7 @@ async def create_worker( return callback - # 声明所有 Queryable + # 声明所有业务 Queryable queryables: List[zenoh.Queryable] = [] for subject, methods in subject_methods.items(): key = f"worker/{worker_type}/{worker_id}/{subject}" @@ -201,63 +532,266 @@ async def create_worker( queryables.append(q) logger.info("[Worker] %s Queryable: '%s' [%s]", worker_id, key, ", ".join(m.method for m in methods)) - # 6. 保持 worker 存活直到被取消 - stop_event = asyncio.Event() + # 声明 admin Queryable（管理指令 > restart / stop / 状态查询）。仅在有 controller 时注册， + # 直接独立调用 create_worker（controller=None）保持原有行为、不暴露管理通道。 + admin_queryable = None + if controller is not None: + admin_key = f"worker/{worker_type}/{worker_id}/__admin__" + + def make_admin_callback(_admin_key: str): + async def _handle_admin(query: zenoh.Query): + await _handle_admin_command( + query, controller, worker_type, worker_id, exit_event, config + ) + + def admin_cb(query: zenoh.Query): + asyncio.run_coroutine_threadsafe(_handle_admin(query), loop) + + return admin_cb + + admin_queryable = session.declare_queryable(admin_key, make_admin_callback(admin_key)) + logger.info("[Worker] %s admin Queryable: '%s'", worker_id, admin_key) + + # 6. 保持 worker 存活直到被取消或收到管理指令 try: - await stop_event.wait() + await exit_event.wait() except asyncio.CancelledError: logger.info("[Worker] %s 正在关闭...", worker_id) except Exception as e: logger.error("[Worker] %s 发生异常: %s", worker_id, e) finally: - # 清理 + # 退出非由管理指令触发（正常返回 / 取消 / 异常）时，标记 stop —— 管理循环据此不自动重启 + if controller is not None and not controller.exit_request: + controller.exit_request = ADMIN_CMD_STOP + # 清理：先摘除 liveliness + 业务 Queryable，隔绝新请求 if worker_manager is not None: worker_manager.unregister_worker(worker_type, worker_id) token.undeclare() for q in queryables: q.undeclare() + # 优雅排空：等正在处理的请求（in-flight）完成后再关闭（带超时，超时强制退出） + if controller is not None and controller.draining: + await _wait_inflight_zero(worker_id, timeout=DRAIN_INFLIGHT_TIMEOUT) + if admin_queryable is not None: + admin_queryable.undeclare() logger.info("[Worker] %s 已关闭", worker_id) + +async def _wait_inflight_zero(worker_id: str, timeout: float = 30.0) -> None: + """等待指定 worker 的 in-flight 请求计数归零（用于优雅排空）。 + + 摘除 liveliness/业务 Queryable 之后调用：此时不再接收新请求，只需等 + 正在处理的协程完成；超时则仅告警、强制返回（避免卡死关闭流程）。 + """ + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + remain = _inflight.get(worker_id, 0) + while remain and remain > 0: + if loop.time() >= deadline: + logger.warning( + "[Worker] %s 优雅排空超时（in-flight=%d），强制关闭", + worker_id, remain, + ) + return + await asyncio.sleep(0.05) + remain = _inflight.get(worker_id, 0) + logger.info("[Worker] %s 在途请求已清空，完成排空", worker_id) + +async def _managed_worker( + session: zenoh.Session, + worker_type: str, + controller: WorkerController, + config: Optional[Dict[str, Any]], + worker_manager=None, + reload_cb: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None, +) -> None: + """单个副本槽位（index）的管理循环：串行驱动该副本多次运行，支持自愈重启。 + + - worker_id 由 controller 按 epoch 生成，重启后变化（liveliness 反映 down/up） + - create_worker 正常返回后，按 controller.exit_request 决定： + restart → 生成新 id 重建该副本；否则（stop/正常/取消）→ 结束循环 + - 外部取消（网关/进程关闭）冒泡传播，保证下游调用方能正常收尾 + """ + while True: + controller.exit_request = "" + controller.current_worker_id = controller.make_worker_id() + controller.reload_cb = reload_cb + # 若存在待应用的新配置（config_set 设置），本次重建使用新配置运行 + run_config = controller.pending_config if controller.pending_config else config + logger.info( + "[Worker] %s#%d 启动管理循环（epoch=%d, id=%s）", + worker_type, controller.index, controller.epoch, controller.current_worker_id, + ) + try: + await create_worker( + session=session, + worker_type=worker_type, + worker_id=controller.current_worker_id, + config=run_config, + worker_manager=worker_manager, + controller=controller, + ) + except asyncio.CancelledError: + logger.info("[Worker] %s#%d 管理循环被取消", worker_type, controller.index) + raise + if controller.exit_request == ADMIN_CMD_RESTART: + logger.info( + "[Worker] %s#%d 副本 %s 收到重启指令，将以新 id 重建", + worker_type, controller.index, controller.current_worker_id, + ) + continue + logger.info( + "[Worker] %s#%d 副本 %s 退出（%s），管理循环结束", + worker_type, controller.index, controller.current_worker_id, + controller.exit_request or "normal", + ) + break + + def create_workers( session: zenoh.Session, config: dict, count: int = 1, worker_manager=None, + reload_cb: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None, ) -> list[asyncio.Task]: """ - 创建一组 Worker 协程任务 + 创建一组 Worker 副本管理任务 - 工厂函数，供 WorkerManager 动态加载使用。 + 工厂函数，供 WorkerManager / 进程插件加载使用。为每个副本槽位（index）创建 + 一个管理循环 task，该 task 串行承载此副本的多次运行（重启后生成新 worker_id）。 + + 返回的 task 与原有契约一致：调用方 crawl 取消（如网关/进程关闭）即可优雅停掉 + 该副本及其后续重启；副本重启不会导致 task 结束。 Args: session: Zenoh 会话 config: Worker 配置字典（至少需要 "type" 字段） count: 启动的 Worker 副本数 worker_manager: WorkerManager 实例（可选） + reload_cb: 配置热重载回调（config_set 且 index0 时触发，重建插件客户端） Returns: - list[asyncio.Task]: Worker 协程任务列表 + list[asyncio.Task]: 每个副本管理循环的 Task 列表 """ worker_type = config.get("type", "default") - pid = os.getpid() - tasks = [] - for i in range(count): - worker_id = f"worker-{i}-{pid}" + # 复用当前进程该类型的池（存在则不动已有槽位），确保 rescale 能统一触达本副本 + pool = get_process_pool(worker_type) + if pool is None: + pool = WorkerPool( + session, worker_type, config, worker_manager=worker_manager, reload_cb=reload_cb + ) + _register_process_pool(worker_type, pool) + return pool.start(count) + + +class WorkerPool: + """某类型在当前进程内的副本槽位池，支持按目标数量动态扩缩容。 + + 管理该类型在本进程的全部副本（index 槽位），每个槽位由一个 _managed_worker 管理 + 循环承载（内含自愈重启）。用于实现「同类 worker 数量管理」： + + - 扩容：新增槽位（新 controller + 管理循环 task） + - 缩容：随机挑选多余槽位，触发**优雅排空**（设置 exit_request=STOP + draining， + 通过 control_event 唤醒副本，摘除 liveliness/业务 Queryable 隔绝新请求， + 等待 in-flight 请求完成后再退出；管理循环见 exit=stop 不再重建） + + inline 与 external（进程插件）共用此语义：网关经 Zenoh 向目标副本下发 rescale + 指令，副本在**其所在进程内**对本类型池 set_target 扩缩——无需重启整个进程。 + """ + + def __init__( + self, + session: zenoh.Session, + worker_type: str, + config: Optional[Dict[str, Any]] = None, + worker_manager=None, + reload_cb: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None, + ): + self.session = session + self.worker_type = worker_type + self.config = config or {} + self.worker_manager = worker_manager + self.reload_cb = reload_cb + # index -> controller / task（仅记录当前活动槽位；drain 结束后由 done 回调移除） + self._slots: Dict[int, WorkerController] = {} + self._tasks: Dict[int, asyncio.Task] = {} + self._next_index = 0 + self._target = 0 + + def _new_slot(self) -> WorkerController: + """为本类型创建/追加一个副本槽位（index 递增，避免与已有槽位、历史 index 冲突）。""" + index = self._next_index + self._next_index += 1 + controller = WorkerController(self.worker_type, index) task = asyncio.create_task( - create_worker( - session=session, - worker_type=worker_type, - worker_id=worker_id, - config=config, - worker_manager=worker_manager, + _managed_worker( + self.session, self.worker_type, controller, self.config, + self.worker_manager, self.reload_cb, ) ) - tasks.append(task) - logger.info("[Worker] 已创建任务: %s (type=%s)", worker_id, worker_type) - return tasks + self._slots[index] = controller + self._tasks[index] = task + task.add_done_callback(lambda _t: self._drop_slot(index)) + logger.info("[Worker] %s 池扩增副本: index=%d", self.worker_type, index) + return controller + + def _drop_slot(self, index: int) -> None: + """槽位管理循环结束（drain/stop 后不重建）时从其登记中移除。""" + self._slots.pop(index, None) + self._tasks.pop(index, None) + logger.info("[Worker] %s 池移除副本: index=%d", self.worker_type, index) + + def start(self, count: int) -> list[asyncio.Task]: + """按 count 初始化目标副本数，返回当前活动槽位管理循环 Task 列表（兼容旧契约）。""" + self._target = count + already = len(self._slots) + for _ in range(max(0, count - already)): + self._new_slot() + return list(self._tasks.values()) + + async def set_target(self, count: int) -> int: + """将本类型副本数调整为 count（扩容新增槽位；缩容随机优雅排空多余槽位）。 + + 返回调整后的活动槽位数。 + """ + count = max(0, int(count)) + self._target = count + # 扩容：补齐缺失槽位 + while len(self._slots) < count: + self._new_slot() + # 缩容：随机挑选多余槽位，逐个优雅排空 + overflow = list(self._slots.keys()) + random.shuffle(overflow) + for index in overflow: + if len(self._slots) <= count: + break + controller = self._slots[index] + controller.exit_request = ADMIN_CMD_STOP + controller.draining = True + controller.control_event.set() + logger.info( + "[Worker] %s 触发副本 %s 优雅排空（目标 %d）", + self.worker_type, controller.current_worker_id or index, count, + ) + task = self._tasks.get(index) + if task is not None: + try: + await asyncio.wait_for(task, timeout=DRAIN_INFLIGHT_TIMEOUT + 5) + except asyncio.TimeoutError: + logger.warning("[Worker] %s index=%d 排空超时", self.worker_type, index) + except asyncio.CancelledError: + raise + return len(self._slots) + + async def shutdown(self) -> None: + """优雅关闭：将所有副本数量缩至 0。""" + await self.set_target(0) __all__ = [ "create_worker", "create_workers", + "WorkerPool", + "get_process_pool", ]