training/skills/fullstack-start-verify/scripts/start_and_verify.py

#!/usr/bin/env python3
"""启动两个本地服务，等待 HTTP 就绪后执行检查并清理进程。"""

from __future__ import annotations

import argparse
import concurrent.futures
import os
import re
import shutil
import signal
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request
import uuid
from dataclasses import dataclass
from pathlib import Path


@dataclass
class Service:
    name: str
    command: str
    cwd: Path
    url: str
    process: subprocess.Popen[bytes] | None = None
    started_by_script: bool = False
    reused: bool = False
    log_path: Path | None = None
    log_file: object | None = None


SENSITIVE_VALUE_PATTERN = re.compile(
    r"(?i)(password|token|secret|authorization|api[_-]?key)"
    r"(\s*[\"']?\s*[:=]\s*[\"']?)([^\"'\s,;}\]]+)"
)
BEARER_PATTERN = re.compile(r"(?i)(Bearer\s+)[^\s,;]+")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="启动前后端命令，检查 HTTP 地址并清理进程。"
    )
    parser.add_argument("--backend-cmd", required=True, help="后端启动命令")
    parser.add_argument("--backend-dir", default=".", help="后端工作目录")
    parser.add_argument("--backend-url", required=True, help="后端健康检查地址")
    parser.add_argument("--frontend-cmd", required=True, help="前端启动命令")
    parser.add_argument("--frontend-dir", default=".", help="前端工作目录")
    parser.add_argument("--frontend-url", required=True, help="前端就绪检查地址")
    parser.add_argument(
        "--check-url",
        action="append",
        default=[],
        metavar="LABEL=URL",
        help="额外的 GET 检查，可重复传入",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=60.0,
        help="每个服务的就绪超时时间，单位为秒，默认 60",
    )
    parser.add_argument(
        "--interval",
        type=float,
        default=0.25,
        help="每次就绪检查的间隔秒数，默认 0.25",
    )
    parser.add_argument(
        "--check-timeout",
        type=float,
        default=5.0,
        help="额外接口检查的单请求超时时间，默认 5 秒",
    )
    parser.add_argument(
        "--reuse-running",
        action="store_true",
        help="健康地址已属于当前项目时复用已运行服务，不启动或清理该服务",
    )
    parser.add_argument(
        "--log-dir",
        type=Path,
        help="日志目录，运行结束后保留",
    )
    parser.add_argument(
        "--keep-logs",
        action="store_true",
        help="运行结束后保留临时日志",
    )
    parser.add_argument(
        "--keep-running",
        action="store_true",
        help="验证结束后保持已启动的服务运行",
    )
    return parser.parse_args()


def resolve_dir(value: str) -> Path:
    path = Path(value).resolve()
    if not path.is_dir():
        raise RuntimeError(f"工作目录不存在：{path}")
    return path


def start_service(service: Service, log_dir: Path) -> None:
    service.log_path = log_dir / f"{service.name}.log"
    service.log_file = service.log_path.open("w", encoding="utf-8", buffering=1)
    creationflags = 0
    start_new_session = False
    if os.name == "nt":
        creationflags = getattr(subprocess, "CREATE_NEW_PROCESS_GROUP", 0)
    else:
        start_new_session = True

    try:
        service.process = subprocess.Popen(
            service.command,
            cwd=service.cwd,
            shell=True,
            stdout=service.log_file,
            stderr=subprocess.STDOUT,
            creationflags=creationflags,
            start_new_session=start_new_session,
        )
        service.started_by_script = True
    except Exception:
        service.log_file.close()
        service.log_file = None
        raise

    print(f"启动 {service.name}：进程={service.process.pid}，工作目录={service.cwd}")


def request(url: str, timeout: float, request_id: str) -> tuple[bool, str]:
    request_obj = urllib.request.Request(
        url,
        headers={
            "User-Agent": "fullstack-start-verify/1.1",
            "X-Request-Id": request_id,
        },
    )
    try:
        with urllib.request.urlopen(request_obj, timeout=timeout) as response:
            status = response.status
            return 200 <= status < 400, str(status)
    except urllib.error.HTTPError as error:
        return 200 <= error.code < 400, str(error.code)
    except (urllib.error.URLError, TimeoutError, OSError) as error:
        reason = getattr(error, "reason", error)
        return False, type(reason).__name__


def wait_for_service(
    service: Service, timeout: float, interval: float, run_id: str
) -> None:
    if service.process is None:
        raise RuntimeError(f"{service.name} 未启动")

    deadline = time.monotonic() + timeout
    last_reason = "not checked"
    while time.monotonic() < deadline:
        exit_code = service.process.poll()
        if exit_code is not None:
            raise RuntimeError(f"{service.name} 在就绪检查前退出，退出码为 {exit_code}")

        request_timeout = min(max(interval * 2, 0.5), 5.0)
        ok, reason = request(service.url, request_timeout, f"{run_id}-{service.name}")
        last_reason = reason
        if ok:
            print(f"通过 {service.name}：{service.url} [{reason}]")
            return
        time.sleep(min(max(interval, 0.05), max(deadline - time.monotonic(), 0)))

    raise RuntimeError(f"{service.name} 在 {timeout:g} 秒内未就绪 [{last_reason}]")


def parse_extra_checks(values: list[str]) -> list[tuple[str, str]]:
    checks: list[tuple[str, str]] = []
    for value in values:
        if "=" in value:
            label, url = value.split("=", 1)
        else:
            label, url = value, value
        label = label.strip() or url.strip()
        url = url.strip()
        if not url.startswith(("http://", "https://")):
            raise RuntimeError(f"检查地址无效：{url}")
        checks.append((label, url))
    return checks


def wait_for_services(
    services: list[Service], timeout: float, interval: float, run_id: str
) -> None:
    """并行等待服务，整体耗时取最慢服务的就绪时间。"""
    executor = concurrent.futures.ThreadPoolExecutor(max_workers=len(services))
    failed = True
    try:
        futures = [
            executor.submit(wait_for_service, service, timeout, interval, run_id)
            for service in services
        ]
        done, _ = concurrent.futures.wait(
            futures, return_when=concurrent.futures.FIRST_EXCEPTION
        )
        for future in done:
            future.result()
        failed = False
    finally:
        executor.shutdown(wait=not failed, cancel_futures=failed)


def reuse_running_services(
    services: list[Service], timeout: float, run_id: str
) -> None:
    """并行探测已运行服务，避免重复启动本项目服务。"""
    if not services:
        return

    probe_timeout = min(max(timeout, 0.5), 1.0)

    def probe(service: Service) -> tuple[Service, bool, str]:
        ok, reason = request(
            service.url, probe_timeout, f"{run_id}-{service.name}-reuse"
        )
        return service, ok, reason

    with concurrent.futures.ThreadPoolExecutor(max_workers=len(services)) as executor:
        results = list(executor.map(probe, services))
    for service, ok, reason in results:
        if ok:
            service.reused = True
            print(f"复用 {service.name}：{service.url} [{reason}]")


def wait_for_services_and_checks(
    services: list[Service],
    checks: list[tuple[str, str]],
    timeout: float,
    interval: float,
    check_timeout: float,
    run_id: str,
) -> None:
    """后端就绪后立即检查 API，同时继续等待前端，缩短总耗时。"""
    if not services:
        run_extra_checks(checks, check_timeout, run_id)
        return
    if not checks:
        wait_for_services(services, timeout, interval, run_id)
        return

    executor = concurrent.futures.ThreadPoolExecutor(
        max_workers=len(services) + max(1, len(checks))
    )
    failed = True
    try:
        service_futures = [
            (
                service,
                executor.submit(wait_for_service, service, timeout, interval, run_id),
            )
            for service in services
        ]
        backend_future = next(
            (
                future
                for service, future in service_futures
                if service.name == "backend"
            ),
            None,
        )
        pending = {future for _, future in service_futures}
        check_future: concurrent.futures.Future[None] | None = None
        if backend_future is None:
            check_future = executor.submit(run_extra_checks, checks, check_timeout, run_id)
            pending.add(check_future)

        while pending:
            done, pending = concurrent.futures.wait(
                pending, return_when=concurrent.futures.FIRST_COMPLETED
            )
            for future in done:
                future.result()
                if future is backend_future and check_future is None:
                    check_future = executor.submit(
                        run_extra_checks, checks, check_timeout, run_id
                    )
                    pending.add(check_future)
        failed = False
    finally:
        executor.shutdown(wait=not failed, cancel_futures=failed)


def run_extra_checks(
    checks: list[tuple[str, str]], timeout: float, run_id: str
) -> None:
    """并行执行只读接口检查，避免多个接口检查串行等待。"""
    if not checks:
        return

    def execute_check(
        item: tuple[int, tuple[str, str]]
    ) -> tuple[str, str, tuple[bool, str]]:
        index, (label, url) = item
        return label, url, request(url, timeout, f"{run_id}-check-{index}")

    with concurrent.futures.ThreadPoolExecutor(max_workers=len(checks)) as executor:
        results = list(executor.map(execute_check, enumerate(checks, start=1)))
    for label, url, (ok, reason) in results:
        if not ok:
            raise RuntimeError(f"检查失败：{label} {url} [{reason}]")
        print(f"通过 {label}：{url} [{reason}]")


def stop_service(service: Service) -> None:
    process = service.process
    if not service.started_by_script or process is None or process.poll() is not None:
        return

    if os.name == "nt":
        try:
            subprocess.run(
                ["taskkill", "/PID", str(process.pid), "/T", "/F"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                check=False,
                timeout=2,
            )
        except subprocess.TimeoutExpired:
            pass
        try:
            process.wait(timeout=1)
        except subprocess.TimeoutExpired:
            process.kill()
            try:
                process.wait(timeout=1)
            except subprocess.TimeoutExpired:
                process.returncode = -1
        except OSError:
            # taskkill 已经结束进程，但 Windows 可能同时使句柄失效。
            process.returncode = -1
    else:
        try:
            os.killpg(process.pid, signal.SIGTERM)
        except ProcessLookupError:
            return
        try:
            process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            os.killpg(process.pid, signal.SIGKILL)


def tail(path: Path | None, lines: int = 20) -> list[str]:
    if path is None or not path.exists():
        return []
    return path.read_text(encoding="utf-8", errors="replace").splitlines()[-lines:]


def redact(line: str) -> str:
    """仅在终端输出日志时脱敏，避免错误摘要泄露凭据。"""
    line = SENSITIVE_VALUE_PATTERN.sub(r"\1\2[REDACTED]", line)
    return BEARER_PATTERN.sub(r"\1[REDACTED]", line)


def main() -> int:
    args = parse_args()
    if args.timeout <= 0 or args.interval <= 0 or args.check_timeout <= 0:
        raise SystemExit("timeout、interval 和 check-timeout 必须大于 0")
    run_id = uuid.uuid4().hex[:12]
    print(f"验证任务 ID：{run_id}")
    checks = parse_extra_checks(args.check_url)
    temp_log_dir: Path | None = None
    log_dir = args.log_dir
    if log_dir is None:
        temp_log_dir = Path(tempfile.mkdtemp(prefix="fullstack-start-verify-"))
        log_dir = temp_log_dir
    log_dir = log_dir.resolve()
    log_dir.mkdir(parents=True, exist_ok=True)

    services = [
        Service("backend", args.backend_cmd, resolve_dir(args.backend_dir), args.backend_url),
        Service("frontend", args.frontend_cmd, resolve_dir(args.frontend_dir), args.frontend_url),
    ]
    failure: Exception | None = None

    try:
        if args.reuse_running:
            reuse_running_services(services, args.check_timeout, run_id)
        for service in services:
            if not service.reused:
                start_service(service, log_dir)
        pending_services = [service for service in services if not service.reused]
        wait_for_services_and_checks(
            pending_services,
            checks,
            args.timeout,
            args.interval,
            args.check_timeout,
            run_id,
        )
        print(f"全栈验证通过（FULLSTACK_VERIFY=PASS，任务 ID：{run_id}）")
        return 0
    except (OSError, RuntimeError, ValueError) as error:
        failure = error
        print(f"全栈验证失败：{error}", file=sys.stderr)
        for service in services:
            for line in tail(service.log_path):
                print(f"日志 {service.name}：{redact(line)}", file=sys.stderr)
        return 1
    finally:
        if not args.keep_running:
            for service in reversed(services):
                stop_service(service)
        for service in services:
            if service.log_file is not None:
                service.log_file.close()
        if args.keep_running:
            print(f"服务仍在运行，日志目录：{log_dir}")
        elif temp_log_dir is not None and not args.keep_logs:
            shutil.rmtree(temp_log_dir, ignore_errors=True)
        elif temp_log_dir is not None:
            print(f"日志已保留：{temp_log_dir}")
        elif failure is not None:
            print(f"日志已保留：{log_dir}")


if __name__ == "__main__":
    raise SystemExit(main())