ykxmz/_build/repro_relay_hang.py

# -*- coding: utf-8 -*-
r'''精确复现「经中转发命令，前几次正常、之后卡死」。

== 为什么要单独写这个 ==

_tests/test_relay_local.py 会卡在「经中转发协议命令」那一步，
但它起的是完整的被控端进程，中间隔了太多层，看不出卡在哪。

这个脚本只用 HTTP 直接和中转对话，把每一步的队列状态都打出来：
    watch 发 -> agent 取 -> agent 回 -> client 取
每一步都断言，卡住的那一步会直接暴露。

**关键点：要模拟真实场景里「agent 的长轮询是开着的」** ——
真实被控端会一直挂着一个 poll 等命令。
如果只在"发完再 poll"，就复现不出竞态。
'''
from __future__ import annotations

import json
import os
import shutil
import subprocess
import sys
import tempfile
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
PORT = 8123
BASE = f"http://127.0.0.1:{PORT}"
UA = {"User-Agent": "LanRemote/1.8.2"}

passed = 0
failed = 0


def check(name, ok, detail=""):
    global passed, failed
    if ok:
        passed += 1
        print("  [PASS] " + name)
    else:
        failed += 1
        print("  [FAIL] " + name + (("  —— " + str(detail)) if detail else ""))


def post(path, body=b"", headers=None, timeout=20.0):
    h = dict(UA)
    h.update(headers or {})
    req = urllib.request.Request(BASE + path, data=body, headers=h,
                                method="POST")
    try:
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            return resp.status, resp.read()
    except urllib.error.HTTPError as exc:
        return exc.code, exc.read()
    except Exception as exc:                              # noqa: BLE001
        return 0, str(exc).encode()


def health():
    req = urllib.request.Request(BASE + "/v1/health", headers=UA)
    with urllib.request.urlopen(req, timeout=10) as resp:
        return json.loads(resp.read().decode("utf-8"))


def main() -> int:
    work = Path(tempfile.mkdtemp(prefix="repro_"))
    relay = None
    try:
        data = work / "d"
        data.mkdir(parents=True, exist_ok=True)
        env = dict(os.environ)
        env.update({"RELAY_PORT": str(PORT), "RELAY_BIND": "127.0.0.1",
                    "RELAY_DATA": str(data), "RELAY_POLL": "5",
                    "PYTHONIOENCODING": "utf-8", "PYTHONUNBUFFERED": "1"})
        logf = open(work / "relay.log", "wb")
        relay = subprocess.Popen(
            [sys.executable, "-u", str(ROOT / "relay" / "relay.py")], env=env,
            stdout=logf, stderr=subprocess.STDOUT)

        for _ in range(40):
            time.sleep(0.25)
            try:
                health()
                break
            except Exception:                             # noqa: BLE001
                continue
        check("中转起来了", True)

        # ---- 注册设备 + 接入手表 ----
        st, body = post("/v1/agent/register", json.dumps({
            "deviceId": "repro-pc", "code": "555555",
            "name": "复现机", "version": "1.8.2",
            "lanIps": [], "port": 0,
        }).encode(), {"Content-Type": "application/json"})
        reg = json.loads(body.decode())
        agent_token = reg.get("token", "")
        check("设备注册拿到 token", bool(agent_token), reg)

        st, body = post("/v1/client/attach", json.dumps({
            "code": "555555", "name": "复现手表"}).encode(),
            {"Content-Type": "application/json"})
        att = json.loads(body.decode())
        session = att.get("session", "")
        check("手表接入拿到 session", bool(session), att)

        # ---- 场景 1：同步往返（agent 不挂着 poll）----
        print()
        print("  场景 1：先发后取（没有并发长轮询）")
        for i in range(1, 4):
            payload = (json.dumps({"cmd": "ping", "n": i}) + "\n").encode()
            st, body = post("/v1/client/send", payload, {"X-Session": session})
            send_ok = st == 200
            # agent 取
            st2, got = post("/v1/agent/poll", b"", {"X-Token": agent_token},
                            timeout=15)
            received = payload.strip() in got
            check("第 %d 次 发->收" % i, send_ok and received,
                  "send=%s poll=%s got=%r" % (send_ok, st2, got[:80]))
            # agent 回
            reply = (json.dumps({"ok": True, "n": i}) + "\n").encode()
            post("/v1/agent/send", reply, {"X-Token": agent_token})
            st3, back = post("/v1/client/poll", b"", {"X-Session": session},
                             timeout=15)
            check("第 %d 次 回->收" % i, reply.strip() in back,
                  "poll=%s got=%r" % (st3, back[:80]))

        # ---- 场景 2：agent 挂着长轮询，watch 这时发 ----
        print()
        print("  场景 2：agent 挂长轮询的同时手表发命令（真实场景）")
        for i in range(1, 4):
            payload = (json.dumps({"cmd": "ping", "n": 100 + i})
                       + "\n").encode()
            result = {}

            def waiter():
                st, got = post("/v1/agent/poll", b"",
                               {"X-Token": agent_token}, timeout=20)
                result["status"] = st
                result["got"] = got

            t = threading.Thread(target=waiter, daemon=True)
            t.start()
            time.sleep(0.4)                       # 确保 poll 已经在挂
            post("/v1/client/send", payload, {"X-Session": session})
            t.join(timeout=12)
            alive = t.is_alive()
            got = result.get("got", b"")
            check("第 %d 次 长轮询期间发命令能收到" % i,
                  (not alive) and (payload.strip() in got),
                  "still_waiting=%s status=%s got=%r"
                  % (alive, result.get("status"), got[:80]))

            if alive:
                print("      ** 卡住了，看中转内部状态 **")
                print("      " + json.dumps(health(), ensure_ascii=False)[:300])
                break

        # ---- 场景 3：连续发多条，agent 一次取 ----
        print()
        print("  场景 3：连发 5 条，agent 连续取")
        for i in range(5):
            payload = (json.dumps({"cmd": "burst", "n": i}) + "\n").encode()
            post("/v1/client/send", payload, {"X-Session": session})
        got_all = b""
        for _ in range(6):
            st, got = post("/v1/agent/poll", b"", {"X-Token": agent_token},
                           timeout=8)
            got_all += got
            if got_all.count(b"\n") >= 5:
                break
        check("5 条都收到了", got_all.count(b"\n") >= 5,
              "收到 %d 条: %r" % (got_all.count(b"\n"), got_all[:200]))

        print()
        print("  ── 中转日志里的关键行 ──")
        logf.flush()
        try:
            for line in (work / "relay.log").read_text(
                    encoding="utf-8", errors="replace").splitlines()[-12:]:
                print("      " + line[:110])
        except Exception:                                 # noqa: BLE001
            pass

        print()
        print("=" * 56)
        print("  通过 %d 项，失败 %d 项" % (passed, failed))
        return 0 if failed == 0 else 1
    finally:
        if relay is not None:
            try:
                relay.terminate()
                relay.wait(timeout=6)
            except Exception:                             # noqa: BLE001
                try:
                    relay.kill()
                except Exception:                         # noqa: BLE001
                    pass
        shutil.rmtree(work, ignore_errors=True)


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