187 lines
6.5 KiB
Python
187 lines
6.5 KiB
Python
"""
|
||
外部进程多实例全局分摊 —— distribute_quota 逻辑单测（无需网关/缓存后端）。
|
||
|
||
用 dict 作缓存桩，模拟同类型多个外部进程实例注册，验证扩缩容时全局 count 按
|
||
「存活实例数 N」分摊（而非每个进程都全额），以及实例退出/租约过期后的接管行为。
|
||
"""
|
||
import os
|
||
import sys
|
||
import time
|
||
|
||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../utils/worker/src"))
|
||
|
||
import chongming_worker.runtime as r
|
||
|
||
PASS = 0
|
||
FAIL = 0
|
||
|
||
|
||
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 fake_cache():
|
||
return {}
|
||
|
||
|
||
def as_instance(iid, cache, worker_type="demo"):
|
||
"""把当前身份设为 iid，并在缓存登记该实例参与方。"""
|
||
old = os.environ.get("CHONGMING_WORKER_INSTANCE")
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = iid
|
||
try:
|
||
r.register_instance(worker_type, cache)
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
|
||
else:
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = old
|
||
|
||
|
||
def quota_as(iid, worker_type, total, cache):
|
||
"""以指定实例身份计算其分摊配额（external 分支）。"""
|
||
old = os.environ.get("CHONGMING_WORKER_INSTANCE")
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = iid
|
||
try:
|
||
return r.distribute_quota(cache, worker_type, total, gw_id=None)
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
|
||
else:
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = old
|
||
|
||
|
||
def unreg_instance(iid, cache, worker_type="demo"):
|
||
"""以指定实例身份从参与方列表移除它。"""
|
||
old = os.environ.get("CHONGMING_WORKER_INSTANCE")
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = iid
|
||
try:
|
||
r.unregister_instance(worker_type, cache)
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
|
||
else:
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = old
|
||
|
||
|
||
def run():
|
||
cache = fake_cache()
|
||
wt = "demo"
|
||
|
||
# 1. 单实例（未启用实例注册，参与方为空）→ 兜底全额（兼容旧行为）
|
||
q = r.distribute_quota(cache, wt, 6, gw_id=None)
|
||
report("单实例无参与方 → 全额承担 6", q == 6, f"got={q}")
|
||
|
||
# 2. 两个实例并存：全局 6 分摊 → 各 3
|
||
as_instance("inst-a", cache, wt)
|
||
as_instance("inst-b", cache, wt)
|
||
qa = quota_as("inst-a", wt, 6, cache)
|
||
qb = quota_as("inst-b", wt, 6, cache)
|
||
report("两个实例并存 → 各分 3", qa == 3 and qb == 3, f"a={qa}, b={qb}")
|
||
|
||
# 3. 三个实例、总数不能整除：7 // 3 余 1 → 按序给第一个补 1
|
||
as_instance("inst-c", cache, wt)
|
||
qs = {i: quota_as(i, wt, 7, cache) for i in ("inst-a", "inst-b", "inst-c")}
|
||
report(
|
||
"三个实例 ≠整除 → 总和=7 且余量补靠前实例",
|
||
sum(qs.values()) == 7 and set(qs.values()) == {2, 2, 3},
|
||
f"{qs}",
|
||
)
|
||
|
||
# 4. 一个实例退出（unregister）→ 其余接管全部
|
||
unreg_instance("inst-c", cache, wt)
|
||
qa2 = quota_as("inst-a", wt, 6, cache)
|
||
qb2 = quota_as("inst-b", wt, 6, cache)
|
||
report(
|
||
"实例退出后由剩余接管 → 各分 3",
|
||
qa2 == 3 and qb2 == 3,
|
||
f"a={qa2}, b={qb2}",
|
||
)
|
||
|
||
# 5. 租约过期清理：伪造一个过期的参与方，读实例列表应排除它
|
||
cache[r._SCALE_INSTANCES_KEY.format(type=wt)] = [
|
||
["inst-stale", time.time() - r._INSTANCE_LEASE - 10],
|
||
["inst-fresh", time.time()],
|
||
]
|
||
ids = r._read_instances(cache, wt)
|
||
report(
|
||
"死实例按租约剔除",
|
||
ids == ["inst-fresh"],
|
||
f"instances={ids}",
|
||
)
|
||
|
||
# 6. 实例数=1 时全额（不误分摊）
|
||
cache[r._SCALE_INSTANCES_KEY.format(type=wt)] = [
|
||
[r.current_instance_id(), time.time()]
|
||
]
|
||
q1 = r.distribute_quota(cache, wt, 5, gw_id=None)
|
||
report("单存活实例 → 全额 5", q1 == 5, f"got={q1}")
|
||
|
||
# 7. inline（有网关身份）路径回归：按在线网关分摊不受影响
|
||
cache[r._SCALE_GATEWAYS_KEY] = [["gw-a", time.time()], ["gw-b", time.time()]]
|
||
qa_in = r.distribute_quota(cache, wt, 6, "gw-a")
|
||
qb_in = r.distribute_quota(cache, wt, 6, "gw-b")
|
||
report(
|
||
"inline 按网关分摊回归 → 各 3",
|
||
qa_in == 3 and qb_in == 3,
|
||
f"gw-a={qa_in}, gw-b={qb_in}",
|
||
)
|
||
|
||
# 8. 逐实例绝对覆盖：写入后能被读回且供启动解析（external）使用
|
||
r._write_own_instance_count(cache, wt, 4, "inst-a")
|
||
r._write_own_instance_count(cache, wt, 2, "inst-b")
|
||
report(
|
||
"instcounts 覆盖写入/读回",
|
||
r._own_instance_count(cache, wt, "inst-a") == 4
|
||
and r._own_instance_count(cache, wt, "inst-b") == 2,
|
||
f"{r._read_instance_counts(cache, wt)}",
|
||
)
|
||
# 启动解析：external（gw_id=None）有绝对覆盖时直接采用，而非被全局均分
|
||
old = os.environ.get("CHONGMING_WORKER_INSTANCE")
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = "inst-a"
|
||
try:
|
||
restored = r.resolve_worker_count(wt, 6, cache, "cache", gw_id=None)
|
||
finally:
|
||
if old is None:
|
||
os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
|
||
else:
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = old
|
||
report("resolve_worker_count 启动时采用本实例绝对数", restored == 4, f"got={restored}")
|
||
|
||
# 9. 清除覆盖后回退全局均分
|
||
r._clear_instance_counts(cache, wt)
|
||
report(
|
||
"清除覆盖后覆盖表为空",
|
||
r._read_instance_counts(cache, wt) == {},
|
||
f"{r._read_instance_counts(cache, wt)}",
|
||
)
|
||
|
||
# 10. 缺省实例身份唯一（hostname:pid）：同机未指定 --instance 的多进程不会并成一个参与方
|
||
r.set_current_instance_id(None)
|
||
old_env = os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
|
||
try:
|
||
default_iid = r.current_instance_id()
|
||
finally:
|
||
if old_env is not None:
|
||
os.environ["CHONGMING_WORKER_INSTANCE"] = old_env
|
||
else:
|
||
os.environ.pop("CHONGMING_WORKER_INSTANCE", None)
|
||
import socket as _socket
|
||
report(
|
||
"缺省实例身份含进程区分（{host}:{pid}）",
|
||
default_iid == f"{_socket.gethostname()}:{os.getpid()}",
|
||
f"got={default_iid}",
|
||
)
|
||
|
||
print("\n" + "=" * 56)
|
||
print(f"测试结果: ✅ {PASS} 通过 ❌ {FAIL} 失败")
|
||
return 0 if FAIL == 0 else 1
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(run()) |