236 lines
7.7 KiB
Python
236 lines
7.7 KiB
Python
"""Per-user 工作目录配额(§7.5 #4 软配额,应用层 gate)。
|
||
|
||
调用入口:
|
||
- `scan_user_dir(user_root) -> (bytes, count)` ── 累加计入用户配额的文件
|
||
- `scan_user_trash(user_root) -> (bytes, count)` ── 单独统计隐藏回收站物理占用
|
||
- `upsert_user_usage(...)` ── 同时落配额占用与回收站快照
|
||
- `check_disk_quota(user_id, limit_bytes) -> Optional[str]` ── 写前查,返 None=放行 /
|
||
str=拒绝原因。`limit_bytes <= 0` 短路放行(不限)
|
||
- `scan_all_users(user_root_base, limit_bytes)` ── lifespan 后台 task 周期跑,
|
||
per user 跑完后下一个,避免 IO 风暴
|
||
|
||
字节单位解析(yaml `disk_bytes_per_user`):
|
||
- 整数字节 / "5gb" / "500mb" / "1.5g" 等 case-insensitive 后缀
|
||
- 失败返 None,caller 视为不限
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import re
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Iterable, List, Optional, Tuple
|
||
from uuid import UUID
|
||
|
||
from sqlalchemy import select
|
||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
|
||
from .engine import session_scope
|
||
from .models import UserDiskUsage
|
||
|
||
|
||
# yaml 字节解析:5gb / 500mb / 1024 / 1.5g
|
||
_SIZE_RE = re.compile(r"^\s*([\d.]+)\s*([kmgt]?b?)?\s*$", re.IGNORECASE)
|
||
_UNIT_FACTORS = {
|
||
"": 1, "b": 1,
|
||
"k": 1024, "kb": 1024,
|
||
"m": 1024 ** 2, "mb": 1024 ** 2,
|
||
"g": 1024 ** 3, "gb": 1024 ** 3,
|
||
"t": 1024 ** 4, "tb": 1024 ** 4,
|
||
}
|
||
|
||
|
||
def parse_bytes(value) -> Optional[int]:
|
||
"""yaml 字节值 → int;无法解析返 None。"""
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, int):
|
||
return value
|
||
if not isinstance(value, str):
|
||
return None
|
||
m = _SIZE_RE.match(value)
|
||
if not m:
|
||
return None
|
||
num_s, unit_s = m.group(1), (m.group(2) or "").lower()
|
||
factor = _UNIT_FACTORS.get(unit_s)
|
||
if factor is None:
|
||
return None
|
||
try:
|
||
return int(float(num_s) * factor)
|
||
except ValueError:
|
||
return None
|
||
|
||
|
||
# 扫描跳过的 dotfile 顶层名(节省 IO,且 /v1/files API 也隐藏)
|
||
_SKIP_TOPLEVEL = frozenset({
|
||
".zcbot_cache",
|
||
".zcbot_external_results",
|
||
".zcbot_artifact_trash",
|
||
".zcbot_tmp",
|
||
".memory",
|
||
})
|
||
_ARTIFACT_TRASH_DIR = ".zcbot_artifact_trash"
|
||
|
||
|
||
def scan_user_dir(user_root: Path) -> Tuple[int, int]:
|
||
"""os.walk 累加 user_root 下所有文件大小,返 (bytes, count)。
|
||
|
||
跳过顶层平台缓存/临时区、隐藏回收站与 .memory(均有独立生命周期,
|
||
不算入产品配额);
|
||
follow_symlinks=False 防 symlink 循环爆。
|
||
"""
|
||
if not user_root.exists() or not user_root.is_dir():
|
||
return 0, 0
|
||
|
||
total_bytes = 0
|
||
total_count = 0
|
||
try:
|
||
for entry in os.scandir(user_root):
|
||
if entry.name in _SKIP_TOPLEVEL:
|
||
continue
|
||
try:
|
||
if entry.is_file(follow_symlinks=False):
|
||
try:
|
||
total_bytes += entry.stat(follow_symlinks=False).st_size
|
||
total_count += 1
|
||
except OSError:
|
||
pass
|
||
elif entry.is_dir(follow_symlinks=False):
|
||
sub_b, sub_c = _walk_dir(Path(entry.path))
|
||
total_bytes += sub_b
|
||
total_count += sub_c
|
||
except OSError:
|
||
pass
|
||
except OSError:
|
||
pass
|
||
return total_bytes, total_count
|
||
|
||
|
||
def scan_user_trash(user_root: Path) -> Tuple[int, int]:
|
||
"""统计隐藏 artifact 回收站的物理占用，不计入用户配额。"""
|
||
trash = user_root / _ARTIFACT_TRASH_DIR
|
||
if not trash.is_dir():
|
||
return 0, 0
|
||
return _walk_dir(trash)
|
||
|
||
|
||
def _walk_dir(d: Path) -> Tuple[int, int]:
|
||
total_b, total_c = 0, 0
|
||
for root, dirs, files in os.walk(d, followlinks=False, onerror=lambda _e: None):
|
||
for f in files:
|
||
try:
|
||
st = os.stat(os.path.join(root, f), follow_symlinks=False)
|
||
total_b += st.st_size
|
||
total_c += 1
|
||
except OSError:
|
||
pass
|
||
return total_b, total_c
|
||
|
||
|
||
def upsert_user_usage(
|
||
user_id: UUID,
|
||
bytes_used: int,
|
||
file_count: int,
|
||
trash_bytes: int = 0,
|
||
trash_file_count: int = 0,
|
||
) -> None:
|
||
"""落 user_disk_usage 单行；首次 INSERT，后续 UPDATE。"""
|
||
from sqlalchemy import func
|
||
with session_scope() as s:
|
||
stmt = pg_insert(UserDiskUsage).values(
|
||
user_id=user_id,
|
||
bytes_used=bytes_used,
|
||
file_count=file_count,
|
||
trash_bytes=trash_bytes,
|
||
trash_file_count=trash_file_count,
|
||
).on_conflict_do_update(
|
||
index_elements=["user_id"],
|
||
set_={
|
||
"bytes_used": bytes_used,
|
||
"file_count": file_count,
|
||
"trash_bytes": trash_bytes,
|
||
"trash_file_count": trash_file_count,
|
||
"scanned_at": func.now(),
|
||
},
|
||
)
|
||
s.execute(stmt)
|
||
|
||
|
||
def get_user_usage(user_id: UUID) -> Optional[Tuple[int, int, Optional[datetime]]]:
|
||
"""读最近一次扫描结果 (bytes, count, scanned_at);无记录返 None。"""
|
||
with session_scope() as s:
|
||
row = s.execute(
|
||
select(
|
||
UserDiskUsage.bytes_used,
|
||
UserDiskUsage.file_count,
|
||
UserDiskUsage.scanned_at,
|
||
).where(UserDiskUsage.user_id == user_id)
|
||
).first()
|
||
if row is None:
|
||
return None
|
||
return int(row[0]), int(row[1]), row[2]
|
||
|
||
|
||
def check_disk_quota(user_id: UUID, limit_bytes: int) -> Optional[str]:
|
||
"""写前 gate:超额返 错误 msg(给 LLM 直读);放行返 None。
|
||
|
||
`limit_bytes <= 0` 短路放行(不限)。无扫描记录(首次,首次扫描前)放行 ──
|
||
避免冷启动期间所有写入卡死。15min 后周期扫到就生效。
|
||
"""
|
||
if limit_bytes <= 0:
|
||
return None
|
||
usage = get_user_usage(user_id)
|
||
if usage is None:
|
||
return None # 首次,放行,首次扫描后下次 gate 才生效
|
||
used, _, _ = usage
|
||
if used >= limit_bytes:
|
||
used_mb = used / (1024 ** 2)
|
||
limit_mb = limit_bytes / (1024 ** 2)
|
||
return (
|
||
f"[Error] 已达磁盘配额上限({used_mb:.1f} MB / {limit_mb:.1f} MB);"
|
||
f"清理旧产物或联系管理员升配后重试"
|
||
)
|
||
return None
|
||
|
||
|
||
def list_user_ids_with_root(user_root_base: Path) -> List[UUID]:
|
||
"""扫 user_root_base 子目录,返合法 UUID 列表(=有 workspace 子目录的 user)。
|
||
|
||
不去 DB 查 users 全表 —— 有些 user 可能从未发消息(无 workspace 目录),无 disk 占用,
|
||
无需 upsert 占位行。
|
||
"""
|
||
if not user_root_base.is_dir():
|
||
return []
|
||
out: List[UUID] = []
|
||
try:
|
||
for entry in os.scandir(user_root_base):
|
||
if not entry.is_dir(follow_symlinks=False):
|
||
continue
|
||
try:
|
||
out.append(UUID(entry.name))
|
||
except ValueError:
|
||
continue
|
||
except OSError:
|
||
pass
|
||
return out
|
||
|
||
|
||
def scan_all_users(user_root_base: Path) -> int:
|
||
"""扫所有 user 落库,返扫描的 user 数。lifespan 后台 task 调。
|
||
|
||
串行(per user 跑完下一个)避免 IO 风暴;单 user 几秒(几百 MB 量级),N user 总耗时
|
||
线性。失败的 user 静默跳过,下次周期再试。
|
||
"""
|
||
count = 0
|
||
for uid in list_user_ids_with_root(user_root_base):
|
||
try:
|
||
user_root = user_root_base / str(uid)
|
||
b, c = scan_user_dir(user_root)
|
||
trash_b, trash_c = scan_user_trash(user_root)
|
||
upsert_user_usage(uid, b, c, trash_b, trash_c)
|
||
count += 1
|
||
except Exception:
|
||
# 单 user 扫挂不阻塞其他 user;下次周期重试。日志靠 caller 注入。
|
||
pass
|
||
return count
|