From aa64e9bd34d14856a652fe87be0932bd896535f0 Mon Sep 17 00:00:00 2001 From: chakcy <947105045@qq.com> Date: Tue, 13 Oct 2026 01:59:16 +0800 Subject: [PATCH] =?UTF-8?q?modify:=20=E6=B4=BE=E5=8F=91=E9=80=BB=E8=BE=91?= =?UTF-8?q?=E7=BB=9F=E4=B8=80=E5=88=B0=20dispatch=EF=BC=8C=E4=BE=9B=20zeno?= =?UTF-8?q?h=20=E5=9B=9E=E8=B0=83=E4=B8=8E=E8=BF=9B=E7=A8=8B=E5=86=85?= =?UTF-8?q?=E7=9B=B4=E8=B0=83=E5=85=B1=E7=94=A8=EF=BC=880.5.0=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 2 +- src/chongming_worker/dispatch.py | 333 +++++++++++++++++++++++++++++++ src/chongming_worker/runtime.py | 165 ++++----------- 3 files changed, 378 insertions(+), 122 deletions(-) create mode 100644 src/chongming_worker/dispatch.py diff --git a/pyproject.toml b/pyproject.toml index b8e9665..d8f105a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = ["hatchling"] [project] name = "chongming-worker" -version = "0.4.5" +version = "0.5.0" description = "chongming worker" readme = "README.md" requires-python = ">=3.12" diff --git a/src/chongming_worker/dispatch.py b/src/chongming_worker/dispatch.py new file mode 100644 index 0000000..83f2928 --- /dev/null +++ b/src/chongming_worker/dispatch.py @@ -0,0 +1,333 @@ +"""进程内派发器：让「zenoh 回调」与「网关直调」共用同一份业务派发逻辑。 + +背景：inline worker 与网关同进程、同事件循环，但业务请求仍要走 +"zenoh 查询 → router → queryable 回调 → 跨线程调度回原 loop"，一次 RPC 多出 +线程 hop 与两次 TCP 往返。本模块把派发逻辑从 zenoh Query 上解耦出来 +（信封解析 → 方法匹配 → 调用 → 回复），由两类 sink 驱动： + + - ZenohQuerySink ：包装 zenoh.Query，回调路径语义与改造前逐字一致 + - _LocalUnarySink / _LocalStreamSink：进程内直调（网关侧快路径） + +只换传输、不改语义：信封协议、错误通道（reply_err）、流式多回复与错误收尾、 +in-flight 计数（优雅排空据此等待在途请求）、user_info 注入全部保持。 +external（独立进程）插件不受影响：它在自己进程里注册，网关侧查不到 → 仍走 zenoh。 +""" + +import asyncio +import inspect +import logging +from typing import Any, AsyncIterator, Callable, Dict, List, Optional, Protocol, Tuple + +import orjson + +logger = logging.getLogger("chongming_worker.dispatch") + + +# ── in-flight 计数：worker_id → 在途请求数（graceful drain 等它归零）── +_inflight: Dict[str, int] = {} + + +def _inc_inflight(worker_id: str) -> None: + _inflight[worker_id] = _inflight.get(worker_id, 0) + 1 + + +def _dec_inflight(worker_id: str) -> None: + _inflight[worker_id] = _inflight.get(worker_id, 0) - 1 + if _inflight[worker_id] <= 0: + _inflight.pop(worker_id, None) + + +# 非信封载荷只告警一次（与改造前同） +_warned_legacy_payload = False + + +class ReplySink(Protocol): + """回复出口：zenoh Query 与进程内直调都归一到这两个方法""" + + def reply(self, key_expr: str, payload: bytes) -> bool: + """单块回复；实现方可用抛出异常表示"接收端已终止"（流式据此停流）""" + + def reply_err(self, payload: bytes) -> None: + """框架错误回复（网关零解析判定失败 → 503）""" + + +class ZenohQuerySink: + """把 zenoh.Query 适配为 ReplySink（语义与改造前逐字一致）""" + + __slots__ = ("_query",) + + def __init__(self, query: Any) -> None: + self._query = query + + def reply(self, key_expr: str, payload: bytes) -> bool: + self._query.reply(key_expr=key_expr, payload=payload) + return True + + def reply_err(self, payload: bytes) -> None: + self._query.reply_err(payload=payload) + + +async def stream_replies_to_sink(sink: ReplySink, key_expr: str, result) -> None: + """把一个生成器结果逐块回复给 sink（zenoh 多回复 与 进程内流式 共用）。 + + 约定（与网关 call_stream / SSE 端点配合）： + - 生成器正常结束 = 流结束，无结束哨兵 + - 生成器中途异常 → 发送 {"error": ...} 收尾块后停流 + - sink.reply 失败（查询已终止/消费方已停）→ 立即停流，避免继续生成 + """ + async def _reply(chunk) -> bool: + try: + sink.reply(key_expr, orjson.dumps(chunk)) + return True + except Exception as e: + logger.debug("[Worker] 流式回复失败（接收端可能已终止）: %s", e) + return False + + try: + if inspect.isasyncgen(result): + async for chunk in result: + if not await _reply(chunk): + break + else: + for chunk in result: + if not await _reply(chunk): + break + except Exception as e: + # 生成器中途异常：以 {"error": ...} 收尾块结束流 + logger.error("[Worker] 流式生成器异常，发送错误块后停止: %s", e) + await _reply({"error": str(e)}) + + +class SubjectDispatch: + """单个 subject 的派发器（一个 subject 可挂 GET/POST 等多个 RPCMethod）""" + + def __init__( + self, + subject: str, + methods: List[Any], + key_expr: str, + session: Any, + worker_type: str, + worker_id: str, + ) -> None: + self.subject = subject + self.methods = methods + self.key_expr = key_expr + self.session = session + self.worker_type = worker_type + self.worker_id = worker_id + + async def dispatch(self, payload_bytes: bytes, sink: ReplySink) -> None: + """信封解析 → 方法匹配 → 调用 → 回复（与改造前 runtime._handle_query 等价）""" + global _warned_legacy_payload + matched = None # 方法匹配结果（异常分支判断流式错误格式用） + # in-flight 计数：供 drain 排在完成后按需退出 + _inc_inflight(self.worker_id) + try: + # 解析请求负载（信封协议: {"meta": {...}, "data": {...}}） + raw = orjson.loads(payload_bytes) if payload_bytes else {} + + if isinstance(raw, dict) and ("meta" in raw or "data" in raw): + meta = raw.get("meta") or {} + data = raw.get("data") or {} + else: + # 非信封载荷: 整包视为业务数据，method 默认 GET（无控制字段语义） + if not _warned_legacy_payload: + logger.warning( + "收到非信封格式载荷（缺少 meta/data），按 method=GET 处理: %s", + self.subject, + ) + _warned_legacy_payload = True + meta = {} + data = raw if isinstance(raw, dict) else {} + http_method = str(meta.get("method", "GET")).upper() + + # 匹配方法（按 HTTP 方法） + matched = next( + (m for m in self.methods if m.method.upper() == http_method), None + ) + if not matched and self.methods: + matched = self.methods[0] # 降级到第一个 + if not matched: + sink.reply_err(orjson.dumps({"error": f"No method for {self.subject}"})) + return + + # 调用 RPC 方法（支持同步和异步；返回生成器时进入流式模式） + # user_info 从 meta 信封注入（框架参数，不进业务 payload） + result = await matched.call( + data, + session=self.session, + worker_id=self.worker_id, + user_info=meta.get("user_info"), + ) + + if inspect.isasyncgen(result) or inspect.isgenerator(result): + # 流式模式：@rpc_method 函数返回生成器，逐块多回复推送 + await stream_replies_to_sink(sink, self.key_expr, result) + else: + # 一次性回复 + sink.reply(self.key_expr, orjson.dumps(result)) + logger.debug( + "[Worker] %s 回复 %s %s", self.worker_id, matched.method, self.subject + ) + except Exception as e: + if matched is not None and matched.stream: + # 流式端点错误：以 {"error": ...} 收尾块结束流（流式语义不变） + sink.reply(self.key_expr, orjson.dumps({"error": str(e)})) + else: + error_response = { + "result": 0.0, + "operation": "error", + "error": str(e), + "worker": self.worker_id, + } + # 框架错误走 reply_err 通道：网关零解析判定失败（映射 503）； + # 载荷保持原错误信封结构，旧网关仍按普通回复透传（兼容） + sink.reply_err(orjson.dumps(error_response)) + logger.error("[Worker] %s 异常 %s: %s", self.worker_id, self.subject, e) + finally: + # in-flight 完成，递减计数（可能已 < 1 允许负数并 clamps） + _dec_inflight(self.worker_id) + + +class LocalInvokeResult: + """进程内直调结果：ok=正常回复；ok=False 对应 zenoh 的 reply_err（网关映射 503）""" + + __slots__ = ("ok", "payload") + + def __init__(self, ok: bool, payload: bytes) -> None: + self.ok = ok + self.payload = payload + + +class _StreamClosed(Exception): + """消费方已停止消费（等价 zenoh 路径里 query.reply 失败）""" + + +class _LocalUnarySink: + """直调-单回复 sink：首个回复胜出（与网关取首个 zenoh reply 一致）""" + + __slots__ = ("result",) + + def __init__(self) -> None: + self.result: Optional[LocalInvokeResult] = None + + def reply(self, key_expr: str, payload: bytes) -> bool: + if self.result is None: + self.result = LocalInvokeResult(True, bytes(payload)) + return True + + def reply_err(self, payload: bytes) -> None: + if self.result is None: + self.result = LocalInvokeResult(False, bytes(payload)) + + +class _LocalStreamSink: + """直调-流式 sink：块投进队列；消费方退出后 reply 抛 _StreamClosed 让生成器尽早停流""" + + __slots__ = ("_queue", "_closed") + + def __init__(self, queue: "asyncio.Queue") -> None: + self._queue = queue + self._closed = False + + def close(self) -> None: + self._closed = True + + def reply(self, key_expr: str, payload: bytes) -> bool: + if self._closed: + raise _StreamClosed("consumer stopped") + self._queue.put_nowait(bytes(payload)) + return True + + def reply_err(self, payload: bytes) -> None: + if self._closed: + raise _StreamClosed("consumer stopped") + self._queue.put_nowait(bytes(payload)) + + +class WorkerDispatcher: + """一个 worker 副本的派发器集合（按 (worker_type, worker_id) 注册到本进程注册表）""" + + def __init__(self, session: Any, worker_type: str, worker_id: str) -> None: + self.session = session + self.worker_type = worker_type + self.worker_id = worker_id + self._subjects: Dict[str, SubjectDispatch] = {} + + def bind(self, subject: str, methods: List[Any], key_expr: str) -> SubjectDispatch: + dispatcher = SubjectDispatch( + subject, methods, key_expr, self.session, self.worker_type, self.worker_id + ) + self._subjects[subject] = dispatcher + return dispatcher + + def has_subject(self, subject: str) -> bool: + return subject in self._subjects + + async def invoke_unary(self, subject: str, payload_bytes: bytes) -> LocalInvokeResult: + """进程内直调（单回复）：返回 ok/err 结果，不抛业务异常（框架内部已兜底）""" + dispatcher = self._subjects.get(subject) + if dispatcher is None: + return LocalInvokeResult(False, orjson.dumps({"error": f"No subject {subject}"})) + sink = _LocalUnarySink() + await dispatcher.dispatch(payload_bytes, sink) + if sink.result is None: # 理论不可达：dispatch 必走 reply 或 reply_err + return LocalInvokeResult(False, orjson.dumps({"error": "no reply"})) + return sink.result + + async def invoke_stream(self, subject: str, payload_bytes: bytes) -> AsyncIterator[bytes]: + """进程内直调（流式）：逐块产出 JSON 字节；错误块沿用 {"error": ...} 约定""" + dispatcher = self._subjects.get(subject) + if dispatcher is None: + yield orjson.dumps({"error": f"No subject {subject}"}) + return + + queue: "asyncio.Queue" = asyncio.Queue() + sink = _LocalStreamSink(queue) + _DONE = object() + + async def _run() -> None: + try: + await dispatcher.dispatch(payload_bytes, sink) + finally: + queue.put_nowait(_DONE) + + task = asyncio.create_task(_run()) + task.add_done_callback(lambda t: None if t.cancelled() else t.exception()) + try: + while True: + item = await queue.get() + if item is _DONE: + return + yield item + finally: + # 消费方提前退出（SSE 断连/网关放弃）：停流并取消在途派发 + sink.close() + if not task.done(): + task.cancel() + + +# ── 进程内注册表（仅本进程有效）── +# 网关进程里的 inline 副本会在此注册；external 副本注册在自己的进程里，网关查不到 → 走 zenoh +_DISPATCHERS: Dict[Tuple[str, str], WorkerDispatcher] = {} + + +def register_dispatcher(dispatcher: WorkerDispatcher) -> None: + _DISPATCHERS[(dispatcher.worker_type, dispatcher.worker_id)] = dispatcher + + +def unregister_dispatcher(worker_type: str, worker_id: str) -> None: + _DISPATCHERS.pop((worker_type, worker_id), None) + + +def get_dispatcher(worker_type: str, worker_id: str) -> Optional[WorkerDispatcher]: + return _DISPATCHERS.get((worker_type, worker_id)) + + +def parse_worker_key(key_expr: str) -> Optional[Tuple[str, str]]: + """'worker//' → (type, id)；格式不符返回 None""" + parts = key_expr.split("/") + if len(parts) != 3 or parts[0] != "worker" or not parts[1] or not parts[2]: + return None + return parts[1], parts[2] diff --git a/src/chongming_worker/runtime.py b/src/chongming_worker/runtime.py index 31f05f5..9815bb1 100644 --- a/src/chongming_worker/runtime.py +++ b/src/chongming_worker/runtime.py @@ -25,13 +25,18 @@ from chongming_worker.decorators import ( RPCMethod, get_rpc_methods, ) +# 派发逻辑统一在 dispatch.py：zenoh 回调与网关进程内直调共用同一份实现 +from chongming_worker.dispatch import ( + _inflight, + register_dispatcher, + stream_replies_to_sink, + unregister_dispatcher, + WorkerDispatcher, + ZenohQuerySink, +) logger = logging.getLogger("chongming_worker.runtime") -# 是否已提示过"非信封载荷"警告（每个进程仅提示一次，避免刷屏） -_warned_legacy_payload = False - - # ============================================================ # 单副本管理协议（admin Queryable） # ============================================================ @@ -53,9 +58,8 @@ ADMIN_CMD_RESCALE = "rescale" # 优雅排空时，等待在途请求完成的最大时长（秒），超时则强制关闭该副本 DRAIN_INFLIGHT_TIMEOUT = 30.0 -# 进程级 in-flight 请求计数（key=worker_id）。单线程事件循环内读写，无需加锁。 -# 用于「优雅排空」：drain 摘除 liveliness/queryable 后，等待正在处理的请求完成再退出。 -_inflight: Dict[str, int] = {} +# 进程级 in-flight 请求计数（key=worker_id）定义在 dispatch.py（与派发逻辑同处一处， +# 直调与 zenoh 两条路径共用同一份计数）；此处从那里导入，供 _wait_inflight_zero 读取。 # 进程级 Worker 池注册表（key=worker_type）。由 create_workers/WorkerPool 在同一进程内 # 建立，供 rescale 管理指令按目标副本数扩缩当前进程内槽位（inline 与 external 共用）。 @@ -171,18 +175,24 @@ def set_current_instance_id(iid: Optional[str]) -> None: def current_instance_id() -> str: - """外部 worker 进程的稳定全局实例标识。 + """worker 进程的稳定全局实例标识。 优先取进程内显式设定的身份（process.py --instance）；其次 env - CHONGMING_WORKER_INSTANCE（编排时显式给定、跨重启稳定）；缺省回退 - "{hostname}:{pid}"——同一台机器上并行多个未指定身份的实例时，凭 pid 区分， - 避免它们被误并为一个参与方、导致副本翻倍或 UI 只显示一个实例。 - 注意：缺省 id 随进程重启变化；若需跨重启稳定且被逐实例精细扩缩持续命中， - 请显式设 --instance / CHONGMING_WORKER_INSTANCE。 + CHONGMING_WORKER_INSTANCE（编排时显式给定、跨重启稳定）；再次 + CHONGMING_GATEWAY_ID——inline 副本就加载在网关进程内，容器环境（如 K8s）通常 + 给网关进程注入了这个 id，此时网关侧的实例定向用同一 id 即可命中这些副本， + 无需再为 worker 侧单独配置；缺省回退 "{hostname}:{pid}"——同一台机器上并行 + 多个未指定身份的实例时，凭 pid 区分，避免它们被误并为一个参与方、导致副本 + 翻倍或 UI 只显示一个实例。 + 注意：缺省 id 随进程重启变化；若需跨重启稳定且被逐实例精细扩缩/实例定向 + 持续命中，请显式设 --instance / CHONGMING_WORKER_INSTANCE。 """ if _current_instance_id: return _current_instance_id - iid = os.environ.get("CHONGMING_WORKER_INSTANCE") + iid = ( + os.environ.get("CHONGMING_WORKER_INSTANCE") + or os.environ.get("CHONGMING_GATEWAY_ID") + ) if not iid: iid = f"{socket.gethostname()}:{os.getpid()}" return iid @@ -790,35 +800,14 @@ async def stream_replies_to_query(query: zenoh.Query, result) -> None: - query.reply 失败（查询已终止，如客户端断开/网关放弃）-> 立即停流， 避免 Worker 对已无接收方的查询继续生成 + 实现已统一到 dispatch.stream_replies_to_sink（与进程内直调共用同一份逻辑）， + 本函数保留为兼容入口。 + Args: query: Zenoh Query 对象（仅使用 key_expr 与 reply） result: 异步或同步生成器，逐块 yield 可 JSON 序列化的对象 """ - async def _reply(chunk) -> bool: - """单块回复；返回 False 表示查询已终止，应停流""" - try: - query.reply( - key_expr=query.key_expr, - payload=orjson.dumps(chunk), - ) - return True - except Exception as e: - logger.debug("[Worker] 流式回复失败（查询可能已终止）: %s", e) - return False - - try: - if inspect.isasyncgen(result): - async for chunk in result: - if not await _reply(chunk): - break - else: - for chunk in result: - if not await _reply(chunk): - break - except Exception as e: - # 生成器中途异常：以 {"error": ...} 收尾块结束流 - logger.error("[Worker] 流式生成器异常，发送错误块后停止: %s", e) - await _reply({"error": str(e)}) + await stream_replies_to_sink(ZenohQuerySink(query), query.key_expr, result) async def create_worker( session: zenoh.Session, @@ -875,88 +864,18 @@ async def create_worker( exit_event = controller.control_event if controller is not None else asyncio.Event() # 5. 为每个 subject 创建回调函数并声明 Queryable + # + # 派发逻辑集中在 WorkerDispatcher（dispatch.py）：zenoh 回调只做 + # "取载荷 → 交给 dispatcher"，网关侧的进程内直调走同一个 dispatcher， + # 两条路径共用信封解析/方法匹配/错误通道/流式语义与 in-flight 计数。 + dispatcher = WorkerDispatcher(session, worker_type, worker_id) - def make_callback(subject: str, methods: List[RPCMethod]): - """创建处理特定 subject 的回调函数""" - 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"" - raw = orjson.loads(payload_bytes) if payload_bytes else {} + def make_callback(subject: str, bound) -> Callable[[zenoh.Query], None]: + """创建处理特定 subject 的回调函数（bound = 已绑定的 SubjectDispatch）""" - if isinstance(raw, dict) and ("meta" in raw or "data" in raw): - meta = raw.get("meta") or {} - data = raw.get("data") or {} - else: - # 非信封载荷: 整包视为业务数据，method 默认 GET（无控制字段语义） - global _warned_legacy_payload - if not _warned_legacy_payload: - logger.warning( - "收到非信封格式载荷（缺少 meta/data），按 method=GET 处理: %s", - subject, - ) - _warned_legacy_payload = True - meta = {} - data = raw if isinstance(raw, dict) else {} - http_method = str(meta.get("method", "GET")).upper() - - # 匹配方法（按 HTTP 方法） - matched = next((m for m in methods if m.method.upper() == http_method), None) - if not matched and methods: - matched = methods[0] # 降级到第一个 - if not matched: - query.reply_err( - payload=orjson.dumps({"error": f"No method for {subject}"}) - ) - return - - # 调用 RPC 方法（支持同步和异步；返回生成器时进入流式模式） - # user_info 从 meta 信封注入（框架参数，不进业务 payload） - result = await matched.call( - data, - session=session, - worker_id=worker_id, - user_info=meta.get("user_info"), - ) - - if inspect.isasyncgen(result) or inspect.isgenerator(result): - # 流式模式：@rpc_method 函数返回生成器，逐块多回复推送 - await stream_replies_to_query(query, result) - else: - # 一次性回复（原有逻辑） - query.reply( - key_expr=query.key_expr, - payload=orjson.dumps(result) - ) - logger.debug("[Worker] %s 回复 %s %s", worker_id, matched.method, subject) - except Exception as e: - if matched is not None and matched.stream: - # 流式端点错误：以 {"error": ...} 收尾块结束流（流式语义不变） - query.reply( - key_expr=query.key_expr, - payload=orjson.dumps({"error": str(e)}) - ) - else: - error_response = { - "result": 0.0, - "operation": "error", - "error": str(e), - "worker": worker_id - } - # 框架错误走 reply_err 通道：网关零解析判定失败（映射 503）； - # 载荷保持原错误信封结构，旧网关仍按普通回复透传（兼容） - query.reply_err( - 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) + async def _handle_query(query: zenoh.Query) -> None: + payload_bytes = bytes(query.payload) if query.payload else b"" + await bound.dispatch(payload_bytes, ZenohQuerySink(query)) def callback(query: zenoh.Query): # 当前 create_worker 运行在事件循环中，可直接 create_task @@ -968,10 +887,13 @@ async def create_worker( queryables: List[zenoh.Queryable] = [] for subject, methods in subject_methods.items(): key = f"worker/{worker_type}/{worker_id}/{subject}" + bound = dispatcher.bind(subject, methods, key) # 使用回调方式声明 - q = session.declare_queryable(key, make_callback(subject, methods)) + q = session.declare_queryable(key, make_callback(subject, bound)) queryables.append(q) logger.info("[Worker] %s Queryable: '%s' [%s]", worker_id, key, ", ".join(m.method for m in methods)) + # 业务 Queryable 就绪后登记直调入口（同进程调用方经此免 zenoh 往返） + register_dispatcher(dispatcher) # 声明 admin Queryable（管理指令 > restart / stop / 状态查询）。仅在有 controller 时注册， # 直接独立调用 create_worker（controller=None）保持原有行为、不暴露管理通道。 @@ -1004,9 +926,10 @@ async def create_worker( # 退出非由管理指令触发（正常返回 / 取消 / 异常）时，标记 stop —— 管理循环据此不自动重启 if controller is not None and not controller.exit_request: controller.exit_request = ADMIN_CMD_STOP - # 清理：先摘除 liveliness + 业务 Queryable，隔绝新请求 + # 清理：先摘除 liveliness + 业务 Queryable + 直调入口，隔绝新请求 if worker_manager is not None: worker_manager.unregister_worker(worker_type, worker_id) + unregister_dispatcher(worker_type, worker_id) token.undeclare() for q in queryables: q.undeclare()