zhifu-platform/scripts/verify_channel_worker.py

import asyncio
import json
import os
import subprocess
import time
from pathlib import Path
from uuid import uuid4

from redis.asyncio import Redis
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import async_sessionmaker

from zhifu.core.config import Settings
from zhifu.core.database import create_engine
from zhifu.modules.merchant.models import ChannelApplication
from zhifu.modules.platform.models import ChannelDelivery

ROOT = Path(__file__).resolve().parents[1]
MAX_SECONDS = 25


def docker(*args):
    return subprocess.run(
        ["docker", *args], check=True, capture_output=True, text=True, timeout=15,
    ).stdout.strip()


async def inspect_environment():
    settings = Settings()
    engine = create_engine(settings)
    redis = Redis.from_url(settings.state_url.get_secret_value(), socket_timeout=3)
    try:
        sessions = async_sessionmaker(engine)
        async with sessions() as session:
            enabled = await session.scalar(select(func.count()).select_from(
                ChannelApplication,
            ).where(ChannelApplication.enabled.is_(True)))
            pending = await session.scalar(select(func.count()).select_from(
                ChannelDelivery,
            ).where(ChannelDelivery.status.in_(["QUEUED", "SENDING"])))
        keys = {key async for key in redis.scan_iter(match="monitor:worker:channel-delivery:*")}
        return enabled, pending, keys
    finally:
        await redis.aclose()
        await engine.dispose()


def main():
    enabled, pending, initial = asyncio.run(inspect_environment())
    if enabled or pending:
        raise RuntimeError("生命周期验收要求当前没有启用渠道和等待投递记录")
    name = "zhifu-channel-lifecycle-" + uuid4().hex
    environment = Path(os.environ.get(
        "ZHIFU_RUNTIME_ENV", str(ROOT / "ops/runtime.local.env.example"),
    )).resolve(strict=True)
    started = time.monotonic()
    created = False
    try:
        # 只运行渠道进程，源码只读挂载；检查前提确保本次不发送外部业务通知。
        docker(
            "run", "--detach", "--name", name,
            "--network", os.environ.get("ZHIFU_DEPENDENCY_NETWORK", "zhifu-platform_default"),
            "--env-file", str(environment), "--env", "PYTHONPATH=/workspace/src",
            "--mount", f"type=bind,source={ROOT / 'src'},target=/workspace/src,readonly",
            "--read-only", "--cap-drop", "ALL", "--security-opt", "no-new-privileges:true",
            os.environ.get("ZHIFU_IMAGE", "zhifu-platform:local"),
            "python", "-m", "zhifu.operations.worker", "channel-delivery",
        )
        created = True
        while time.monotonic() - started < MAX_SECONDS:
            _, _, current = asyncio.run(inspect_environment())
            if current - initial:
                break
            time.sleep(0.25)
        else:
            raise RuntimeError("渠道工作进程未在 25 秒期限内登记心跳")
        assert len(current - initial) == 1, "渠道工作进程应只登记一份心跳"
        docker("stop", "--time", "10", name)
        exit_code = int(docker("inspect", "--format", "{{.State.ExitCode}}", name))
        assert exit_code == 0, f"渠道进程停止状态异常：{exit_code}"
        _, _, after = asyncio.run(inspect_environment())
        assert after == initial, "正常停止后应立即清除本进程心跳"
        assert "worker=channel-delivery stopped" in docker("logs", name)
        report = {
            "role": "channel-delivery", "heartbeat_count": 1, "signal": "SIGTERM",
            "exit_code": exit_code, "heartbeat_removed": True,
            "seconds": round(time.monotonic() - started, 2),
            "generated_model_requests": 0, "external_delivery_requests": 0,
        }
        target = ROOT / ".work/channel-worker-lifecycle.json"
        target.parent.mkdir(exist_ok=True)
        target.write_text(json.dumps(report, indent=2), encoding="utf-8")
        print(json.dumps(report))
    finally:
        if created:
            docker("rm", "--force", name)


if __name__ == "__main__":
    main()