#!/usr/bin/env python3
"""Sweep concurrent-cycling sandbox counts and measure pause/resume behavior.

针对"高并发 pause/resume 积压"问题的探针脚本:建立一组沙箱,逐级提高
同时循环 pause -> resume 的沙箱数 R,每一级记录:

  - 客户端视角:每次 pause / resume / --command 的墙钟时延(p50/p95/均值);
    --command 额外记录三层分解:客户端全链路、guest 内纯执行耗时
    (bash 纳秒计时回显)、服务端代理 envd RPC 的耗时(/metrics);
  - 服务端视角:/metrics 中 pause stage 的次数与总耗时增量(服务端均值)、
    pause in-flight 峰值;
  - 进程视角(/proc,需与被测 server 同机或共享 PID namespace):server、
    ublk daemon、firecracker 三组进程的 CPU(核,user/sys 分列)与 RSS
    均值/峰值;firecracker 进程数即在线 VM 数;
  - Little 定律交叉验证:max_inflight / 实测吞吐 ≈ 端到端平均时延。

用法示例:

  # 本机默认 4 级:1/2/4/8 个并发循环沙箱,每级 8 轮
  python measure_pause_resume_concurrency.py \
      --api-url http://127.0.0.1:8000 --template my-template

  # 更大压力,带 1s 驻留(模拟轨迹 delay),结果写到文件
  python measure_pause_resume_concurrency.py \
      --levels 1,2,4,8,16,32 --rounds 10 --think-sec 1.0 \
      --output pause-resume-sweep.json

  # 不限并发:整池同时循环(--max-pool 即并发数)
  python measure_pause_resume_concurrency.py \
      --template my-template --levels all --max-pool 384 --rounds 20 \
      --command 'dd if=/dev/urandom of=/testbed/probe.bin bs=1048576 count=80 status=none' \
      --output pause-resume-unlimited.json

结果输出:默认写入 runs/pause-resume-<时间戳>.json,并在旁边生成同名
.csv(每个 level 一行,嵌套字段拍平成列);--output 可显式指定 JSON 路径,
--output-dir 改默认目录(相对当前工作目录)。

判读:客户端 p50 稳定、p95 随 R 上升 => 排队;服务端均值突增的 R 即
服务端 pause 服务能力(μ)附近的拐点。把并发压到拐点以下即可同时拿到
低时延和高吞吐。
"""

from __future__ import annotations

import argparse
import asyncio
import csv
import json
import logging
import os
import random
import re
import statistics
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Any

LOGGER = logging.getLogger("pause-resume-probe")

PAUSE_STAGE_METRIC = "agentenv_sandbox_stage_duration_seconds"
PAUSE_INFLIGHT_METRIC = "agentenv_sandbox_stage_inflight"
HTTP_DURATION_METRIC = "agentenv_http_request_duration_seconds"

# --command 在 guest 内由 bash 包一层纳秒计时,stdout 里回显该标记。
PROBE_GUEST_NS_RE = re.compile(r"__aenv_probe_ns=(\d+)")

# /proc 采样换算用:CPU 节拍频率与内存页大小。
CLK_TCK = float(os.sysconf("SC_CLK_TCK"))
PAGE_BYTES = int(os.sysconf("SC_PAGE_SIZE"))


def now_iso() -> str:
    return datetime.now().astimezone().isoformat(timespec="milliseconds")


def percentile(values: list[float], q: float) -> float | None:
    if not values:
        return None
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, round(q / 100 * (len(ordered) - 1))))
    return ordered[index]


def summarize(values: list[float]) -> dict[str, Any]:
    if not values:
        return {"n": 0, "p50": None, "p95": None, "avg": None, "max": None}
    return {
        "n": len(values),
        "p50": round(percentile(values, 50), 4),
        "p95": round(percentile(values, 95), 4),
        "avg": round(statistics.fmean(values), 4),
        "max": round(max(values), 4),
    }


# ---------------------------------------------------------------------------
# Minimal Prometheus exposition-format scraper (counters/gauges only).
# ---------------------------------------------------------------------------


def parse_prom_text(text: str) -> dict[tuple[str, tuple[tuple[str, str], ...]], float]:
    """Return {(metric_name, sorted labels): value} for non-bucket samples."""

    samples: dict[tuple[str, tuple[tuple[str, str], ...]], float] = {}
    for raw_line in text.splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        name_part, _, value_part = line.rpartition(" ")
        if not name_part:
            continue
        try:
            value = float(value_part)
        except ValueError:
            continue  # NaN / +Inf / timestamps are irrelevant for this probe.
        labels: tuple[tuple[str, str], ...] = ()
        if name_part.endswith("}") and "{" in name_part:
            name, _, label_text = name_part.partition("{")
            label_text = label_text[: -1]
            parsed = []
            for item in label_text.split('",'):
                if not item:
                    continue
                key, _, val = item.partition('="')
                parsed.append((key.strip(), val.rstrip('"')))
            labels = tuple(sorted(parsed))
            metric_name = name
        else:
            metric_name = name_part
        if metric_name.endswith("_bucket") or metric_name.endswith("_created"):
            continue
        samples[(metric_name, labels)] = samples.get((metric_name, labels), 0.0) + value
    return samples


def stage_value(
    samples: dict[tuple[str, tuple[tuple[str, str], ...]], float],
    metric: str,
    suffix: str,
    operation: str,
) -> float | None:
    """Sum a stage-series sample across its status/stage label variants."""

    key_name = f"{metric}{suffix}"
    total = 0.0
    found = False
    for (name, labels), value in samples.items():
        if name != key_name:
            continue
        label_map = dict(labels)
        if label_map.get("operation") != operation:
            continue
        if "status" in label_map and label_map["status"] not in ("ok", "error", "canceled"):
            continue
        total += value
        found = True
    return total if found else None


def inflight_value(
    samples: dict[tuple[str, tuple[tuple[str, str], ...]], float],
    operation: str,
) -> float | None:
    total = 0.0
    found = False
    for (name, labels), value in samples.items():
        if name != PAUSE_INFLIGHT_METRIC:
            continue
        if dict(labels).get("operation") != operation:
            continue
        total += value
        found = True
    return total if found else None


def unmatched_http_totals(
    samples: dict[tuple[str, tuple[tuple[str, str], ...]], float],
    route_source: str,
) -> tuple[float, float]:
    """(count, sum) of unmatched POST HTTP-duration series from route_source.

    route_source="control_plane" 即 SDK connect(resume)请求;
    "proxy_header" 即 envd RPC(--command 转发)请求。按 level 差分可得到
    服务端视角的 resume / command 耗时,绕开 Prometheus rate 窗口稀释。
    """

    count = 0.0
    total = 0.0
    for (name, labels), value in samples.items():
        if name not in (
            f"{HTTP_DURATION_METRIC}_count",
            f"{HTTP_DURATION_METRIC}_sum",
        ):
            continue
        label_map = dict(labels)
        if label_map.get("route") != "unmatched":
            continue
        if label_map.get("route_source") != route_source:
            continue
        if label_map.get("method") != "POST":
            continue
        if name.endswith("_count"):
            count += value
        else:
            total += value
    return count, total


# ---------------------------------------------------------------------------
# /proc 进程组采样:server / ublk daemon / firecracker 的 CPU 与 RSS。
# ---------------------------------------------------------------------------


def _split_proc_names(value: str) -> tuple[str, ...]:
    return tuple(item.strip() for item in value.split(",") if item.strip())


def _read_proc_stat(
    proc_root: Path,
    pid: int,
) -> tuple[str, str, int, int, int] | None:
    """(comm, state, utime, stime, rss_pages) from /proc/<pid>/stat, or None."""

    try:
        data = (proc_root / str(pid) / "stat").read_bytes()
    except OSError:
        return None
    close = data.rfind(b")")
    if close < 0:
        return None
    comm = data[data.find(b"(") + 1 : close].decode("utf-8", "replace")
    fields = data[close + 2 :].split()
    # fields[0] 是第 3 列 state;utime/stime/rss 是第 14/15/24 列。
    if len(fields) < 22:
        return None
    try:
        utime, stime, rss_pages = int(fields[11]), int(fields[12]), int(fields[21])
    except ValueError:
        return None
    return comm, fields[0].decode("ascii", "replace"), utime, stime, rss_pages


@dataclass
class _ProcGroupAccum:
    samples: int = 0
    user_ticks: int = 0
    sys_ticks: int = 0
    rss_sum_mib: float = 0.0
    rss_peak_mib: float = 0.0
    count_sum: int = 0
    count_max: int = 0


class ProcStatsSampler:
    """Per-level CPU/RSS sampling of the server/ublk/firecracker process groups.

    直接读本机 /proc,因此必须与目标进程同 PID namespace(即探针跑在
    server 所在主机或共享其 /proc),否则用 --no-proc-stats 关闭。CPU 按
    tick 差分累计:中途退出的进程只丢最后一个不满 tick 的量(pause 正是
    firecracker 退出前做内存落盘,故默认轮询间隔 0.25s 远快于 /metrics 间隔);
    中途新出现的进程自启动起的 CPU 全部归入它首次出现的那个 tick,语义
    正好覆盖 VM 生灭开销。周期扫描在线程池执行,不阻塞探针事件循环。
    """

    def __init__(
        self,
        matches: dict[str, tuple[str, ...]],
        interval_sec: float,
        proc_root: str = "/proc",
    ):
        self.matches = {group: frozenset(names) for group, names in matches.items()}
        self.interval_sec = interval_sec
        self.proc_root = Path(proc_root)
        self._task: asyncio.Task[None] | None = None
        self._last: dict[int, tuple[int, int]] | None = None
        self._accum: dict[str, _ProcGroupAccum] = {}

    def _scan(self, accumulate: bool = True) -> dict[str, int]:
        current: dict[int, tuple[int, int]] = {}
        found: dict[str, list[tuple[int, int, int, int]]] = {
            group: [] for group in self.matches
        }
        try:
            entries = os.listdir(self.proc_root)
        except OSError as exc:
            LOGGER.warning("proc scan failed: %s: %r", type(exc).__name__, exc)
            return {group: 0 for group in self.matches}
        for entry in entries:
            if not entry.isdigit():
                continue
            pid = int(entry)
            record = _read_proc_stat(self.proc_root, pid)
            if record is None:
                continue
            comm, state, utime, stime, rss_pages = record
            if state in ("Z", "X", "x"):
                continue
            current[pid] = (utime, stime)
            for group, patterns in self.matches.items():
                if comm in patterns:
                    found[group].append((pid, utime, stime, rss_pages))
                    break
        if accumulate and self._last is not None:
            for group, procs in found.items():
                acc = self._accum[group]
                acc.samples += 1
                acc.count_sum += len(procs)
                acc.count_max = max(acc.count_max, len(procs))
                rss_mib = 0.0
                for pid, utime, stime, rss_pages in procs:
                    base_utime, base_stime = self._last.get(pid, (0, 0))
                    # PID 复用可能让差分为负,钳到 0。
                    acc.user_ticks += max(0, utime - base_utime)
                    acc.sys_ticks += max(0, stime - base_stime)
                    rss_mib += rss_pages * PAGE_BYTES / (1024 * 1024)
                acc.rss_sum_mib += rss_mib
                acc.rss_peak_mib = max(acc.rss_peak_mib, rss_mib)
        self._last = current
        return {group: len(procs) for group, procs in found.items()}

    def start_level(self) -> None:
        self._accum = {group: _ProcGroupAccum() for group in self.matches}
        # 基线扫描只记录当前计数不计差分:server 等常驻进程 level 开始前
        # 累积的 CPU 时间不应算进本 level。
        counts = self._scan(accumulate=False)
        LOGGER.info("proc groups matched: %s", counts)
        if "server" in self.matches and counts.get("server", 0) == 0:
            LOGGER.warning(
                "no process matched server patterns; if the daemon runs under a "
                "different name pass --proc-match-server"
            )

        async def sample_loop() -> None:
            while True:
                await asyncio.sleep(self.interval_sec)
                # 扫上千个进程的 /proc 是几十毫秒级阻塞 I/O,必须丢进线程:
                # 同步执行会周期性卡住事件循环,把所有 worker 的请求发出/回调
                # 压成每 interval 一波的突发,既污染客户端计时也放大服务端排队。
                await asyncio.to_thread(self._scan)

        self._task = asyncio.create_task(sample_loop())

    async def finish_level(self, elapsed_sec: float) -> dict[str, Any]:
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        # 收尾再扫一次,补上最后一个不满 tick 的量。
        self._scan()
        report: dict[str, Any] = {}
        for group, acc in self._accum.items():
            user_sec = acc.user_ticks / CLK_TCK
            sys_sec = acc.sys_ticks / CLK_TCK
            report[group] = {
                "cpu_cores_avg": (
                    round((user_sec + sys_sec) / elapsed_sec, 3)
                    if elapsed_sec > 0
                    else None
                ),
                "cpu_user_sec": round(user_sec, 2),
                "cpu_sys_sec": round(sys_sec, 2),
                "rss_mib_avg": (
                    round(acc.rss_sum_mib / acc.samples, 1) if acc.samples else None
                ),
                "rss_mib_peak": round(acc.rss_peak_mib, 1) if acc.samples else None,
                "proc_count_avg": (
                    round(acc.count_sum / acc.samples, 1) if acc.samples else None
                ),
                "proc_count_max": acc.count_max if acc.samples else None,
            }
        return report

    async def close(self) -> None:
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None


@dataclass
class LevelServerStats:
    count_delta: float | None = None
    sum_delta: float | None = None
    max_inflight: float | None = None
    inflight_samples: int = 0
    resume_count_delta: float | None = None
    resume_sum_delta: float | None = None
    resume_max_inflight: float | None = None
    resume_inflight_samples: int = 0
    envd_count_delta: float | None = None
    envd_sum_delta: float | None = None


class MetricsWatcher:
    """Background /metrics sampler capturing per-level server-side stats.

    Uses plain urllib on a worker thread so the probe has no dependency
    beyond the e2b SDK itself.
    """

    def __init__(self, metrics_url: str, interval_sec: float = 1.0):
        self.metrics_url = metrics_url
        self.interval_sec = interval_sec
        self._task: asyncio.Task[None] | None = None
        self._baseline: dict[tuple[str, tuple[tuple[str, str], ...]], float] | None = None
        self._end: dict[tuple[str, tuple[tuple[str, str], ...]], float] | None = None
        self._inflight_peak: float = 0.0
        self._inflight_samples = 0
        self._resume_inflight_peak: float = 0.0
        self._resume_inflight_samples = 0

    async def _fetch(self) -> dict[tuple[str, tuple[tuple[str, str], ...]], float] | None:
        def fetch_blocking() -> str:
            with urllib.request.urlopen(self.metrics_url, timeout=5) as response:
                return response.read().decode("utf-8", errors="replace")

        try:
            text = await asyncio.to_thread(fetch_blocking)
            return parse_prom_text(text)
        except (urllib.error.URLError, OSError, ValueError) as exc:
            LOGGER.warning("metrics scrape failed: %s: %r", type(exc).__name__, exc)
            return None

    async def start_level(self) -> None:
        self._baseline = await self._fetch()
        self._end = None
        self._inflight_peak = 0.0
        self._inflight_samples = 0
        self._resume_inflight_peak = 0.0
        self._resume_inflight_samples = 0

        async def sample_loop() -> None:
            while True:
                await asyncio.sleep(self.interval_sec)
                samples = await self._fetch()
                if samples is None:
                    continue
                value = inflight_value(samples, "pause")
                if value is not None:
                    self._inflight_peak = max(self._inflight_peak, value)
                    self._inflight_samples += 1
                # resume inflight 需要服务端部署 resume stage timer 补丁;
                # 未部署时该系列不存在,保持 None。
                resume_value = inflight_value(samples, "resume")
                if resume_value is not None:
                    self._resume_inflight_peak = max(
                        self._resume_inflight_peak, resume_value
                    )
                    self._resume_inflight_samples += 1

        self._task = asyncio.create_task(sample_loop())

    async def finish_level(self) -> LevelServerStats:
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None
        self._end = await self._fetch()
        stats = LevelServerStats(
            max_inflight=self._inflight_peak if self._inflight_samples else None,
            inflight_samples=self._inflight_samples,
            resume_max_inflight=(
                self._resume_inflight_peak if self._resume_inflight_samples else None
            ),
            resume_inflight_samples=self._resume_inflight_samples,
        )
        if self._baseline is not None and self._end is not None:
            count_start = stage_value(self._baseline, PAUSE_STAGE_METRIC, "_count", "pause")
            count_end = stage_value(self._end, PAUSE_STAGE_METRIC, "_count", "pause")
            sum_start = stage_value(self._baseline, PAUSE_STAGE_METRIC, "_sum", "pause")
            sum_end = stage_value(self._end, PAUSE_STAGE_METRIC, "_sum", "pause")
            if count_start is not None and count_end is not None:
                stats.count_delta = count_end - count_start
            if sum_start is not None and sum_end is not None:
                stats.sum_delta = sum_end - sum_start
            resume_c0, resume_s0 = unmatched_http_totals(self._baseline, "control_plane")
            resume_c1, resume_s1 = unmatched_http_totals(self._end, "control_plane")
            if resume_c1 > resume_c0:
                stats.resume_count_delta = resume_c1 - resume_c0
                stats.resume_sum_delta = resume_s1 - resume_s0
            envd_c0, envd_s0 = unmatched_http_totals(self._baseline, "proxy_header")
            envd_c1, envd_s1 = unmatched_http_totals(self._end, "proxy_header")
            if envd_c1 > envd_c0:
                stats.envd_count_delta = envd_c1 - envd_c0
                stats.envd_sum_delta = envd_s1 - envd_s0
        return stats

    async def wait_pause_drained(self, timeout_sec: float) -> bool:
        """Block until no pause is in flight (level isolation)."""

        deadline = time.monotonic() + timeout_sec
        drained_streak = 0
        while time.monotonic() < deadline:
            samples = await self._fetch()
            value = inflight_value(samples, "pause") if samples is not None else None
            if value is not None and value <= 0:
                drained_streak += 1
                if drained_streak >= 2:
                    return True
            else:
                drained_streak = 0
            await asyncio.sleep(1.0)
        return False

    async def close(self) -> None:
        if self._task is not None:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass


# ---------------------------------------------------------------------------
# Probe: sandbox pool + per-level cycling workers.
# ---------------------------------------------------------------------------


@dataclass
class ProbeOptions:
    template: str
    rounds: int
    think_sec: float
    jitter_sec: float
    command: str | None
    request_timeout_sec: float
    sandbox_timeout_sec: float
    lifecycle_retries: int = 1


@dataclass
class SandboxSlot:
    sandbox: Any
    sandbox_id: str
    dead: bool = False
    pause_samples: list[float] = field(default_factory=list)
    resume_samples: list[float] = field(default_factory=list)
    command_samples: list[float] = field(default_factory=list)
    command_guest_samples: list[float] = field(default_factory=list)
    errors: list[str] = field(default_factory=list)


class PauseResumeProbe:
    def __init__(self, options: ProbeOptions, active_count: int):
        self.options = options
        self.active_count = active_count
        self.slots: list[SandboxSlot] = []

    async def create_pool(self, pool_size: int, metadata_key: str) -> None:
        from e2b import AsyncSandbox  # requires E2B_* env vars set by the caller

        semaphore = asyncio.Semaphore(4)

        async def create_one(index: int) -> SandboxSlot:
            async with semaphore:
                sandbox = await AsyncSandbox.create(
                    template=self.options.template,
                    metadata={
                        "runner": "pause-resume-probe",
                        "probe_key": metadata_key,
                        "slot": str(index),
                    },
                    timeout=self.options.sandbox_timeout_sec,
                    request_timeout=self.options.request_timeout_sec,
                )
            slot = SandboxSlot(sandbox=sandbox, sandbox_id=sandbox.sandbox_id)
            # Warm the pause path once and leave the pool parked in Paused.
            await self._pause(slot)
            return slot

        self.slots = list(await asyncio.gather(*(create_one(i) for i in range(pool_size))))
        LOGGER.info("pool ready: %d sandboxes", len(self.slots))

    async def _pause(self, slot: SandboxSlot) -> float:
        started = time.perf_counter()
        for attempt in range(1, self.options.lifecycle_retries + 2):
            try:
                await slot.sandbox.beta_pause(
                    request_timeout=self.options.request_timeout_sec,
                )
                elapsed = time.perf_counter() - started
                slot.pause_samples.append(elapsed)
                return elapsed
            except Exception as exc:  # noqa: BLE001
                if attempt > self.options.lifecycle_retries:
                    slot.errors.append(f"pause:{type(exc).__name__}")
                    slot.dead = True
                    raise
                LOGGER.warning(
                    "pause retry slot=%s attempt=%d error=%r",
                    slot.sandbox_id, attempt, exc,
                )
        raise RuntimeError("unreachable")

    async def _resume(self, slot: SandboxSlot) -> float:
        # connect() 的返回语义在超时时是二义的:服务端可能已经 resume 成功。
        # 失败重试一次,仍失败才标记死亡——与 replay harness 的处理一致。
        started = time.perf_counter()
        for attempt in range(1, self.options.lifecycle_retries + 2):
            try:
                slot.sandbox = await slot.sandbox.connect(
                    timeout=self.options.sandbox_timeout_sec,
                    request_timeout=self.options.request_timeout_sec,
                )
                elapsed = time.perf_counter() - started
                slot.resume_samples.append(elapsed)
                return elapsed
            except Exception as exc:  # noqa: BLE001
                if attempt > self.options.lifecycle_retries:
                    slot.errors.append(f"resume:{type(exc).__name__}")
                    slot.dead = True
                    raise
                LOGGER.warning(
                    "resume retry slot=%s attempt=%d error=%r",
                    slot.sandbox_id, attempt, exc,
                )
        raise RuntimeError("unreachable")

    async def _worker(self, slot: SandboxSlot, start_gate: asyncio.Event) -> None:
        await start_gate.wait()
        if self.options.jitter_sec > 0:
            await asyncio.sleep(random.uniform(0, self.options.jitter_sec))
        try:
            for _ in range(self.options.rounds):
                await self._resume(slot)
                if self.options.command:
                    # envd 用 /bin/bash -l -c 执行:包一层纳秒计时,guest 内
                    # 纯执行耗时从 stdout 回显标记里解析(退出码保持透传)。
                    wrapped = (
                        "__t0=$(date +%s%N); "
                        f"{self.options.command}; "
                        "__rc=$?; __t1=$(date +%s%N); "
                        'echo "__aenv_probe_ns=$((__t1-__t0))"; exit $__rc'
                    )
                    started = time.perf_counter()
                    result = await slot.sandbox.commands.run(
                        cmd=wrapped,
                        user="root",
                        request_timeout=self.options.request_timeout_sec,
                    )
                    slot.command_samples.append(time.perf_counter() - started)
                    match = PROBE_GUEST_NS_RE.search(result.stdout or "")
                    if match:
                        slot.command_guest_samples.append(
                            int(match.group(1)) / 1e9
                        )
                await self._pause(slot)
                if self.options.think_sec > 0:
                    await asyncio.sleep(self.options.think_sec)
        except Exception as exc:  # noqa: BLE001 - one dead slot must not stop a level
            LOGGER.error("worker aborted slot=%s error=%r", slot.sandbox_id, exc)

    async def run_level(self) -> None:
        active = [slot for slot in self.slots if not slot.dead][: self.active_count]
        for slot in active:
            slot.pause_samples.clear()
            slot.resume_samples.clear()
            slot.command_samples.clear()
            slot.command_guest_samples.clear()
            slot.errors.clear()
        LOGGER.info("level starting: active=%d rounds=%d", len(active), self.options.rounds)
        gate = asyncio.Event()
        workers = [asyncio.create_task(self._worker(slot, gate)) for slot in active]
        # 所有 worker 同时起步:默认(think=0)复现最大突发对齐的到达模式。
        gate.set()
        await asyncio.gather(*workers)
        LOGGER.info("level finished: active=%d", len(active))

    async def kill_pool(self) -> None:
        from e2b import AsyncSandbox

        async def kill_one(slot: SandboxSlot) -> None:
            if slot.sandbox is None:
                return
            try:
                await asyncio.wait_for(
                    slot.sandbox.kill(), timeout=self.options.request_timeout_sec
                )
            except Exception as exc:  # noqa: BLE001
                LOGGER.warning(
                    "kill failed slot=%s error=%r, falling back to class kill",
                    slot.sandbox_id, exc,
                )
                try:
                    await AsyncSandbox.kill(slot.sandbox_id)
                except Exception:  # noqa: BLE001
                    pass

        if self.slots:
            await asyncio.gather(*(kill_one(slot) for slot in self.slots))


# ---------------------------------------------------------------------------
# Orchestration and reporting.
# ---------------------------------------------------------------------------


async def run_sweep(args: argparse.Namespace) -> dict[str, Any]:
    options = ProbeOptions(
        template=args.template,
        rounds=args.rounds,
        think_sec=args.think_sec,
        jitter_sec=args.jitter_sec,
        command=args.command,
        request_timeout_sec=args.request_timeout_sec,
        sandbox_timeout_sec=args.sandbox_timeout_sec,
        lifecycle_retries=args.lifecycle_retries,
    )
    if args.levels.strip().lower() == "all":
        levels = [args.max_pool]
    else:
        levels = [int(item) for item in args.levels.split(",")]
    pool_size = max(levels)
    if pool_size > args.max_pool:
        pool_size = args.max_pool
        levels = [level for level in levels if level <= pool_size]
        LOGGER.warning("pool capped at %d; levels trimmed to %s", pool_size, levels)

    metadata_key = uuid.uuid4().hex[:12]
    probe = PauseResumeProbe(options, active_count=0)
    watcher: MetricsWatcher | None = None
    if not args.no_metrics:
        watcher = MetricsWatcher(args.metrics_url, interval_sec=args.sample_interval_sec)
    proc_sampler: ProcStatsSampler | None = None
    proc_matches: dict[str, tuple[str, ...]] | None = None
    if not args.no_proc_stats:
        api_host = urllib.parse.urlparse(args.api_url).hostname or ""
        if api_host not in ("", "localhost", "127.0.0.1", "::1"):
            LOGGER.warning(
                "API host %r is not local; /proc sampling cannot see the remote "
                "server/ublk/firecracker processes (disable with --no-proc-stats)",
                api_host,
            )
        proc_matches = {
            "server": _split_proc_names(args.proc_match_server),
            "ublk": _split_proc_names(args.proc_match_ublk),
            "firecracker": _split_proc_names(args.proc_match_firecracker),
        }
        proc_sampler = ProcStatsSampler(
            proc_matches,
            interval_sec=args.proc_interval_sec,
            proc_root=args.proc_root,
        )

    results: dict[str, Any] = {
        "started_at": now_iso(),
        "config": {
            "api_url": args.api_url,
            "metrics_url": args.metrics_url,
            "template": args.template,
            "levels": levels,
            "rounds": args.rounds,
            "think_sec": args.think_sec,
            "jitter_sec": args.jitter_sec,
            "command": args.command,
            "request_timeout_sec": args.request_timeout_sec,
            "probe_key": metadata_key,
            "proc_matches": proc_matches,
        },
        "levels": [],
    }

    try:
        await probe.create_pool(pool_size, metadata_key)
        for level in levels:
            if watcher is not None:
                await watcher.start_level()
            if proc_sampler is not None:
                proc_sampler.start_level()
            probe.active_count = level
            level_started = time.perf_counter()
            await probe.run_level()
            elapsed = time.perf_counter() - level_started
            server_stats = (
                await watcher.finish_level() if watcher is not None else LevelServerStats()
            )
            proc_stats = (
                await proc_sampler.finish_level(elapsed)
                if proc_sampler is not None
                else None
            )

            active = [s for s in probe.slots if not s.dead][:level]
            pauses = [v for slot in active for v in slot.pause_samples]
            resumes = [v for slot in active for v in slot.resume_samples]
            command_runs = [v for slot in active for v in slot.command_samples]
            command_guest = [v for slot in active for v in slot.command_guest_samples]
            errors = [err for slot in active for err in slot.errors]
            client_pause_ops = len(pauses) + len(errors)
            client_rate = client_pause_ops / elapsed if elapsed > 0 else None
            server_rate = (
                server_stats.count_delta / elapsed
                if server_stats.count_delta is not None and elapsed > 0
                else None
            )
            server_avg = (
                server_stats.sum_delta / server_stats.count_delta
                if server_stats.sum_delta is not None and server_stats.count_delta
                else None
            )
            server_resume_avg = (
                server_stats.resume_sum_delta / server_stats.resume_count_delta
                if server_stats.resume_sum_delta is not None
                and server_stats.resume_count_delta
                else None
            )
            server_envd_avg = (
                server_stats.envd_sum_delta / server_stats.envd_count_delta
                if server_stats.envd_sum_delta is not None and server_stats.envd_count_delta
                else None
            )
            server_resume_rate = (
                server_stats.resume_count_delta / elapsed
                if server_stats.resume_count_delta is not None and elapsed > 0
                else None
            )
            little_estimate = (
                server_stats.max_inflight / server_rate
                if server_stats.max_inflight is not None and server_rate
                else None
            )

            level_result = {
                "level": level,
                "active": len(active),
                "dead_slots": sum(1 for s in probe.slots if s.dead),
                "elapsed_sec": round(elapsed, 3),
                "client_pause": summarize(pauses),
                "client_resume": summarize(resumes),
                "client_command": summarize(command_runs),
                "client_command_guest": summarize(command_guest),
                "errors": errors,
                "client_pause_ops_per_sec": round(client_rate, 3) if client_rate else None,
                "server_pause_count_delta": server_stats.count_delta,
                "server_pause_ops_per_sec": round(server_rate, 3) if server_rate else None,
                "server_pause_avg_sec": round(server_avg, 4) if server_avg else None,
                "server_pause_max_inflight": server_stats.max_inflight,
                "server_resume_count_delta": server_stats.resume_count_delta,
                "server_resume_ops_per_sec": (
                    round(server_resume_rate, 3) if server_resume_rate else None
                ),
                "server_resume_avg_sec": (
                    round(server_resume_avg, 4) if server_resume_avg else None
                ),
                "server_resume_max_inflight": server_stats.resume_max_inflight,
                "server_envd_count_delta": server_stats.envd_count_delta,
                "server_envd_avg_sec": (
                    round(server_envd_avg, 4) if server_envd_avg else None
                ),
                "proc": proc_stats,
                "little_estimate_sec": round(little_estimate, 3) if little_estimate else None,
            }
            results["levels"].append(level_result)
            print_level(level_result)

            # 级间隔离:等 pause in-flight 清零再进入下一级,避免上一级队列
            # 污染下一级测量。
            if watcher is not None and level != levels[-1]:
                drained = await watcher.wait_pause_drained(timeout_sec=args.cooldown_sec)
                if not drained:
                    LOGGER.warning("pause inflight did not drain within %ss", args.cooldown_sec)
                await asyncio.sleep(2.0)
        results["completed_at"] = now_iso()
    finally:
        if not args.keep_sandboxes:
            LOGGER.info("cleaning up pool (%d sandboxes)", len(probe.slots))
            await probe.kill_pool()
        if watcher is not None:
            await watcher.close()
        if proc_sampler is not None:
            await proc_sampler.close()
    return results


def print_level(result: dict[str, Any]) -> None:
    pause = result["client_pause"]
    resume = result["client_resume"]
    command = result["client_command"]
    guest = result["client_command_guest"]
    print(
        f"\n[ level {result['level']:>3} ] active={result['active']} "
        f"errors={len(result['errors'])} elapsed={result['elapsed_sec']}s\n"
        f"  client pause : n={pause['n']} p50={pause['p50']} p95={pause['p95']} "
        f"avg={pause['avg']} max={pause['max']}\n"
        f"  client resume: n={resume['n']} p50={resume['p50']} p95={resume['p95']} "
        f"avg={resume['avg']} max={resume['max']}\n"
        f"  client cmd   : n={command['n']} p50={command['p50']} p95={command['p95']} "
        f"avg={command['avg']} max={command['max']}\n"
        f"  guest cmd    : n={guest['n']} p50={guest['p50']} p95={guest['p95']} "
        f"avg={guest['avg']} max={guest['max']}\n"
        f"  client ops/s : {result['client_pause_ops_per_sec']}   "
        f"server ops/s: {result['server_pause_ops_per_sec']}\n"
        f"  server avg  : {result['server_pause_avg_sec']}s   "
        f"max inflight: {result['server_pause_max_inflight']}   "
        f"little est: {result['little_estimate_sec']}s\n"
        f"  server rsm  : avg={result['server_resume_avg_sec']}s "
        f"ops/s={result['server_resume_ops_per_sec']} "
        f"inflight={result['server_resume_max_inflight']} "
        f"n={result['server_resume_count_delta']}   "
        f"server envd: avg={result['server_envd_avg_sec']}s "
        f"n={result['server_envd_count_delta']}",
        flush=True,
    )
    proc = result.get("proc") or {}
    proc_lines = []
    for group in ("server", "ublk", "firecracker"):
        stats = proc.get(group)
        if not stats:
            continue
        proc_lines.append(
            f"  proc {group:<11}: cpu={stats['cpu_cores_avg']} cores "
            f"(u={stats['cpu_user_sec']}s s={stats['cpu_sys_sec']}s) "
            f"rss avg={stats['rss_mib_avg']} peak={stats['rss_mib_peak']} MiB "
            f"procs avg={stats['proc_count_avg']} max={stats['proc_count_max']}"
        )
    if proc_lines:
        print("\n".join(proc_lines), flush=True)


# ---------------------------------------------------------------------------
# CSV 报告:每个 level 一行,嵌套字段拍平成列,便于表格工具直接分析。
# ---------------------------------------------------------------------------


def flatten_level(result: dict[str, Any], prefix: str = "") -> dict[str, Any]:
    flat: dict[str, Any] = {}
    for key, value in result.items():
        name = f"{prefix}_{key}" if prefix else key
        if isinstance(value, dict):
            flat.update(flatten_level(value, name))
        elif isinstance(value, list):
            flat[f"{name}_count"] = len(value)
        else:
            flat[name] = value
    return flat


def write_levels_csv(results: dict[str, Any], path: Path) -> None:
    rows = [flatten_level(entry) for entry in results["levels"]]
    fieldnames: list[str] = []
    for row in rows:
        for key in row:
            if key not in fieldnames:
                fieldnames.append(key)
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(fieldnames)
        for row in rows:
            writer.writerow(
                ["" if row.get(name) is None else row.get(name) for name in fieldnames]
            )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Measure pause/resume latency vs concurrent cycling sandboxes.",
    )
    parser.add_argument(
        "--api-url",
        default=os.getenv("E2B_API_URL", "http://127.0.0.1:8000"),
        help="AgentENV E2B-compatible API URL",
    )
    parser.add_argument(
        "--metrics-url",
        default=None,
        help="Prometheus /metrics URL (default: <api-url>/metrics)",
    )
    parser.add_argument("--template", required=True)
    parser.add_argument(
        "--levels",
        default="1,2,4,8",
        help="Comma-separated concurrent cycling sandbox counts to sweep, "
        "or 'all' for a single run with the entire pool cycling at once",
    )
    parser.add_argument(
        "--rounds", type=int, default=8, help="pause/resume cycles per sandbox per level"
    )
    parser.add_argument(
        "--think-sec", type=float, default=0.0,
        help="Dwell after each pause before the next resume (trajectory delay model)",
    )
    parser.add_argument(
        "--jitter-sec", type=float, default=0.0,
        help="Random per-worker start offset; 0 keeps bursts fully aligned",
    )
    parser.add_argument("--command", default=None, help="Optional command between resume and pause")
    parser.add_argument("--request-timeout-sec", type=float, default=300.0)
    parser.add_argument("--sandbox-timeout-sec", type=int, default=86_400)
    parser.add_argument("--lifecycle-retries", type=int, default=1)
    parser.add_argument("--max-pool", type=int, default=64)
    parser.add_argument("--sample-interval-sec", type=float, default=1.0)
    parser.add_argument("--cooldown-sec", type=float, default=120.0)
    parser.add_argument("--no-metrics", action="store_true", help="Skip /metrics sampling")
    parser.add_argument(
        "--no-proc-stats",
        action="store_true",
        help="Skip local /proc CPU/RSS sampling of the server, ublk daemon and "
        "firecracker processes (use when the probe is not on the server host)",
    )
    parser.add_argument(
        "--proc-root",
        default="/proc",
        help="procfs mount used for process stats (e.g. a host /proc bind-mount)",
    )
    parser.add_argument(
        "--proc-interval-sec",
        type=float,
        default=0.25,
        help="/proc polling interval; bounds the CPU lost when a process exits "
        "mid-interval (firecracker exits on every pause)",
    )
    parser.add_argument(
        "--proc-match-server",
        default="server,agentenv-server",
        help="Comma-separated process comm names for the AgentENV server",
    )
    parser.add_argument(
        "--proc-match-ublk",
        default="uvm-ublk-daemon",
        help="Comma-separated process comm names for the ublk daemon",
    )
    parser.add_argument(
        "--proc-match-firecracker",
        default="firecracker,firecracker-nex",
        help="Comma-separated process comm names for firecracker VM processes "
        "(PVM builds are firecracker-next, truncated to firecracker-nex in comm)",
    )
    parser.add_argument("--keep-sandboxes", action="store_true")
    parser.add_argument(
        "--output-dir",
        default="runs",
        help="Directory for default run outputs, relative to the working directory",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="Explicit JSON results path (a per-level CSV table is written "
        "alongside); default: <output-dir>/pause-resume-<timestamp>.json",
    )
    parser.add_argument("--verbose", action="store_true")
    return parser


def main() -> int:
    args = build_parser().parse_args()
    if args.metrics_url is None:
        args.metrics_url = args.api_url.rstrip("/") + "/metrics"
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    # 传输层会在 INFO 级打印每条 HTTP 请求(e2b>=2.46 是 pyqwest,<=2.37 是
    # httpx),刷屏且与测量无关,压到 WARNING。
    logging.getLogger("pyqwest").setLevel(logging.WARNING)
    logging.getLogger("httpx").setLevel(logging.WARNING)
    try:
        results = asyncio.run(run_sweep(args))
    except KeyboardInterrupt:
        LOGGER.warning("interrupted")
        return 130
    if args.output:
        output_path = Path(args.output).expanduser().resolve()
    else:
        stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
        output_path = Path(args.output_dir) / f"pause-resume-{stamp}.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(results, ensure_ascii=False, indent=2))
    csv_path = output_path.with_suffix(".csv")
    write_levels_csv(results, csv_path)
    print(f"\nresults written to {output_path}", flush=True)
    print(f"csv written to {csv_path}", flush=True)
    return 0


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