zhifu-platform/scripts/recovery_exercise.py

import argparse
import asyncio
import json
import os
import subprocess
import sys
from contextlib import suppress
from datetime import UTC, datetime
from pathlib import Path
from time import monotonic
from uuid import UUID, uuid4

import mysql_recovery
import psutil
import redis_recovery
from sqlalchemy.ext.asyncio import async_sessionmaker

from zhifu.core.config import Settings
from zhifu.core.database import create_engine
from zhifu.modules.knowledge.config import KnowledgeSettings
from zhifu.modules.knowledge.recovery import rebuild, validate_target

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


class RecoveryTimeout(TimeoutError):
    def __init__(self, process_ids):
        super().__init__("恢复阶段达到最长执行时间，子进程树已终止")
        self.process_ids = process_ids


def output_path(value):
    path = Path(value).resolve()
    if not path.is_relative_to((ROOT / ".work").resolve()):
        raise ValueError("恢复编排报告必须保存在项目 .work 内")
    return path


async def knowledge(arguments):
    engine = create_engine(Settings())
    try:
        return await rebuild(
            async_sessionmaker(engine, expire_on_commit=False), arguments.target_collection,
            tenant_id=arguments.tenant_id.bytes if arguments.tenant_id else None,
            maximum_documents=arguments.maximum_documents,
            maximum_chunks=arguments.maximum_chunks,
            timeout_seconds=min(arguments.timeout_seconds, 600),
        )
    finally:
        await engine.dispose()


def phase(arguments):
    if arguments.phase == "mysql":
        directory = mysql_recovery.backup()
        report = mysql_recovery.restore(directory)
    elif arguments.phase == "redis":
        directory = redis_recovery.backup()
        report = redis_recovery.verify(directory)
    else:
        return asyncio.run(knowledge(arguments))
    return {"backup_directory": str(directory.relative_to(ROOT)), **report}


def containers(run_id):
    result = subprocess.run(
        ["docker", "ps", "-aq", "--filter", "label=zhifu.recovery-run=" + run_id],
        check=True, capture_output=True, text=True, timeout=15,
    )
    return result.stdout.split()


def terminate_tree(process):
    """暂停明确子进程树后终止，保证恢复阶段不能在容器清理后继续创建资源。"""
    root = psutil.Process(process.pid)
    suspended = {}
    # 暂停已发现进程后再次检查后代，最多扫描五轮覆盖阶段启动器和 Docker 子进程。
    for _ in range(5):
        candidates = [root, *root.children(recursive=True)]
        additions = [item for item in candidates if item.pid not in suspended]
        if not additions:
            break
        for item in additions:
            with suppress(psutil.NoSuchProcess):
                item.suspend()
                suspended[item.pid] = item
    for item in reversed(list(suspended.values())):
        with suppress(psutil.NoSuchProcess):
            item.kill()
    _, alive = psutil.wait_procs(list(suspended.values()), timeout=10)
    if alive:
        raise RuntimeError("恢复阶段子进程未在十秒内终止")
    process.wait(timeout=5)
    return sorted(suspended)


def run_phase(command, *, log_path, timeout, environment):
    with log_path.open("xb") as log:
        process = subprocess.Popen(
            command, env=environment, cwd=ROOT, stdout=log, stderr=subprocess.STDOUT,
        )
        try:
            returncode = process.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            identifiers = terminate_tree(process)
            raise RecoveryTimeout(identifiers) from None
        if returncode:
            raise subprocess.CalledProcessError(returncode, command)


def exercise(arguments):
    if not 60 <= arguments.timeout_seconds <= 1800:
        raise ValueError("恢复编排最长执行时间必须在 60–1800 秒之间")
    if (
        not 1 <= arguments.maximum_documents <= 100000
        or not 1 <= arguments.maximum_chunks <= 1000000
    ):
        raise ValueError("资料与分段数量上限无效")
    validate_target(KnowledgeSettings(), arguments.target_collection)
    if not 1 <= arguments.phase_timeout_seconds <= 1800:
        raise ValueError("单个恢复阶段最长执行时间必须在 1–1800 秒之间")
    output = output_path(arguments.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    run_id = uuid4().hex
    directory = output.parent / ("recovery-phases-" + run_id)
    directory.mkdir()
    report = {
        "run_id": run_id, "started_at": datetime.now(UTC).isoformat(), "status": "RUNNING",
        "maximum_seconds": arguments.timeout_seconds, "phases": {},
        "maximum_phase_seconds": arguments.phase_timeout_seconds,
        "generation_attempts": 0, "business_targets_switched": False,
        "consistent_cross_system_snapshot": False,
    }
    started = monotonic()
    environment = {**os.environ, "ZHIFU_RECOVERY_RUN_ID": run_id}
    try:
        for name in ("mysql", "redis", "knowledge"):
            remaining = arguments.timeout_seconds - (monotonic() - started)
            if remaining <= 0:
                raise TimeoutError("恢复编排达到最长执行时间")
            result = directory / (name + ".json")
            command = [
                sys.executable, str(Path(__file__).resolve()), "--phase", name,
                "--output", str(result), "--target-collection", arguments.target_collection,
                "--maximum-documents", str(arguments.maximum_documents),
                "--maximum-chunks", str(arguments.maximum_chunks),
                "--timeout-seconds", str(max(1, int(remaining))),
            ]
            if arguments.tenant_id:
                command.extend(["--tenant-id", str(arguments.tenant_id)])
            # 超时先终止阶段的整个进程树，再按随机运行标识清理独立恢复容器。
            run_phase(command, log_path=directory / (name + ".log"),
                      timeout=min(remaining, arguments.phase_timeout_seconds),
                      environment=environment)
            report["phases"][name] = json.loads(result.read_text(encoding="utf-8"))
        report["status"] = "VERIFIED"
    except RecoveryTimeout as error:
        report["terminated_process_ids"] = error.process_ids
        raise
    finally:
        remaining_containers = containers(run_id)
        report["temporary_containers_cleaned"] = remaining_containers
        if remaining_containers:
            subprocess.run(
                ["docker", "rm", "-fv", *remaining_containers], check=True,
                capture_output=True, timeout=30,
            )
        report["temporary_containers_remaining"] = containers(run_id)
        report["seconds"] = monotonic() - started
        report["finished_at"] = datetime.now(UTC).isoformat()
        if report["status"] == "RUNNING":
            report["status"] = "INCOMPLETE"
        output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    return report


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="顺序验证 MySQL、Redis 和独立知识集合恢复")
    parser.add_argument("--phase", choices=["mysql", "redis", "knowledge"])
    parser.add_argument("--output", default=".work/recovery-exercise.json")
    parser.add_argument("--target-collection", required=True)
    parser.add_argument("--tenant-id", type=UUID)
    parser.add_argument("--maximum-documents", type=int, default=1000)
    parser.add_argument("--maximum-chunks", type=int, default=10000)
    parser.add_argument("--timeout-seconds", type=int, default=600)
    parser.add_argument("--phase-timeout-seconds", type=int, default=600)
    arguments = parser.parse_args()
    if arguments.phase:
        output_path(arguments.output).write_text(
            json.dumps(phase(arguments), ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
        )
    else:
        print(json.dumps(exercise(arguments), ensure_ascii=False, indent=2))