295 lines
11 KiB
Python
295 lines
11 KiB
Python
"""Rust 进程插件端到端协议验证（多语言演进规划阶段 1 验收）
|
||
|
||
前置条件:
|
||
1. 本地 Zenoh Router 已运行（scripts/zenoh_router.py）
|
||
2. Rust 参考插件已编译（cargo build --examples 或 workers/rsdemo）
|
||
|
||
验证内容（docs/plugin-wire-protocol.md 全协议面）:
|
||
1. 路由描述符发布（keyexpr 编码、字段、schema_version/type）
|
||
2. liveliness 服务发现（worker/rsdemo/*）
|
||
3. 普通 RPC（GET 信封 / POST user_info 透传）
|
||
4. 非信封载荷兼容（method 默认 GET）
|
||
5. 流式多回复（逐块）+ 流中异常 {"error": ...} 收尾块
|
||
6. 非流式异常回复格式（result/operation/error/worker）
|
||
7. SIGTERM 优雅退出 → 描述符 DELETE 摘除
|
||
8. 重启 → 10s 内重发布自愈
|
||
|
||
用法（项目根目录）:
|
||
.venv/bin/python utils/worker-rs/scripts/verify_e2e.py [--type rsdemo] [--count 1] [--bin <路径>]
|
||
"""
|
||
import argparse
|
||
import orjson
|
||
import os
|
||
import signal
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
import zenoh
|
||
|
||
PASS = 0
|
||
FAIL = 0
|
||
|
||
DEFAULT_BIN = os.path.abspath(
|
||
os.path.join(os.path.dirname(__file__), "..", "target", "debug", "examples", "demo")
|
||
)
|
||
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
|
||
ROUTES_PREFIX = "chongming/routes"
|
||
|
||
|
||
def report(name, cond, detail=""):
|
||
global PASS, FAIL
|
||
if cond:
|
||
PASS += 1
|
||
print(f"✅ {name}" + (f" — {detail}" if detail else ""))
|
||
else:
|
||
FAIL += 1
|
||
print(f"❌ {name}" + (f" — {detail}" if detail else ""))
|
||
|
||
|
||
def start_worker(worker_type, count=1, bin_path=None):
|
||
"""启动 Rust 插件进程（CWD=项目根，读取全局 config.toml）"""
|
||
binary = bin_path or DEFAULT_BIN
|
||
proc = subprocess.Popen(
|
||
[binary, "--type", worker_type, "--count", str(count)],
|
||
cwd=ROOT,
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.STDOUT,
|
||
)
|
||
# 等待 Queryable 就绪（日志出现"进程就绪"）
|
||
deadline = time.monotonic() + 20
|
||
while time.monotonic() < deadline:
|
||
if proc.poll() is not None:
|
||
out = proc.stdout.read().decode() if proc.stdout else ""
|
||
raise RuntimeError(f"Worker 启动失败，退出码 {proc.returncode}:\n{out}")
|
||
time.sleep(0.2)
|
||
line = proc.stdout.readline() if proc.stdout else b""
|
||
if "等待退出信号" in line.decode("utf-8", "ignore"):
|
||
return proc
|
||
raise RuntimeError("Worker 20s 内未就绪")
|
||
|
||
|
||
def wait_worker_exit(proc, timeout=10):
|
||
try:
|
||
proc.wait(timeout=timeout)
|
||
return True
|
||
except subprocess.TimeoutExpired:
|
||
return False
|
||
|
||
|
||
def make_subscriber(session, worker_type):
|
||
"""建立常驻订阅，持续记录 PUT 描述符与 DELETE 事件（不重放历史值，
|
||
调用方须在事件发生前建立订阅）"""
|
||
samples = {}
|
||
|
||
def cb(sample):
|
||
key = str(sample.key_expr)
|
||
if sample.kind == zenoh.SampleKind.DELETE:
|
||
samples[key] = None # 摘除标记
|
||
else:
|
||
samples[key] = orjson.loads(bytes(sample.payload))
|
||
|
||
sub = session.declare_subscriber(f"{ROUTES_PREFIX}/{worker_type}/**", cb)
|
||
return sub, samples
|
||
|
||
|
||
def rpc(session, worker_expr, subject, payload, timeout=10.0):
|
||
"""按信封协议发起 Zenoh Query，返回所有回复块 JSON 列表"""
|
||
full_key = f"{worker_expr}/{subject}"
|
||
replies = session.get(
|
||
full_key,
|
||
payload=orjson.dumps(payload),
|
||
consolidation=zenoh.ConsolidationMode.NONE,
|
||
timeout=timeout,
|
||
)
|
||
result = []
|
||
for reply in replies:
|
||
if reply.ok:
|
||
result.append(orjson.loads(bytes(reply.ok.payload)))
|
||
else:
|
||
result.append({"__zenoh_err__": str(reply.err.payload)})
|
||
return result
|
||
|
||
|
||
def get_worker_exprs(session, worker_type):
|
||
"""liveliness 发现 worker/worker_type/* 副本列表"""
|
||
exprs = []
|
||
for reply in session.liveliness().get(f"worker/{worker_type}/*"):
|
||
if reply.ok:
|
||
exprs.append(str(reply.ok.key_expr))
|
||
return exprs
|
||
|
||
|
||
def main():
|
||
global PASS, FAIL
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--type", default="rsdemo")
|
||
parser.add_argument("--count", type=int, default=1)
|
||
parser.add_argument("--bin", default=None, help="插件二进制路径（默认 SDK 示例 demo）")
|
||
args = parser.parse_args()
|
||
wt = args.type
|
||
|
||
session = zenoh.open(zenoh.Config.from_file("config/zenoh_config.json5"))
|
||
print(f"=== Rust 进程插件协议验证 (type={wt}, bin={args.bin or DEFAULT_BIN}) ===")
|
||
|
||
# ========== 1. 路由描述符发布（先订阅再启动，zenoh 不重放历史值） ==========
|
||
sub1, samples1 = make_subscriber(session, wt)
|
||
proc = start_worker(wt, args.count, args.bin)
|
||
time.sleep(1.5)
|
||
|
||
expect_keys = {
|
||
"rsdemo": {
|
||
"chongming/routes/rsdemo/rsdemo_hello__GET",
|
||
"chongming/routes/rsdemo/rsdemo_echo__POST",
|
||
"chongming/routes/rsdemo/rsdemo_stream__POST",
|
||
"chongming/routes/rsdemo/rsdemo_error__GET",
|
||
},
|
||
"demo": {
|
||
"chongming/routes/demo/demo_hello__GET",
|
||
"chongming/routes/demo/demo_storage_upload_bytes__POST",
|
||
"chongming/routes/demo/demo_storage_download_bytes__GET",
|
||
"chongming/routes/demo/demo_storage_presign_upload__POST",
|
||
"chongming/routes/demo/demo_storage_presign_download__GET",
|
||
},
|
||
}[wt]
|
||
missing = expect_keys - set(samples1.keys())
|
||
report(
|
||
"路由描述符 keyexpr 编码正确",
|
||
not missing,
|
||
f"缺失: {missing or '无'}",
|
||
)
|
||
|
||
if "chongming/routes/rsdemo/rsdemo_hello__GET" in samples1:
|
||
d = samples1["chongming/routes/rsdemo/rsdemo_hello__GET"]
|
||
report("描述符 schema_version=1", d.get("schema_version") == 1, f"实际: {d.get('schema_version')}")
|
||
report("描述符 type=插件名", d.get("type") == wt, f"实际: {d.get('type')}")
|
||
report(
|
||
"描述符字段完整（subject/method/path/params/response_model）",
|
||
d.get("subject") == "rsdemo.hello"
|
||
and d.get("method") == "GET"
|
||
and d.get("path") == "/rsdemo/hello"
|
||
and d.get("params") == ["name: str"]
|
||
and isinstance(d.get("response_model"), dict),
|
||
)
|
||
report(
|
||
"defaults 合并（timeout=30.0/auth_required=false）",
|
||
d.get("timeout") == 30.0 and d.get("auth_required") is False,
|
||
)
|
||
sub1.undeclare()
|
||
|
||
# ========== 2. liveliness 服务发现 ==========
|
||
exprs = get_worker_exprs(session, wt)
|
||
report(
|
||
"liveliness 服务发现副本",
|
||
len(exprs) == args.count,
|
||
f"实际: {exprs}",
|
||
)
|
||
worker_expr = exprs[0] if exprs else f"worker/{wt}/missing"
|
||
|
||
# ========== 3. 普通 RPC：GET 信封 ==========
|
||
replies = rpc(session, worker_expr, "rsdemo.hello", {"meta": {"method": "GET"}, "data": {"name": "rust"}})
|
||
ok = (
|
||
len(replies) == 1
|
||
and replies[0].get("greeting") == "Hello, rust!"
|
||
and replies[0].get("worker_id") == worker_expr.split("/")[-1]
|
||
and isinstance(replies[0].get("timestamp"), (int, float))
|
||
)
|
||
report("RPC GET（信封分派 + worker_id 注入）", ok, f"实际: {replies}")
|
||
|
||
# ========== 4. 非信封载荷兼容（method 默认 GET） ==========
|
||
replies = rpc(session, worker_expr, "rsdemo.hello", {"name": "legacy"})
|
||
report(
|
||
"非信封载荷兼容（按 GET 处理）",
|
||
len(replies) == 1 and replies[0].get("greeting") == "Hello, legacy!",
|
||
f"实际: {replies}",
|
||
)
|
||
|
||
# ========== 5. POST + user_info 透传 ==========
|
||
replies = rpc(
|
||
session, worker_expr, "rsdemo.echo",
|
||
{"meta": {"method": "POST", "user_info": {"username": "alice"}}, "data": {"msg": "hi"}},
|
||
)
|
||
ok = (
|
||
len(replies) == 1
|
||
and replies[0].get("echo") == "hi"
|
||
and replies[0].get("user_info") == {"username": "alice"}
|
||
)
|
||
report("RPC POST（user_info 信封透传）", ok, f"实际: {replies}")
|
||
|
||
# ========== 6. 流式多回复 ==========
|
||
replies = rpc(
|
||
session, worker_expr, "rsdemo.stream",
|
||
{"meta": {"method": "POST"}, "data": {"prompt": "x", "tokens": 5}},
|
||
timeout=15.0,
|
||
)
|
||
report(
|
||
"流式多回复（5 块逐块推送）",
|
||
len(replies) == 5
|
||
and all(r.get("index") == i for i, r in enumerate(replies))
|
||
and all("error" not in r for r in replies),
|
||
f"块数: {len(replies)}",
|
||
)
|
||
|
||
# 流中异常 → {"error": ...} 收尾块（第 26 块抛错，共 26 块）
|
||
replies = rpc(
|
||
session, worker_expr, "rsdemo.stream",
|
||
{"meta": {"method": "POST"}, "data": {"prompt": "x", "tokens": 60}},
|
||
timeout=20.0,
|
||
)
|
||
ok = (
|
||
len(replies) == 26
|
||
and all("error" not in r for r in replies[:25])
|
||
and isinstance(replies[25].get("error"), str)
|
||
and "模拟流中异常" in replies[25]["error"]
|
||
)
|
||
report("流中异常 → error 收尾块终止", ok, f"块数: {len(replies)}, 末块: {replies[-1] if replies else None}")
|
||
|
||
# ========== 7. 非流式异常回复格式 ==========
|
||
replies = rpc(session, worker_expr, "rsdemo.error", {"meta": {"method": "GET"}, "data": {}})
|
||
ok = (
|
||
len(replies) == 1
|
||
and replies[0].get("result") == 0.0
|
||
and replies[0].get("operation") == "error"
|
||
and "模拟业务异常" in replies[0].get("error", "")
|
||
and replies[0].get("worker") == worker_expr.split("/")[-1]
|
||
)
|
||
report("非流式异常回复格式（result/operation/error/worker）", ok, f"实际: {replies}")
|
||
|
||
# ========== 8. SIGTERM 优雅退出 → 描述符摘除（先订阅再 kill） ==========
|
||
sub2, samples2 = make_subscriber(session, wt)
|
||
os.kill(proc.pid, signal.SIGTERM)
|
||
exited = wait_worker_exit(proc, timeout=10)
|
||
report("SIGTERM 优雅退出", exited, f"退出码: {proc.returncode if exited else '超时'}")
|
||
time.sleep(1.5)
|
||
removed = [k for k, v in samples2.items() if v is None]
|
||
report(
|
||
"退出时 DELETE 摘除全部描述符",
|
||
len(removed) == len(expect_keys),
|
||
f"摘除: {sorted(removed)}",
|
||
)
|
||
sub2.undeclare()
|
||
|
||
# ========== 9. 重启 → 10s 内重发布自愈 ==========
|
||
sub3, samples3 = make_subscriber(session, wt)
|
||
proc = start_worker(wt, args.count, args.bin)
|
||
time.sleep(11.5) # 覆盖一个完整重发布周期
|
||
ok = all(k in samples3 and samples3[k] is not None for k in expect_keys)
|
||
report(
|
||
"重启后周期重发布自愈（≤10s）",
|
||
ok,
|
||
f"恢复: {sorted(k for k in samples3 if samples3[k] is not None)}",
|
||
)
|
||
sub3.undeclare()
|
||
|
||
# 清理
|
||
os.kill(proc.pid, signal.SIGTERM)
|
||
wait_worker_exit(proc, timeout=10)
|
||
session.close()
|
||
|
||
print(f"\n=== 结果: {PASS} 通过, {FAIL} 失败 ===")
|
||
sys.exit(1 if FAIL else 0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|