chongming-core/test/test_inline_leader.py

"""
inline worker 领导选举 + 自愈接管测试（gunicorn 多 worker 兼容）

基于共享 chongming-cache（DistributedLock）验证 WorkerManager 的领导选举语义：
  1. 多"进程"各自 acquire_inline_owner：同一时刻只有一个持有者（无分裂）
  2. 不同 gateway_id（多实例）key 隔离，各选各的 leader，互不抢占
  3. 持有者在场且副本存活时不触发接管（不重复 init）
  4. 持有者崩溃（副本消失）后，非持有者经 _maybe_heal_inline 抢锁接管并调用注册的
     starter 拉起 worker
  5. 接管后其它进程不再重复接管；再次崩溃可继续 failover
  6. 关停释放锁；保护性 release 不误删他人锁

运行: uv run python test/test_inline_leader.py
"""
import asyncio
import os
import sys
import tempfile
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_cache import Cache
from chongming.app.core.worker_manager import WorkerManager

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 DummySession:
    """占位 session：测试不涉及真实 Zenoh，_maybe_heal_inline 不触碰 session"""


async def try_heal(loser):
    """非持有者进程尝试接管"""
    await loser._maybe_heal_inline()


async def main():
    tmp = tempfile.TemporaryDirectory()
    c1 = Cache(backend="diskcache", cache_store_path=tmp.name)
    c2 = Cache(backend="diskcache", cache_store_path=tmp.name + "_g2")

    # ── 1. 领导选举：唯一持有者 ──────────────────────────────
    m1 = WorkerManager(session=DummySession(), cache=c1)
    m2 = WorkerManager(session=DummySession(), cache=c1)
    a1 = m1.acquire_inline_owner()
    a2 = m2.acquire_inline_owner()
    check("1.1 同一时刻仅一进程获得持有权", a1 != a2, f"m1={a1}, m2={a2}")
    owner, loser = (m1, m2) if a1 else (m2, m1)
    check("1.2 持有者 is_inline_owner()=True", owner.is_inline_owner())
    check("1.3 失权者 is_inline_owner()=False", not loser.is_inline_owner())

    # ── 2. 释放后可再次竞选 ──────────────────────────────────
    await owner.release_inline_owner()
    check("2.1 持锁者释放后不再持有", not owner.is_inline_owner())
    ra = m1.acquire_inline_owner()
    rb = m2.acquire_inline_owner()
    check("2.2 释放后可再次成功选举（唯一持有者）", ra != rb, f"m1={ra}, m2={rb}")
    await m1.release_inline_owner()
    await m2.release_inline_owner()

    # ── 3. 不同 gateway_id 隔离（多实例各选各 leader） ────────
    # gateway_id 默认取 hostname；同宿主多实例须显式设置 CHONGMING_GATEWAY_ID 区分，
    # 这儿模拟两个不同实例（gwA / gwB），验证锁 key 隔离、互不抢占。
    _orig_gw = os.environ.get("CHONGMING_GATEWAY_ID")
    os.environ["CHONGMING_GATEWAY_ID"] = "gwA"
    ga1 = WorkerManager(session=DummySession(), cache=c2)
    ga2 = WorkerManager(session=DummySession(), cache=c2)
    os.environ["CHONGMING_GATEWAY_ID"] = "gwB"
    gb1 = WorkerManager(session=DummySession(), cache=c2)
    if _orig_gw is None:
        os.environ.pop("CHONGMING_GATEWAY_ID", None)
    else:
        os.environ["CHONGMING_GATEWAY_ID"] = _orig_gw

    acc_a = ga1.acquire_inline_owner() or ga2.acquire_inline_owner()
    acc_b = gb1.acquire_inline_owner()
    check("3.1 不同 gateway 互不抢占（各自选到 leader）", acc_a and acc_b)
    check("3.2 gateway 内部仍只有一个持有者",
          ga1.is_inline_owner() != ga2.is_inline_owner())
    await ga1.release_inline_owner()
    await ga2.release_inline_owner()
    await gb1.release_inline_owner()

    # ── 4. 自愈接管 ──────────────────────────────────────────
    ha = WorkerManager(session=DummySession(), cache=c1)
    hb = WorkerManager(session=DummySession(), cache=c1)
    got = hb.acquire_inline_owner()
    check("4.1 初始 hb 持锁、ha 未持锁", got and not ha.is_inline_owner())

    started = []

    async def fake_starter():
        started.append("doubao")
        return []

    # 与真实 load_workers 一致：所有进程都注册 starter
    ha.register_worker_starter("doubao", fake_starter)
    hb.register_worker_starter("doubao", fake_starter)

    # 4.2 持有者在场、副本存活 → 不接管
    hb._local_workers["doubao"] = ["worker/doubao/worker-0"]
    ha._workers_cache["doubao"] = ["worker/doubao/worker-0"]
    await try_heal(ha)
    check("4.2 持有者在场且副本存活：不接管、不拉起", not ha.is_inline_owner())

    # 4.3 持有者崩溃（模拟：锁租约过期 + 副本消失）→ ha 接管并拉起
    await hb.release_inline_owner()          # 模拟原持有者崩溃后锁租约过期
    ha._workers_cache.clear()
    hb._local_workers.clear()
    await try_heal(ha)
    check("4.3 副本消失后非持有者接管", ha.is_inline_owner())
    check("4.4 接管时调用 starter 拉起 worker", "doubao" in started)

    # 4.5 接管后其它进程不再重复接管/拉起
    n_started = len(started)
    # 模拟场景：新副本已能被"所有进程"看到
    hb._workers_cache["doubao"] = ["worker/doubao/worker-0"]
    await try_heal(hb)
    check("4.5 接管后其它进程不再重复拉起",
          len(started) == n_started or hb.is_inline_owner())

    # ── 5. 再次崩溃可继续 failover ───────────────────────────
    await ha.release_inline_owner()          # 模拟 ha 崩溃后锁租约过期
    ha._workers_cache.clear()
    hb._workers_cache.clear()
    n_started = len(started)
    await try_heal(hb)
    check("5.1 再次崩溃后 hb 接管", hb.is_inline_owner())
    check("5.2 failover 期间新增一次拉起", len(started) > n_started)

    # ── 6. 保护性释放：失权者 release 不误删他人锁 ───────────
    # 当前持有者为 hb，由 ha(非持有者) 释放——锁 token 不匹配，不应删除
    await ha.release_inline_owner()
    check("6.1 非持锁者 release 不误删（hb 仍持有）", hb.is_inline_owner())
    # 之后真正由持锁者安装
    await hb.release_inline_owner()
    check("6.2 真正持锁者释放后不再持有", not hb.is_inline_owner())

    for c in (c1, c2):
        c.close()

    print()
    if failures:
        print(f"FAILED {len(failures)}: {failures}")
        return 1
    print("ALL PASS")
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))