modify: worker 管理端点
This commit is contained in:
+2
-1
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
+559
-25
@@ -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/<type>/<worker_id>/__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-<index>-<epoch>-<pid>，重启 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/<type>/<worker_id>/__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",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user