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))