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

"""方式丙压测执行器（任务书 §20）：模拟器是负载发起方。

能力映射：
- 并发 worker（ThreadPoolExecutor，plan.maxConcurrency）；
- QPS 限制（槽位法节流器，实际 QPS 接近目标且偏差可解释）；
- 批大小（bulk 模式每请求行数 = plan.batchSize）；
- 持续时间 / 阶梯配置（rampSteps，缺省固定档）；
- 错误数据比例（errorRate 概率缺必填字段，与 N3 partial-reject 同语义）；
- 幂等重复（idemReplay：复用同一 idempotencyKey 重发）；
- 立即停止（guard.stop 信号，不再调度新请求，等待在途收尾）；
- 达到连续错误/超时/429 阈值自动停止（AutoStopTracker）；
- 每轮唯一 runId；行构造剔除租户字段（§20 安全限制）。
"""
from __future__ import annotations

import random
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

_SIM_ROOT = Path(__file__).resolve().parents[1]
if str(_SIM_ROOT) not in sys.path:
    sys.path.insert(0, str(_SIM_ROOT))

from catalog import sample_adapters as sa  # noqa: E402
from load import controls  # noqa: E402
from load.metrics import LoadMetrics  # noqa: E402


class _RateLimiter:
    """固定 QPS 节流器（槽位法：请求间隔 = 1/QPS）。

    近似实现，实际 QPS 受网络与并发影响，偏差在报告中可解释（§20 验收）。
    """

    def __init__(self, qps: float):
        self.interval = 1.0 / qps if qps > 0 else 0.0
        self._next = 0.0
        self._lock = threading.Lock()

    def wait(self) -> None:
        if self.interval <= 0:
            return
        with self._lock:
            now = time.monotonic()
            slot = self._next if now < self._next else now
            self._next = slot + self.interval
        delay = slot - now
        if delay > 0:
            time.sleep(delay)


class P3Runner(threading.Thread):
    """后台线程运行方式丙压测；结束后把 LoadGuard 收尾（含无残留断言）。"""

    def __init__(self, client, entity: str, template_rows: list[dict], plan: dict,
                 guard: controls.LoadGuard, metrics: LoadMetrics, *,
                 contract_version: str | None = None):
        super().__init__(name=f"p3-runner-{plan.get('runId', 'x')}", daemon=True)
        self.client = client
        self.entity = entity
        self.template_rows = [dict(r) for r in template_rows]
        self.plan = plan
        self.guard = guard
        self.metrics = metrics
        self.contract_version = contract_version
        self.tracker = controls.AutoStopTracker(
            int(plan.get("maxConsecutiveErrors") or 5),
            int(plan.get("maxTimeouts") or 5),
            int(plan.get("maxRateLimited") or 5))
        self.mode = plan.get("mode") or "push"            # push | bulk
        self.batch_size = max(1, int(plan.get("batchSize") or 1))
        self.error_rate = float(plan.get("errorRate") or 0.0)
        self.idem_replay = bool(plan.get("idemReplay"))
        self._first_required = self._resolve_first_required_field()
        self._replay_idem = f"LOAD-REPLAY-{random.randrange(1 << 32):08x}"

    # -- 行构造 --------------------------------------------------------------

    def _resolve_first_required_field(self) -> str | None:
        spec = sa.INBOUND_SPECS.get(self.entity) or []
        for field, required, _ in spec:
            if required:
                return field
        return None

    def _build_row(self, seq: int) -> dict:
        """固定/随机样例 + 序号唯一化 + 租户字段剔除。"""
        row = dict(random.choice(self.template_rows)) if self.plan.get("sampleMode") == "random" \
            else dict(self.template_rows[seq % len(self.template_rows)])
        row = controls.strip_tenant_fields(row)
        for key in ("bizKey", "Id"):
            if key in row:
                row[key] = f"{row[key]}#{self.guard.run_id}-{seq}"
        row["sourceUpdatedAt"] = sa.now_iso()
        return row

    def _one_request(self, seq: int) -> None:
        rows = [self._build_row(seq)]
        bad = self.error_rate > 0 and random.random() < self.error_rate
        if bad and self._first_required:
            # 部分错误数据：缺必填字段（同 N3 partial-reject 语义）
            rows[0].pop(self._first_required, None)
        started = time.monotonic()
        kwargs: dict = {}
        if self.idem_replay:
            kwargs["idem"] = self._replay_idem
        try:
            if self.mode == "bulk":
                if self.batch_size > 1:
                    rows.extend(self._build_row(seq * 1000 + i) for i in range(1, self.batch_size))
                lines = [{"idempotencyKey": f"{self.guard.run_id}-L{seq}-{i}",
                          "data": {"list": [r]}} for i, r in enumerate(rows)]
                status, _, _ = self.client.op_bulk(self.entity, lines,
                                                   contract_version=self.contract_version)
            else:
                envelope = sa.inbound_envelope(rows)
                status, _, _ = self.client.op_push(self.entity, envelope,
                                                   contract_version=self.contract_version, **kwargs)
            latency = time.monotonic() - started
            self.metrics.record(status, latency, rows=len(rows))
            if self.tracker.record(status):
                self.guard.request_stop(f"自动停止：{self.tracker.tripped_by} 达到阈值")
        except Exception as ex:  # noqa: BLE001 —— 压测收集异常而非中断
            self.metrics.record(0, time.monotonic() - started, rows=len(rows))
            self.guard.stop_reason = self.guard.stop_reason or f"worker 异常：{type(ex).__name__}"
            if self.tracker.record(0):
                self.guard.request_stop(f"自动停止：{self.tracker.tripped_by} 达到阈值")

    # -- 阶梯 ----------------------------------------------------------------

    def _ramp_steps(self) -> list[dict]:
        steps = self.plan.get("rampSteps") or []
        if steps:
            return steps
        return [{
            "duration": float(self.plan.get("maxDurationSeconds") or 10),
            "qps": float(self.plan.get("maxQps") or 1),
            "concurrency": int(self.plan.get("maxConcurrency") or 1),
        }]

    def _run_step(self, step: dict) -> None:
        duration = max(0.5, float(step.get("duration") or 1))
        qps = max(0.0, float(step.get("qps") or 0))
        concurrency = max(1, int(step.get("concurrency") or 1))
        limiter = _RateLimiter(qps)
        deadline = time.monotonic() + duration
        futures = []
        seq = 0
        with ThreadPoolExecutor(max_workers=concurrency) as pool:
            while time.monotonic() < deadline and not self.guard.stop.is_set():
                limiter.wait()
                if self.guard.stop.is_set():
                    break
                seq += 1
                futures.append(pool.submit(self._one_request, seq))
                if len(futures) >= concurrency * 8:
                    futures = [f for f in futures if not f.done()]
            # 停止后：取消排队未启动的任务，只等待真正在途的请求收尾（§20）
            pool.shutdown(wait=True, cancel_futures=True)

    def run(self) -> None:
        self.guard.register_worker(self)
        warmup = max(0, int(self.plan.get("warmupSeconds") or 0))
        if warmup and not self.guard.stop.is_set():
            # §17 预热时间：等待环境稳定后再发压（可被立即停止打断）
            self.guard.stop.wait(timeout=warmup)
        self.metrics.mark_start()
        try:
            for step in self._ramp_steps():
                if self.guard.stop.is_set():
                    break
                self._run_step(step)
        finally:
            self.metrics.mark_finish()
            result = self.guard.finish()
            if not result.get("ok"):
                # 在途线程未收尾：记录到 stop_reason 供报告如实呈现
                self.guard.stop_reason = (self.guard.stop_reason or "") + \
                    f"（残留线程：{result.get('residualThreads')}）"