codex-app-server/cloud/chv/resident_pool.py

"""每用户常驻 VM 池（生产模块，P1 对齐 Muse，见 ADR-17）。

每用户一台常驻 VM：任务到时 resume 已暂停的 VM，跑完
再 pause，不销毁；池空或任何一步失败都回落冷启动，
绝不因池坏而丢任务。常驻 VM 的盘是每用户持久工作区盘
（state_dir/chv-resident-<用户>/disk.raw），不是每任务
新盘——隔离单元在常驻路径上从“每任务”变成“每用户”
（单租户可接受，多租户重审前不得把常驻当租户隔离）。

状态机：cold → provisioning → paused → running → paused。
acquire(user)：有 paused 条目时 resume_vm + 等 SSH 就绪
（复用执行器 wait_guest_ready，不另写判定），返回
mode=resident；无可用/失败时返回 mode=cold。release(user)：
pause_vm，不销毁；pause 失败时条目留在 running（不删
不销毁，worker 关闭时 reconcile 清理）。provision(user)：
wake_vm + pause_vm 建 paused 条目（惰性建池，失败回落）。
reclaim_idle()：paused 超 CHV_RESIDENT_IDLE_MINUTES（默认
0=不自动关机）时 hibernate（pause + snapshot + shutdown）
并删条目，快照文件留作恢复物，restore 是后续工作。

Registry（state_dir/resident-pool.json）：文件 IO 与损坏
语义复用 cloud.chv.orphans（原子写、损坏 quarantine、
单条目非法只丢该条）；重启后读到的条目一律 orphan，不
复用，只能 reconcile 清理。条目按用户键隔离：release
绝不把用过的 VM 放回给另一用户；同一用户同一时间只跑
一个常驻任务，并发第二个任务回落冷启动。

诚实边界：resume 失败时 wait_guest_ready 的失败路径会
stop_vm（含删工作区盘），条目随之丢弃——该路径的工作区
以 Gitea 同步为准；带密钥 extra_env 的任务不走常驻
（seed 是建池时定的，密钥只能冷启动随 seed 投递）。
"""

import logging
import os
import re
import threading
import time

from cloud.chv.orphans import _read_registry, _write_registry

_LOGGER = logging.getLogger("cloud.cloud_hypervisor")

_REGISTRY_FILENAME = "resident-pool.json"
_RESIDENT_STATES = ("paused", "running", "provisioning")
_DEFAULT_USER = "default"


def resident_user_key(user_id: str | None) -> str:
    """用户键：非空即原样消毒（非法字符变 -），空/缺省
    即 "default"（当前单租户：任务 dict 无 user_id 时全
    部落到同一常驻 VM；多租户重审前不得当租户边界用）。"""
    raw = _DEFAULT_USER if user_id is None else str(user_id)
    key = re.sub(r"[^A-Za-z0-9_.-]", "-", raw).strip(".-")
    return key or _DEFAULT_USER


def resident_task_for_user(user_id: str | None) -> dict:
    """常驻伪任务：执行器按它算 VM 名/known_hosts/socket。
    vm_name({"id": "resident-alice"}) = "chv-resident-alice"。"""
    return {"id": "resident-" + resident_user_key(user_id)}


def resident_registry_path_for_executor(executor: object) -> str | None:
    """常驻 registry 路径：state_dir/resident-pool.json。
    无 state_dir 时返回 None（调用方按“无残留”处理）。"""
    state_dir = getattr(executor, "state_dir", None)
    if not state_dir:
        return None
    return os.path.join(state_dir, _REGISTRY_FILENAME)


def reconcile_resident_orphans_for_executor(executor: object) -> list[str]:
    """生产路径的常驻残留清理入口：返回已清理的 VM 名。
    只做清理：读 resident-pool.json → 每个上一世条目尽力
    shutdown_vm → 落盘为空 → best-effort 审计。绝不构造
    ResidentPool、绝不把任何条目发给任务。无 registry 路
    径时返回 []。文件 IO 与损坏语义复用 orphans。"""
    path = resident_registry_path_for_executor(executor)
    if not path:
        return []
    data, _corrupt = _read_registry(path)
    if data is None:
        return []
    cleaned: list[str] = []
    for entry in data.values():
        if not isinstance(entry, dict):
            continue
        if not entry.get("vm"):
            continue
        if entry.get("state") not in _RESIDENT_STATES:
            continue
        vm = str(entry["vm"])
        try:
            executor.shutdown_vm(vm)  # type: ignore[attr-defined]
        except Exception:
            pass
        cleaned.append(vm)
    if cleaned:
        _write_registry(path, {})
        try:
            executor._emit(  # type: ignore[attr-defined]
                "resident_pool_reconcile", {"id": None},
                "cleaned=%s" % ",".join(cleaned))
        except Exception:
            pass
    return cleaned


class ResidentPool:
    """每用户常驻 VM 池（生产路径，见模块 docstring）。

    同一 state_dir 只活一个池实例：registry 只在构造时
    读一次，重启后条目一律 orphan（见 _load_registry）。
    公共方法全部加锁：并发 acquire/release/provision 只
    放行一路，其余回 cold/False，不把同一 VM 发给两个任务。
    """

    def __init__(self, executor: object, idle_minutes: int | None = None,
                 registry_path: str | None = None) -> None:
        """建池并读 registry：读到的条目一律 orphan（见模块
        docstring）。idle_minutes 未给时读执行器属性或
        CHV_RESIDENT_IDLE_MINUTES，非法/负数即 0（不自动关机）。"""
        self.executor = executor
        self._lock = threading.Lock()
        self._busy: set[str] = set()
        self._entries: dict = {}
        self.registry_path = registry_path
        if self.registry_path is None:
            self.registry_path = \
                resident_registry_path_for_executor(executor)
        if idle_minutes is None:
            idle_minutes = self._idle_minutes_from_env()
        self.idle_minutes = idle_minutes if idle_minutes > 0 else 0
        self.registry_corrupt = False
        self.registry_corrupt_entries = 0
        self._load_registry()

    # -- registry ------------------------------------------------------
    def _load_registry(self) -> None:
        """载入 registry：读到的条目一律标 orphan（绝不直接
        当 paused 发给任务，旧 VMM 进程已随旧 worker 消失）。
        无文件时为空池；损坏时 quarantine 并从空开始；单条
        目非法只丢该条并计数，不连坐其它条目。"""
        if not self.registry_path:
            return
        data, corrupt = _read_registry(self.registry_path)
        if corrupt:
            self.registry_corrupt = True
            return
        if data is None:
            return
        for key, entry in data.items():
            if not isinstance(entry, dict) \
                    or not entry.get("vm") \
                    or entry.get("state") not in _RESIDENT_STATES:
                self.registry_corrupt_entries += 1
                continue
            self._entries[key] = {
                "state": "orphan",
                "vm": entry["vm"],
                "user_id": entry.get("user_id") or key,
                "paused_at": None,
                "last_used": None,
            }

    def _save_registry(self) -> None:
        """落盘当前条目（只写 paused/running/provisioning；
        orphan 是待清理的上一世条目，reconcile 后即删，不写
        回）。调用方须已持锁；写失败只记日志，不炸任务流。"""
        if not self.registry_path:
            return
        data = {}
        for key, entry in self._entries.items():
            if entry["state"] not in _RESIDENT_STATES:
                continue
            data[key] = {
                "state": entry["state"],
                "vm": entry.get("vm"),
                "user_id": entry.get("user_id"),
                "paused_at": entry.get("paused_at"),
                "last_used": entry.get("last_used"),
            }
        _write_registry(self.registry_path, data)

    def _emit(self, event: str, task: dict, detail: str) -> None:
        """best-effort 审计：执行器无 _emit 或抛错都不炸
        池操作（与 WarmPool 同款）。"""
        emit = getattr(self.executor, "_emit", None)
        if not callable(emit):
            return
        try:
            emit(event, task, detail)
        except Exception:
            pass

    def _idle_minutes_from_env(self) -> int:
        """闲置阈值：执行器属性优先（P1-B 若经 config 落
        属性，直接复用），否则读 CHV_RESIDENT_IDLE_MINUTES；
        非数字/负数一律 0（不自动关机，不炸建池）。"""
        raw = getattr(self.executor, "resident_idle_minutes", None)
        if raw is None:
            env = getattr(self.executor, "_env", None)
            if env is None:
                env = os.environ
            raw = env.get("CHV_RESIDENT_IDLE_MINUTES", 0)
        if raw is None:
            return 0
        try:
            minutes = int(raw)
        except (TypeError, ValueError):
            return 0
        return minutes if minutes > 0 else 0

    def _vm_name(self, resident_task: dict) -> str:
        """常驻 VM 名：优先执行器 vm_name（与冷启动同款派
        生），鸭子型执行器无该方法时退回 chv- 前缀拼法。"""
        vm_name = getattr(self.executor, "vm_name", None)
        if callable(vm_name):
            try:
                return str(vm_name(resident_task))
            except Exception:
                pass
        return "chv-" + str(resident_task.get("id"))

    # -- 计数 ----------------------------------------------------------
    def paused_count(self) -> int:
        """当前 paused（可分配）的常驻条目数。"""
        with self._lock:
            return sum(1 for e in self._entries.values()
                       if e.get("state") == "paused")

    def orphan_count(self) -> int:
        """上一世残留条目数（待 reconcile，不可分配）。"""
        with self._lock:
            return sum(1 for e in self._entries.values()
                       if e.get("state") == "orphan")

    # -- acquire / release ----------------------------------------------
    def _claim_paused(self, user_key: str) -> dict | None:
        """占一条 paused 条目：加忙标记并返回其快照；无条目/
        非 paused/正忙时返回 None（调用方回落冷启动）。"""
        with self._lock:
            entry = self._entries.get(user_key)
            if entry is None or entry.get("state") != "paused":
                return None
            if user_key in self._busy:
                return None
            self._busy.add(user_key)
            return dict(entry)

    def _resume_and_wait(self, resident_task: dict) -> bool:
        """resume + 等 SSH 就绪（复用执行器判定）。
        resume 缺失的鸭子型执行器跳过 resume（与 WarmPool
        同款宽容）；SSH 未就绪即 False（wait_guest_ready
        的失败路径已自行 stop_vm，调用方只须删条目）。"""
        resume = getattr(self.executor, "resume_vm", None)
        if callable(resume):
            try:
                resume(resident_task)
            except Exception:
                return False
        wait = getattr(self.executor, "wait_guest_ready", None)
        if not callable(wait):
            return True
        try:
            return bool(wait(resident_task))
        except Exception:
            return False

    def _settle_acquire(self, user_key: str, ok: bool) -> None:
        """acquire 收尾：成功标 running，失败删条目；都清
        忙标记并落盘。"""
        with self._lock:
            self._busy.discard(user_key)
            entry = self._entries.get(user_key)
            if entry is None:
                return
            if ok:
                entry["state"] = "running"
                entry["last_used"] = time.time()
            else:
                self._entries.pop(user_key, None)
            self._save_registry()

    def acquire(self, user_id: str | None) -> dict:
        """取常驻 VM：paused 条目 → resume + SSH 就绪 →
        mode=resident；无条目/正忙/失败 → mode=cold（调用
        方回落冷启动）。失败条目直接丢弃（VM 本身已被
        wait_guest_ready 的失败路径关掉，或 resume 已拒）。"""
        user_key = resident_user_key(user_id)
        entry = self._claim_paused(user_key)
        if entry is None:
            return {"mode": "cold", "vm": None, "task": None}
        resident_task = resident_task_for_user(user_id)
        ok = self._resume_and_wait(resident_task)
        self._settle_acquire(user_key, ok)
        if not ok:
            return {"mode": "cold", "vm": None, "task": None}
        self._emit("resident_pool_acquire", resident_task,
                   "vm=%s user=%s" % (entry.get("vm"), user_key))
        return {"mode": "resident", "vm": entry.get("vm"),
                "task": resident_task}

    def _pause(self, resident_task: dict) -> bool:
        """pause_vm 尽力而为：失败返回 False，不抛。"""
        pause = getattr(self.executor, "pause_vm", None)
        if not callable(pause):
            return False
        try:
            pause(resident_task)
            return True
        except Exception:
            return False

    def release(self, user_id: str | None) -> bool:
        """还常驻 VM：running 条目 → pause_vm，不销毁。
        无条目/非 running/正忙 → False；pause 失败时条目
        留在 running（不删不销毁，worker 关闭时 reconcile
        清理），返回 False。只动本用户的条目。"""
        user_key = resident_user_key(user_id)
        resident_task = resident_task_for_user(user_id)
        with self._lock:
            entry = self._entries.get(user_key)
            if entry is None or entry.get("state") != "running":
                return False
            if user_key in self._busy:
                return False
            self._busy.add(user_key)
        paused = self._pause(resident_task)
        with self._lock:
            self._busy.discard(user_key)
            entry = self._entries.get(user_key)
            if entry is None or not paused:
                return False
            entry["state"] = "paused"
            entry["paused_at"] = time.time()
            self._save_registry()
        self._emit("resident_pool_release", resident_task,
                   "vm=%s user=%s" % (entry.get("vm"), user_key))
        return True

    # -- provision -------------------------------------------------------
    def _wake_and_pause(self, resident_task: dict) -> bool:
        """建池：wake_vm（冷启动全套，含 SSH 就绪）再
        pause_vm。任一步失败都尽力 shutdown，不留半醒的
        VM；调用方随后删条目回落冷启动。"""
        try:
            wake = getattr(self.executor, "wake_vm", None)
            pause = getattr(self.executor, "pause_vm", None)
            if not callable(wake) or not callable(pause):
                return False
            wake(resident_task)
            pause(resident_task)
            return True
        except Exception:
            shutdown = getattr(self.executor, "shutdown_vm", None)
            if callable(shutdown):
                try:
                    shutdown(resident_task)
                except Exception:
                    pass
            return False

    def _finish_provision(self, user_key: str, ok: bool) -> str | None:
        """建池收尾：成功标 paused 并记时间，失败删条目；
        都清忙标记并落盘。返回 VM 名或 None。"""
        with self._lock:
            self._busy.discard(user_key)
            entry = self._entries.get(user_key)
            if entry is None:
                return None
            if not ok:
                self._entries.pop(user_key, None)
                self._save_registry()
                return None
            entry["state"] = "paused"
            entry["paused_at"] = time.time()
            self._save_registry()
            return str(entry.get("vm"))

    def provision(self, user_id: str | None) -> dict:
        """建常驻 VM：已有 paused/running 条目直接返回（不
        重 wake）；否则 wake + pause，成功落 paused 条目，
        失败删条目回 cold。并发建池只放行一路，其余回 cold。"""
        user_key = resident_user_key(user_id)
        resident_task = resident_task_for_user(user_id)
        with self._lock:
            entry = self._entries.get(user_key)
            if entry is not None and entry.get("state") in (
                    "paused", "running"):
                return {"mode": "resident", "vm": entry.get("vm"),
                        "task": resident_task}
            if user_key in self._busy:
                return {"mode": "cold", "vm": None, "task": None}
            self._busy.add(user_key)
            self._entries[user_key] = {
                "state": "provisioning",
                "vm": self._vm_name(resident_task),
                "user_id": user_id if user_id is not None
                else _DEFAULT_USER,
                "paused_at": None, "last_used": None}
            self._save_registry()
        ok = self._wake_and_pause(resident_task)
        vm = self._finish_provision(user_key, ok)
        if vm is None:
            return {"mode": "cold", "vm": None, "task": None}
        self._emit("resident_pool_provision", resident_task,
                   "vm=%s user=%s" % (vm, user_key))
        return {"mode": "resident", "vm": vm, "task": resident_task}

    # -- 空闲回收 --------------------------------------------------------
    def _hibernate_one(self, user_key: str) -> str | None:
        """单条目休眠：hibernate_vm（快照到 state_dir 内
        resident.snapshot，成功后删条目）。失败留条目下次
        再试，返回 None；无 state_dir 或无 hibernate 入口
        同样留条目。"""
        with self._lock:
            entry = self._entries.get(user_key)
            if entry is None or entry.get("state") != "paused":
                return None
            if user_key in self._busy:
                return None
            self._busy.add(user_key)
            resident_task = resident_task_for_user(
                entry.get("user_id"))
            vm = str(entry.get("vm"))
        state_dir = getattr(self.executor, "state_dir", None)
        hibernate = getattr(self.executor, "hibernate_vm", None)
        if not state_dir or not callable(hibernate):
            with self._lock:
                self._busy.discard(user_key)
            return None
        snapshot_url = "file://" + os.path.join(
            state_dir, vm, "resident.snapshot")
        try:
            hibernate(resident_task, snapshot_url)
        except Exception:
            with self._lock:
                self._busy.discard(user_key)
            return None
        with self._lock:
            self._busy.discard(user_key)
            self._entries.pop(user_key, None)
            self._save_registry()
        self._emit("resident_pool_reclaim", resident_task,
                   "vm=%s snapshot=%s" % (vm, snapshot_url))
        return vm

    def reclaim_idle(self, now: float | None = None) -> list[str]:
        """空闲回收：paused 超阈值的条目 hibernate（pause +
        snapshot + shutdown）并删条目。阈值 0 即关闭回收；
        未到阈值/正忙/无 paused_at 的条目不动。返回回收的
        VM 名。"""
        if self.idle_minutes <= 0:
            return []
        if now is None:
            now = time.time()
        due: list[str] = []
        with self._lock:
            for key, entry in self._entries.items():
                if entry.get("state") != "paused":
                    continue
                if key in self._busy:
                    continue
                paused_at = entry.get("paused_at")
                if paused_at is None:
                    continue
                if now - paused_at > self.idle_minutes * 60:
                    due.append(key)
        reclaimed: list[str] = []
        for key in due:
            vm = self._hibernate_one(key)
            if vm is not None:
                reclaimed.append(vm)
        return reclaimed

    # -- 孤儿清理 --------------------------------------------------------
    def _shutdown_one(self, user_key: str) -> str | None:
        """单条目尽力 shutdown 并从内存删掉（失败仍计入
        已清理，与 orphans 旧语义一致：不因单条卡住整批）。
        返回 VM 名或 None（条目已不在）。"""
        with self._lock:
            entry = self._entries.get(user_key)
            if entry is None:
                return None
            vm = str(entry.get("vm"))
        shutdown = getattr(self.executor, "shutdown_vm", None)
        if callable(shutdown):
            try:
                shutdown(vm)
            except Exception:
                pass
        with self._lock:
            self._entries.pop(user_key, None)
        return vm

    def reconcile_orphans(self) -> list[str]:
        """清理本池内存里的 orphan 条目（重启后读到的上一
        世条目）：每个尽力 shutdown_vm，条目清掉并落盘。
        生产启动/关闭不走本方法，走模块级 facade。"""
        with self._lock:
            orphans = [k for k, e in self._entries.items()
                       if e.get("state") == "orphan"]
        cleaned: list[str] = []
        for key in orphans:
            vm = self._shutdown_one(key)
            if vm is not None:
                cleaned.append(vm)
        if cleaned:
            with self._lock:
                self._save_registry()
            self._emit("resident_pool_reconcile", {"id": None},
                       "cleaned=%s" % ",".join(cleaned))
        return cleaned