AiDOPWarehouse/tools/integration-simulator/load/controls.py

"""压力测试安全控制（任务书 §16 / §20 / §21 验收表）。

不变量（对应第三阶段验收表）：
- 压力测试默认关闭：LoadGuard 初始 idle，无完整计划不可启动；
- §16 执行前十项必填，缺一拒绝；共享环境须负责人确认（confirmed）；
- 参数保守上限，提高须 escalationConfirm 显式二次确认，且不得越过硬上限；
- 达到连续错误 / 超时 / 429 阈值自动停止；不自动无限重试；
- 立即停止：stop 信号 + 在途请求收尾等待 + 无残留负载断言；
- 每轮唯一 runId；P3 行构造禁止携带 tenantId（§20 安全限制）。
"""
from __future__ import annotations

import threading
import time
import uuid

CHANNELS = ("DB_SYNC", "API_PULL", "API_INBOUND")

# §16 执行前必填项（十项）
REQUIRED_PLAN_FIELDS = (
    "environment",          # 环境
    "tenant",               # 测试租户
    "owner",                # 负责人
    "timeWindow",           # 时间窗口
    "maxDataRows",          # 最大数据量
    "maxConcurrency",       # 最大并发
    "maxQps",               # 最大 QPS
    "maxDurationSeconds",   # 最大持续时间
    "stopThreshold",        # 停止阈值（如：连续 5 次错误/超时/429 自动停止）
    "cleanupPlan",          # 清理方案
)

# 保守默认上限（§20 安全限制）。提高须 escalationConfirm=True，且不得超过 HARD_LIMITS
DEFAULT_LIMITS = {
    "maxDataRows": 100_000,        # 与 L 档一致；XL 须审批
    "maxConcurrency": 20,
    "maxQps": 50,
    "maxDurationSeconds": 600,
    "maxBatchSize": 100,
    "errorRate": 0.5,
    "maxRetries": 3,
    "timeoutSeconds": 30,
    "maxConsecutiveErrors": 5,
    "maxTimeouts": 5,
    "maxRateLimited": 5,
    "rateLimit": 200,          # P2 Mock 每秒最大请求数
    "randomDelay": 10.0,       # P2 随机延迟上限（秒）
}

# 硬上限：即使显式二次确认也不允许越过（非测试环境阻断的兜底）
HARD_LIMITS = {
    "maxDataRows": 200_000,
    "maxConcurrency": 50,
    "maxQps": 100,
    "maxDurationSeconds": 1800,
    "maxBatchSize": 500,
    "errorRate": 0.9,
    "maxRetries": 5,
    "timeoutSeconds": 120,
    "maxConsecutiveErrors": 100,
    "maxTimeouts": 100,
    "maxRateLimited": 100,
    "rateLimit": 1000,
    "randomDelay": 60.0,
}

# 行内禁止携带的租户字段（§20：请求体不得携带任意 tenantId 改变目标租户）
_TENANT_FIELD_NAMES = ("tenantid", "tenant_id", "tenantcode", "tenant_code",
                       "domainid", "domain_id", "orgid", "org_id")


def strip_tenant_fields(row: dict) -> dict:
    """剔除行内租户/组织字段，防止 payload 切换目标租户。"""
    return {k: v for k, v in row.items() if str(k).lower() not in _TENANT_FIELD_NAMES}


class ValidationError(ValueError):
    pass


def validate_plan(plan: dict) -> list[str]:
    """校验压力计划，返回错误列表（空列表 = 通过）。"""
    errors: list[str] = []
    if plan.get("channel") not in CHANNELS:
        errors.append("channel 必须是 DB_SYNC / API_PULL / API_INBOUND")
    for field in REQUIRED_PLAN_FIELDS:
        if plan.get(field) in (None, "", 0):
            errors.append(f"缺少必填项：{field}（§16 执行前必须填写）")
    if not bool(plan.get("confirmed")):
        errors.append("共享开发/UAT 环境须由负责人确认后执行（confirmed=true）")
    escalated = bool(plan.get("escalationConfirm"))
    float_keys = ("errorRate", "randomDelay")
    for key, cap in DEFAULT_LIMITS.items():
        raw = plan.get(key)
        if raw in (None, "", 0):
            continue
        try:
            value = float(raw) if key in float_keys else int(raw)
        except (TypeError, ValueError):
            errors.append(f"{key} 必须是数字，收到：{raw!r}")
            continue
        if key == "errorRate" and not 0 <= value <= 1:
            errors.append(f"errorRate 必须在 0-1 之间")
            continue
        if value > cap and not escalated:
            errors.append(f"{key}={value} 超过保守上限 {cap}，需勾选「提高上限二次确认」")
        if value > HARD_LIMITS[key]:
            errors.append(f"{key}={value} 超过硬上限 {HARD_LIMITS[key]}，即使二次确认也不允许")
    return errors


class AutoStopTracker:
    """自动停止判定（§20 安全限制：连续错误、超时或 429 达阈值自动停止）。"""

    def __init__(self, consecutive_errors: int = 5, timeouts: int = 5, rate_limited: int = 5):
        self.max_consecutive_errors = max(1, consecutive_errors)
        self.max_timeouts = max(1, timeouts)
        self.max_rate_limited = max(1, rate_limited)
        self._consecutive_errors = 0
        self._timeouts = 0
        self._rate_limited = 0
        self._tripped_by: str | None = None
        self._lock = threading.Lock()

    def record(self, status: int) -> bool:
        """记录一次请求结果；返回 True 表示应自动停止。

        status 约定：0 = 网络超时/不可达；429 = 限流；其余为 HTTP 状态码。
        """
        with self._lock:
            if status == 0:
                self._timeouts += 1
                self._consecutive_errors += 1
            elif status == 429:
                self._rate_limited += 1
                self._consecutive_errors += 1
            elif status >= 400:
                self._consecutive_errors += 1
            else:
                self._consecutive_errors = 0
            if self._tripped_by is None:
                if self._consecutive_errors >= self.max_consecutive_errors:
                    self._tripped_by = "consecutiveErrors"
                elif self._timeouts >= self.max_timeouts:
                    self._tripped_by = "timeouts"
                elif self._rate_limited >= self.max_rate_limited:
                    self._tripped_by = "rateLimited"
            return self._tripped_by is not None

    @property
    def tripped_by(self) -> str | None:
        return self._tripped_by

    def summary(self) -> dict:
        return {"consecutiveErrors": self._consecutive_errors,
                "timeouts": self._timeouts, "rateLimited": self._rate_limited,
                "thresholds": {"maxConsecutiveErrors": self.max_consecutive_errors,
                               "maxTimeouts": self.max_timeouts,
                               "maxRateLimited": self.max_rate_limited},
                "trippedBy": self._tripped_by}


class LoadGuard:
    """压力测试生命周期守卫（进程内单例，重启即默认关闭）。

    状态机：idle -> ready -> running -> stopping -> stopped
    """

    def __init__(self):
        self.state = "idle"            # idle | ready | running | stopping | stopped
        self.plan: dict | None = None
        self.run_id: str | None = None
        self.started_at: str | None = None
        self.finished_at: str | None = None
        self.stop_reason: str | None = None
        self.stop = threading.Event()
        self._workers: list[threading.Thread] = []
        self._lock = threading.Lock()

    # -- 计划 ----------------------------------------------------------------

    def submit_plan(self, plan: dict) -> dict:
        """保存并校验计划。校验失败不改变当前状态。"""
        errors = validate_plan(plan)
        if errors:
            return {"ok": False, "errors": errors, "state": self.state}
        with self._lock:
            self.plan = dict(plan)
            self.state = "ready"
            self.stop.clear()
            self.stop_reason = None
            self.finished_at = None
        return {"ok": True, "state": self.state, "runId": self.plan.get("runId") or ""}

    # -- 生命周期 ------------------------------------------------------------

    def start(self, run_id: str | None = None) -> str:
        """ready -> running；返回本轮 runId（每轮唯一）。"""
        with self._lock:
            if self.state != "ready":
                raise ValidationError(f"当前状态 {self.state}，无法开始（须先提交并通过计划校验）")
            self.run_id = run_id or f"LOAD-{uuid.uuid4().hex[:12]}"
            self.started_at = time.strftime("%Y-%m-%dT%H:%M:%S")
            self.state = "running"
            self.stop.clear()
            return self.run_id

    def request_stop(self, reason: str = "manual") -> dict:
        """请求停止：置信号，状态 -> stopping（在途请求由执行器收尾）。"""
        with self._lock:
            self.stop_reason = self.stop_reason or reason
            if self.state in ("ready", "running"):
                self.state = "stopping"
        self.stop.set()
        return {"ok": True, "state": self.state, "reason": self.stop_reason}

    def finish(self, reason: str | None = None) -> dict:
        """执行器收尾调用：stopping/running -> stopped，并断言无残留负载。"""
        with self._lock:
            if self.state not in ("running", "stopping"):
                raise ValidationError(f"当前状态 {self.state}，无法收尾")
            # 排除当前线程：P1/P3 执行器线程会注册自身为 worker，
            # 它在 finish() 调用瞬间必然存活，不应误报为残留
            current = threading.current_thread()
            residual = [t.name for t in self._workers
                        if t.is_alive() and t is not current]
            if residual:
                # 在途收尾不彻底：如实上报，不伪装干净
                self.state = "stopped"
                self.finished_at = time.strftime("%Y-%m-%dT%H:%M:%S")
                self.stop_reason = self.stop_reason or reason or "finished"
                return {"ok": False, "state": self.state,
                        "reason": self.stop_reason or "finished",
                        "residualThreads": residual}
            self.state = "stopped"
            self.finished_at = time.strftime("%Y-%m-%dT%H:%M:%S")
            self.stop_reason = self.stop_reason or reason or "finished"
            return {"ok": True, "state": self.state, "reason": self.stop_reason,
                    "residualThreads": []}

    def register_worker(self, worker: threading.Thread) -> None:
        with self._lock:
            self._workers.append(worker)

    def no_residual(self) -> bool:
        with self._lock:
            return all(not t.is_alive() for t in self._workers)

    # -- 快照 ----------------------------------------------------------------

    def snapshot(self) -> dict:
        with self._lock:
            return {
                "state": self.state,
                "runId": self.run_id,
                "startedAt": self.started_at,
                "finishedAt": self.finished_at,
                "stopReason": self.stop_reason,
                "planChannel": (self.plan or {}).get("channel"),
                "planConfirmed": bool((self.plan or {}).get("confirmed")),
                # 持锁内直接检查，避免重入 no_residual() 死锁（Lock 非重入）
                "noResidual": all(not t.is_alive() for t in self._workers),
            }