221 lines
9.8 KiB
Python
221 lines
9.8 KiB
Python
"""
|
||
worker 实例定向（把请求固定到某个进程实例）测试
|
||
|
||
验证 WorkerManager 的实例定向语义（HTTP 的 X-Chongming-Instance 头 / MQTT 的
|
||
target_instance 字段 / 程序调用 instance_id 参数最终都走这里）：
|
||
1. 不指定实例：与以往一致，在全部候选副本里随机挑
|
||
2. 指定实例：只在其实例的在线副本里挑，不跨实例
|
||
3. 身份来源：副本经 admin status 上报 instance_id，网关按需查询并缓存
|
||
（本地注册的 inline 副本直接用网关自身实例，无需查询）
|
||
4. 目标实例不在线：报错而不回退到别的实例（定向即严格）
|
||
5. 调用失败重试：限定在同一实例内换副本，绝不跨实例
|
||
6. 实例内已无其它副本可换：立即失败，不重复打同一个副本
|
||
7. 实例维度缓存键：定向调用的缓存与其它实例隔离
|
||
|
||
运行: .venv/bin/python test/test_instance_targeting.py
|
||
"""
|
||
import asyncio
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
_HERE = Path(__file__).resolve().parent.parent # chongming-zenoh/
|
||
_PROJ = _HERE
|
||
for _p in (_PROJ / "gateway" / "src", _PROJ / "utils" / "worker" / "src"):
|
||
sys.path.insert(0, str(_p))
|
||
|
||
from chongming.app.core.worker_manager import WorkerManager # noqa: E402
|
||
|
||
failures = []
|
||
|
||
|
||
def check(name, cond, detail=""):
|
||
status = "PASS" if cond else "FAIL"
|
||
print(f"[{status}] {name}" + (f" — {detail}" if detail else ""))
|
||
if not cond:
|
||
failures.append(name)
|
||
|
||
|
||
class _OkPayload:
|
||
def __init__(self, data: bytes):
|
||
self.payload = data
|
||
|
||
|
||
class _Ok:
|
||
def __init__(self, data: bytes):
|
||
self.ok = _OkPayload(data)
|
||
|
||
|
||
class _Reply:
|
||
"""一次 get 的回复：ok 与 err 互斥（与 Zenoh reply 的形态一致）。"""
|
||
|
||
def __init__(self, data: bytes = b"{}", err: bytes | None = None):
|
||
if err is None:
|
||
self.ok = _OkPayload(data)
|
||
self.err = None
|
||
else:
|
||
self.ok = None
|
||
self.err = _OkPayload(err)
|
||
|
||
|
||
class FakeSession:
|
||
"""记录每次 get 的 key，并按 key 决定回复。
|
||
|
||
admin 查询（worker/<type>/<id>/__admin__）由 admin_instances 决定 instance_id；
|
||
业务调用（worker/<type>/<id>/<subject>）可按 call_errors 注入错误回复。
|
||
"""
|
||
|
||
def __init__(self, admin_instances: dict | None = None,
|
||
call_errors: set | None = None):
|
||
self.admin_instances = admin_instances or {}
|
||
self.call_errors = set(call_errors or ())
|
||
self.calls: list = []
|
||
|
||
def get(self, key, payload=None, **kwargs):
|
||
self.calls.append(key)
|
||
if key.endswith("/__admin__"):
|
||
wid = key.split("/")[2]
|
||
iid = self.admin_instances.get(wid)
|
||
body = (
|
||
f'{{"worker_id":"{wid}","instance_id":"{iid}","status":"running"}}'.encode()
|
||
if iid else f'{{"worker_id":"{wid}","status":"running"}}'.encode()
|
||
)
|
||
return iter([_Reply(body)])
|
||
if key in self.call_errors:
|
||
return iter([_Reply(err=b'{"error":"boom"}')])
|
||
return iter([_Reply(b'{"ok":true}')])
|
||
|
||
def liveliness(self):
|
||
raise RuntimeError("本测试不触及 liveliness")
|
||
|
||
def close(self):
|
||
pass
|
||
|
||
|
||
def new_manager(session, gateway_id="gw-local"):
|
||
m = WorkerManager(session=session)
|
||
m._gateway_id = gateway_id
|
||
return m
|
||
|
||
|
||
async def main():
|
||
# ── 1. 不定向：与以往一致，全候选集随机 ─────────────────
|
||
sess = FakeSession(admin_instances={"w-a": "edge-01", "w-b": "edge-02"})
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = ["worker/demo/w-a", "worker/demo/w-b"]
|
||
picked = {await m.get_worker("demo") for _ in range(60)}
|
||
check("1.1 不定向时可命中全部副本", picked == set(m._workers_cache["demo"]),
|
||
f"picked={sorted(picked)}")
|
||
check("1.2 不定向不触发实例身份查询（零额外 RPC）",
|
||
not [c for c in sess.calls if c.endswith("__admin__")], f"calls={sess.calls}")
|
||
check("1.3 不定向调用走随机副本",
|
||
await m.call_raw("demo", "ping", b"{}") == b'{"ok":true}')
|
||
|
||
# ── 2. 定向：只在目标实例内挑副本 ────────────────────────
|
||
sess = FakeSession(admin_instances={"w-a": "edge-01", "w-b": "edge-02",
|
||
"w-c": "edge-01"})
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = [
|
||
"worker/demo/w-a", "worker/demo/w-b", "worker/demo/w-c",
|
||
]
|
||
picked = {await m.get_worker("demo", instance_id="edge-01") for _ in range(60)}
|
||
check("2.1 定向只命中目标实例的副本",
|
||
picked == {"worker/demo/w-a", "worker/demo/w-c"}, f"picked={sorted(picked)}")
|
||
|
||
# ── 3. 身份查询：按需查询并缓存，实例归属正确 ─────────────
|
||
sess = FakeSession(admin_instances={"w-a": "edge-01", "w-b": "edge-02",
|
||
"w-c": "edge-01"})
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = [
|
||
"worker/demo/w-a", "worker/demo/w-b", "worker/demo/w-c",
|
||
]
|
||
index = await m._instance_workers("demo", "edge-01")
|
||
check("3.1 实例→副本反向索引正确",
|
||
sorted(index) == ["worker/demo/w-a", "worker/demo/w-c"], f"index={sorted(index)}")
|
||
admin_calls = [c for c in sess.calls if c.endswith("__admin__")]
|
||
check("3.2 每个副本只查询一次身份", len(admin_calls) == 3, f"n={len(admin_calls)}")
|
||
await m._refresh_instance_map("demo")
|
||
check("3.3 重复刷新复用缓存不重复查询",
|
||
len([c for c in sess.calls if c.endswith("__admin__")]) == len(admin_calls))
|
||
|
||
# ── 3b. inline 本地注册副本：直接用网关实例，零查询 ──────
|
||
sess = FakeSession()
|
||
m = new_manager(sess, gateway_id="gw-1")
|
||
m.register_worker("demo", "worker-0-1-999")
|
||
index = await m._instance_workers("demo", "gw-1")
|
||
check("3.4 inline 副本归入网关实例", index == ["worker/demo/worker-0-1-999"],
|
||
f"index={index}")
|
||
check("3.5 inline 副本身份查询零 RPC",
|
||
not [c for c in sess.calls if c.endswith("__admin__")], f"calls={sess.calls}")
|
||
|
||
# ── 4. 定向实例不在线：严格报错，不回退 ──────────────────
|
||
sess = FakeSession(admin_instances={"w-a": "edge-01"})
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = ["worker/demo/w-a"]
|
||
check("4.1 目标实例无副本时 get_worker 返回 None",
|
||
await m.get_worker("demo", instance_id="edge-99") is None)
|
||
try:
|
||
await m.call_raw("demo", "ping", b"{}", instance_id="edge-99")
|
||
check("4.2 定向到离线实例应报错", False, "未抛错")
|
||
except RuntimeError as e:
|
||
msg = str(e)
|
||
check("4.2 定向到离线实例报错且说明不跨实例回退",
|
||
"edge-99" in msg and "不" in msg, f"msg={msg[:90]}")
|
||
|
||
# ── 5. 重试锚定同实例：绝不跨实例 ────────────────────────
|
||
sess = FakeSession(
|
||
admin_instances={"w-a": "edge-01", "w-b": "edge-02", "w-c": "edge-01"},
|
||
call_errors={"worker/demo/w-b/ping"}, # 只让别实例的副本出错，验证不会被选中
|
||
)
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = [
|
||
"worker/demo/w-a", "worker/demo/w-b", "worker/demo/w-c",
|
||
]
|
||
# 目标实例 edge-01 内的 w-a 失败 -> 应换到同实例的 w-c，不碰 edge-02 的 w-b
|
||
sess.call_errors = {"worker/demo/w-a/ping"}
|
||
out = await m.call_raw("demo", "ping", b"{}", instance_id="edge-01", retries=2)
|
||
business = [c for c in sess.calls if not c.endswith("__admin__")]
|
||
check("5.1 重试不跨实例（未打其它实例的副本）",
|
||
all("/w-b/" not in c for c in business), f"calls={business}")
|
||
check("5.2 同实例内换副本后成功", out == b'{"ok":true}', f"out={out!r}")
|
||
|
||
# ── 6. 实例内无其它副本可换：立即失败 ────────────────────
|
||
sess = FakeSession(admin_instances={"w-a": "edge-01"},
|
||
call_errors={"worker/demo/w-a/ping"})
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = ["worker/demo/w-a", "worker/demo/w-b"]
|
||
sess.admin_instances["w-b"] = "edge-02"
|
||
raised = None
|
||
try:
|
||
await m.call_raw("demo", "ping", b"{}", instance_id="edge-01", retries=2)
|
||
except RuntimeError as e:
|
||
raised = e
|
||
check("6.1 实例内无副本可换应报错", raised is not None, f"raised={raised!r}")
|
||
n = len([c for c in sess.calls if not c.endswith("__admin__")])
|
||
check("6.2 不重复打同一个副本（只调用一次）", n == 1, f"业务调用次数={n}")
|
||
check("6.3 错误信息点名实例与原因",
|
||
raised is not None and "edge-01" in str(raised), f"msg={str(raised)[:90]}")
|
||
|
||
# ── 7. 定向调用的缓存按实例隔离 ──────────────────────────
|
||
sess = FakeSession(admin_instances={"w-a": "edge-01", "w-b": "edge-02"})
|
||
m = new_manager(sess)
|
||
m._workers_cache["demo"] = ["worker/demo/w-a", "worker/demo/w-b"]
|
||
r1 = await m.call_raw("demo", "state", b"{}", cacheable=True, instance_id="edge-01")
|
||
r2 = await m.call_raw("demo", "state", b"{}", cacheable=True, instance_id="edge-02")
|
||
keys = [f"{iid}@demo:state:" for iid in ("edge-01", "edge-02")]
|
||
check("7.1 不同实例的定向缓存键互不影响",
|
||
all(any(k.startswith(p) for k in m._rpc_mem_cache._d) for p in keys),
|
||
f"cache_keys={list(m._rpc_mem_cache._d)}")
|
||
check("7.2 定向调用结果正常返回", r1 == b'{"ok":true}' and r2 == b'{"ok":true}',
|
||
f"r1={r1!r}, r2={r2!r}")
|
||
|
||
print()
|
||
if failures:
|
||
print(f"FAILED: {len(failures)} 项 -> {failures}")
|
||
return 1
|
||
print("全部通过")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(asyncio.run(main()))
|