269 lines
12 KiB
Python
269 lines
12 KiB
Python
"""宿主级 Sandbox 重型执行容量。
|
|
|
|
状态保存在 workspace/.sandbox 下并由 advisory file lock 串行化，因此蓝绿和多
|
|
Web 进程共享同一组槽位。业务 DB 不承载实时状态；后台排队事实源仍是
|
|
``.zcbot_procs``，这里只记录已经获得槽位的租约和短暂的前台候选。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
import uuid
|
|
from contextlib import contextmanager, suppress
|
|
from pathlib import Path
|
|
from typing import Callable, Dict, Iterator, Optional
|
|
|
|
from core.file_store import atomic_write_text, interprocess_file_lock
|
|
|
|
DEFAULT_MAX_ACTIVE_EXECS = 6
|
|
DEFAULT_MAX_BACKGROUND_EXECS = 4
|
|
DEFAULT_MAX_ACTIVE_EXECS_PER_USER = 3
|
|
DEFAULT_MIN_MEM_AVAILABLE_BYTES = 1024 ** 3
|
|
|
|
|
|
def _parse_bytes(value: object, default: int) -> int:
|
|
text = str(value or "").strip().lower()
|
|
if not text:
|
|
return default
|
|
units = {"k": 1024, "kb": 1024, "m": 1024**2, "mb": 1024**2,
|
|
"g": 1024**3, "gb": 1024**3}
|
|
for suffix, factor in sorted(units.items(), key=lambda x: -len(x[0])):
|
|
if text.endswith(suffix):
|
|
return int(float(text[:-len(suffix)]) * factor)
|
|
return int(text)
|
|
|
|
|
|
def mem_available_bytes() -> Optional[int]:
|
|
try:
|
|
for line in Path("/proc/meminfo").read_text(encoding="ascii").splitlines():
|
|
if line.startswith("MemAvailable:"):
|
|
return int(line.split()[1]) * 1024
|
|
except (OSError, ValueError, IndexError):
|
|
return None
|
|
return None
|
|
|
|
|
|
class ExecCapacity:
|
|
def __init__(self, state_dir: Path, cfg: Optional[dict] = None) -> None:
|
|
cfg = cfg or {}
|
|
self.state_dir = Path(state_dir)
|
|
self.state_path = self.state_dir / "exec-capacity.json"
|
|
self.lock_path = self.state_dir / "exec-capacity.lock"
|
|
self.max_active = max(1, min(DEFAULT_MAX_ACTIVE_EXECS, int(os.getenv("ZCBOT_MAX_ACTIVE_EXECS") or cfg.get("max_active_execs") or DEFAULT_MAX_ACTIVE_EXECS)))
|
|
self.max_background = max(1, min(DEFAULT_MAX_BACKGROUND_EXECS, int(os.getenv("ZCBOT_MAX_BACKGROUND_EXECS") or cfg.get("max_background_execs") or DEFAULT_MAX_BACKGROUND_EXECS)))
|
|
self.max_per_user = max(1, min(DEFAULT_MAX_ACTIVE_EXECS_PER_USER, int(os.getenv("ZCBOT_MAX_ACTIVE_EXECS_PER_USER") or cfg.get("max_active_execs_per_user") or DEFAULT_MAX_ACTIVE_EXECS_PER_USER)))
|
|
self.min_mem_available = _parse_bytes(
|
|
os.getenv("ZCBOT_MIN_MEM_AVAILABLE") or cfg.get("min_mem_available"),
|
|
DEFAULT_MIN_MEM_AVAILABLE_BYTES,
|
|
)
|
|
|
|
def _read(self) -> dict:
|
|
try:
|
|
data = json.loads(self.state_path.read_text(encoding="utf-8"))
|
|
if isinstance(data, dict):
|
|
data.setdefault("leases", {})
|
|
data.setdefault("foreground_queue", [])
|
|
data.setdefault("containers", {})
|
|
return data
|
|
except (OSError, ValueError):
|
|
pass
|
|
return {"leases": {}, "foreground_queue": [], "containers": {}}
|
|
|
|
def _write(self, state: dict) -> None:
|
|
atomic_write_text(self.state_path, json.dumps(state, ensure_ascii=False, sort_keys=True))
|
|
|
|
def _prune(self, state: dict) -> None:
|
|
leases = state["leases"]
|
|
for key, lease in list(leases.items()):
|
|
if lease.get("kind") == "foreground":
|
|
pid = int(lease.get("owner_pid") or 0)
|
|
if pid and not self._pid_alive(pid):
|
|
leases.pop(key, None)
|
|
state["foreground_queue"] = [
|
|
q for q in state["foreground_queue"]
|
|
if self._pid_alive(int(q.get("owner_pid") or 0))
|
|
]
|
|
|
|
@staticmethod
|
|
def _pid_alive(pid: int) -> bool:
|
|
try:
|
|
os.kill(pid, 0)
|
|
return True
|
|
except PermissionError:
|
|
return True
|
|
except OSError:
|
|
return False
|
|
|
|
def _can_admit(self, state: dict, user_id: str, kind: str) -> bool:
|
|
leases = list(state["leases"].values())
|
|
if len(leases) >= self.max_active:
|
|
return False
|
|
if sum(1 for x in leases if x.get("user_id") == user_id) >= self.max_per_user:
|
|
return False
|
|
if kind == "background" and sum(1 for x in leases if x.get("kind") == "background") >= self.max_background:
|
|
return False
|
|
available = mem_available_bytes()
|
|
return available is None or available >= self.min_mem_available
|
|
|
|
def try_acquire(self, user_id: str, kind: str, *, lease_id: Optional[str] = None, proc_id: Optional[str] = None) -> Optional[str]:
|
|
lease_id = lease_id or uuid.uuid4().hex
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read()
|
|
self._prune(state)
|
|
if lease_id in state["leases"]:
|
|
return None
|
|
if not self._can_admit(state, str(user_id), kind):
|
|
self._write(state)
|
|
return None
|
|
state["leases"][lease_id] = {
|
|
"lease_id": lease_id, "user_id": str(user_id), "kind": kind,
|
|
"proc_id": proc_id, "owner_pid": os.getpid(), "started_ts": time.time(),
|
|
}
|
|
self._write(state)
|
|
return lease_id
|
|
|
|
def acquire_foreground(
|
|
self,
|
|
user_id: str,
|
|
cancel_check: Optional[Callable[[], bool]] = None,
|
|
wait_notify: Optional[Callable[[dict], None]] = None,
|
|
) -> Optional[str]:
|
|
ticket = uuid.uuid4().hex
|
|
wait_notified = False
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read(); self._prune(state)
|
|
state["foreground_queue"].append({"ticket": ticket, "user_id": str(user_id), "created_ts": time.time(), "owner_pid": os.getpid()})
|
|
self._write(state)
|
|
while True:
|
|
admitted = False
|
|
notification = None
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read(); self._prune(state)
|
|
queue = state["foreground_queue"]
|
|
eligible = next((q for q in queue if q.get("user_id") == str(user_id)), None)
|
|
if eligible and eligible.get("ticket") == ticket and self._can_admit(state, str(user_id), "foreground"):
|
|
queue[:] = [q for q in queue if q.get("ticket") != ticket]
|
|
state["leases"][ticket] = {"lease_id": ticket, "user_id": str(user_id), "kind": "foreground", "owner_pid": os.getpid(), "started_ts": time.time()}
|
|
self._write(state)
|
|
admitted = True
|
|
else:
|
|
self._write(state)
|
|
if not admitted and not wait_notified and wait_notify is not None:
|
|
leases = list(state["leases"].values())
|
|
user_running = sum(
|
|
1 for lease in leases
|
|
if lease.get("user_id") == str(user_id)
|
|
)
|
|
available = mem_available_bytes()
|
|
if user_running >= self.max_per_user:
|
|
reason = "per_user_limit"
|
|
elif len(leases) >= self.max_active:
|
|
reason = "global_limit"
|
|
elif available is not None and available < self.min_mem_available:
|
|
reason = "memory_pressure"
|
|
else:
|
|
reason = "queue_order"
|
|
notification = {
|
|
"state": "waiting",
|
|
"reason": reason,
|
|
"user_running": user_running,
|
|
"user_limit": self.max_per_user,
|
|
"global_running": len(leases),
|
|
"global_limit": self.max_active,
|
|
}
|
|
wait_notified = True
|
|
if admitted:
|
|
if wait_notified and wait_notify is not None:
|
|
with suppress(Exception):
|
|
wait_notify({"state": "admitted"})
|
|
return ticket
|
|
if notification is not None and wait_notify is not None:
|
|
with suppress(Exception):
|
|
wait_notify(notification)
|
|
if cancel_check is not None and cancel_check():
|
|
self.cancel_waiter(ticket)
|
|
return None
|
|
time.sleep(0.1)
|
|
|
|
def cancel_waiter(self, ticket: str) -> None:
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read()
|
|
state["foreground_queue"] = [q for q in state["foreground_queue"] if q.get("ticket") != ticket]
|
|
self._write(state)
|
|
|
|
def release(self, lease_id: Optional[str]) -> None:
|
|
if not lease_id:
|
|
return
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read(); state["leases"].pop(lease_id, None); self._write(state)
|
|
|
|
def reconcile_background(self, running_proc_ids: set[str]) -> None:
|
|
"""服务重启后保留 Docker 仍在跑的租约，清启动窗口遗留的孤儿租约。"""
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read()
|
|
for key, lease in list(state["leases"].items()):
|
|
if (lease.get("kind") == "background"
|
|
and str(lease.get("proc_id") or "") not in running_proc_ids
|
|
and not self._pid_alive(int(lease.get("owner_pid") or 0))):
|
|
state["leases"].pop(key, None)
|
|
self._write(state)
|
|
|
|
def touch_container(self, name: str, *, active_delta: int = 0) -> None:
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read()
|
|
row = state["containers"].setdefault(name, {"active_execs": 0})
|
|
row["active_execs"] = max(0, int(row.get("active_execs") or 0) + active_delta)
|
|
row["last_active_ts"] = time.time()
|
|
self._write(state)
|
|
|
|
def remove_container(self, name: str) -> None:
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read(); state["containers"].pop(name, None); self._write(state)
|
|
|
|
def reconcile_containers(self, running_names: set[str]) -> None:
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read()
|
|
for name in list(state["containers"]):
|
|
if name not in running_names:
|
|
state["containers"].pop(name, None)
|
|
self._write(state)
|
|
|
|
@contextmanager
|
|
def foreground(
|
|
self,
|
|
user_id: str,
|
|
cancel_check: Optional[Callable[[], bool]] = None,
|
|
wait_notify: Optional[Callable[[dict], None]] = None,
|
|
) -> Iterator[bool]:
|
|
lease = self.acquire_foreground(user_id, cancel_check, wait_notify)
|
|
try:
|
|
yield lease is not None
|
|
finally:
|
|
self.release(lease)
|
|
|
|
def snapshot(self) -> dict:
|
|
with interprocess_file_lock(self.lock_path, timeout_seconds=None):
|
|
state = self._read(); self._prune(state); self._write(state)
|
|
leases = list(state["leases"].values())
|
|
fg = sum(x.get("kind") == "foreground" for x in leases)
|
|
bg = sum(x.get("kind") == "background" for x in leases)
|
|
by_user: Dict[str, int] = {}
|
|
for x in leases:
|
|
by_user[x["user_id"]] = by_user.get(x["user_id"], 0) + 1
|
|
available_mem = mem_available_bytes()
|
|
memory_paused = available_mem is not None and available_mem < self.min_mem_available
|
|
containers = state.get("containers", {})
|
|
now = time.time()
|
|
return {
|
|
"limits": {"active": self.max_active, "background": self.max_background, "per_user": self.max_per_user},
|
|
"foreground_running": fg, "foreground_queued": len(state["foreground_queue"]),
|
|
"background_running": bg, "per_user": by_user,
|
|
"admit_available": 0 if memory_paused else max(0, self.max_active - len(leases)),
|
|
"memory_paused": memory_paused,
|
|
"mem_available_bytes": available_mem,
|
|
"cpu_load": list(os.getloadavg()) if hasattr(os, "getloadavg") else None,
|
|
"active_sandbox_containers": len(containers),
|
|
"idle_reap_candidates": sum(int(x.get("active_execs") or 0) == 0 and float(x.get("last_active_ts") or now) < now - getattr(self, "idle_ttl_seconds", 600) for x in containers.values()),
|
|
}
|