diff --git a/Dockerfile b/Dockerfile index c3bd9ba7..54359a04 100644 --- a/Dockerfile +++ b/Dockerfile @@ -76,6 +76,8 @@ SH WORKDIR /app COPY LICENSE COMMERCIAL-LICENSE.md /app/ COPY backend/app /app/backend/app +# 锁文件留一份：进程登记的代码版本把它算进源码指纹（app/cluster/version.py：升了依赖的进程也算版本不同） +COPY backend/requirements.lock /app/backend/requirements.lock COPY --from=web /src/frontend/dist /app/web ENV AGENTLAB_WEB_DIST=/app/web \ diff --git a/README.md b/README.md index a9062297..8cd85445 100644 --- a/README.md +++ b/README.md @@ -1202,14 +1202,14 @@ asyncpg 方言一律这样做，同一个事务里、在同一条服务端连接 SQL 形状。各进程各自的并发闸门（口径比对、自动对比、起草的模型调用、上传解析）按进程算，进程多了总量跟着多。 **多台机器。** 几台机器要用同一份数据目录：放在共享存储（NFS 之类）上，每台挂在**同一个绝对路径**（库里记着上传表格各版本库文件的 -绝对路径），各台机器的时钟要对时（NTP；租约按墙钟算）。文件都是先写临时文件再改名，读的一方看不到写了一半的；写上传版本库、记录表 +绝对路径），各台机器的时钟仍要对时（NTP）：租约按库的时钟算，机器之间差多少都不影响谁算活着，但「停了两个租约时长」这类判断看的是各机器按本机时钟写下的时间。库服务器自己的时钟往前跳超过一个租约时（少见），所有进程续租时都会发现自己的租约已经过期、一起隔离：执行到一半的运行标成中断，在运行详情里点「继续运行」从断点接着执行；后台任务按「属主不在了」的规矩处理；各进程自动换身份重新登记，之后发起的照常。往回跳时不隔离，只是真没了的进程要多等跳回去的那么久才被接手。本机墙钟跳了（NTP 一步校时、手动改时间）不会隔离：判续没续晚看的是计入休眠的单调钟，对不上时以库里的租约为准。文件都是先写临时文件再改名，读的一方看不到写了一半的；写上传版本库、记录表 在库里的锁下一个接一个。但要说实话：数据目录里还有 SQLite 文件——上传表格的各版本库（发布后只读）、填报的记录表（`records/*.db`， 原地追加写）、沙箱里代码自己建的库——SQLite 文件放在网络文件系统上，读写一致性靠文件锁，而 NFS 上的文件锁不可靠：一台机器在追加时 另一台读，可能读到不一致的页。所以多台机器时，记录表写得多的部署建议用一台机器多进程，或者把数据目录放在支持可靠文件锁的存储上。 沙箱会话（microVM）在各进程自己手上，不跨进程共享。 -**升级。** 不支持新旧版本的进程同时连着一个库，代码里拦着：每个进程登记时带上代码版本（包版本、构建号或 git 提交、源码指纹）和 -库结构版本，库里还有活着的、版本不同的进程时，新版本的进程在补列之前就拒绝启动，报错里写明是哪几个进程、什么版本；同版本的照常 +**升级。** 不支持新旧版本的进程同时连着一个库，代码里拦着：每个进程登记时带上代码版本（包版本、构建号或 git 提交、源码指纹——源码加上依赖锁文件，只升了依赖也算版本不同）和 +库结构版本，库里还有活着的、版本不同的进程时，新版本的进程在补列之前就拒绝启动，报错里列出全部这样的进程、各是什么版本、心跳还在不在走；同版本的照常 加入。升级步骤： 1. 停掉全部旧版本的进程（平稳关停的马上删掉自己的登记；强杀的等它的租约过期，新进程启动时会等）； diff --git a/backend/app/apps/resolve.py b/backend/app/apps/resolve.py index 95bb9d58..f033259f 100644 --- a/backend/app/apps/resolve.py +++ b/backend/app/apps/resolve.py @@ -487,7 +487,8 @@ async def run_view(session: AsyncSession, *, app: Any, version: AppVersion, view return int(await s.scalar(select(func.count(Run.id)).where( Run.app_id == app.id, Run.origin_kind == "app_view", Run.created_at >= start)) or 0) - if not await quota.take(f"app_view:{app.id}", limit, seed=used_today): + day = await quota.take(f"app_view:{app.id}", limit, seed=used_today) + if day is None: raise ViewRefused(429, f"这个应用今天的按需运行已达到上限（{limit} 次），请明天再试，或请管理员调整上限", DAILY_LIMIT) try: @@ -498,7 +499,7 @@ async def run_view(session: AsyncSession, *, app: Any, version: AppVersion, view origin={"view": view_id, "spec_hash": version.spec_hash, "actor_verified": principal.verified, "auth_mode": principal.mode, **_key_trace(principal)}) except BaseException: - await quota.give_back(f"app_view:{app.id}") + await quota.give_back(f"app_view:{app.id}", day=day) raise ttl = source.get("ttl_s") if ttl is None and await reads_live(session, state.graph): diff --git a/backend/app/cluster/fence.py b/backend/app/cluster/fence.py index 60da513b..ec1eecc7 100644 --- a/backend/app/cluster/fence.py +++ b/backend/app/cluster/fence.py @@ -6,16 +6,28 @@ - 续租约在单独的任务、单独的连接上做（cluster/service.py），不被同一拍别的活拖住；最近一次续上的时刻记在这里（renewed）。 - 本进程的执行在动手写之前问一句 check()：离上次续上已经超过一个租约时长（别的进程按两个租约时长才接手，中间留着一段），或者这段 - 执行是在上一次隔离之前开始的（世代不对），就抛 Fenced——写事件、写断点、写终态、写待办和动作、发外部写入之前都问。 + 执行是在上一次隔离之前开始的（世代不对），就抛 Fenced——写事件、写断点、写终态、写待办和动作、发外部写入之前都问。只对绑定了 + 世代的执行（运行的一段执行、后台任务）抛；别的调用者问到时只隔离、不抛。 - 一旦发现续晚了（trip）：世代加一，手上的执行全部停掉（停的时候不写终态，那些行交给接手的进程按「属主不在了」处理），换一个新的 - 身份重新登记，之后发起的照常。 + 身份重新登记；新身份登记上之前新开始的执行先等着（ready），登记上之后照常。 + +**用哪只钟（复核 R-b7、R-b14）。** 别的进程判「租约过期」用的是库的时钟（cluster/workers：续租写 lease_until = 库的 now() + +租约），跟各台机器的墙钟差多少无关。这边判「续晚了」用计入休眠的钟：time.monotonic() 在 macOS（mach_absolute_time）、Linux +（CLOCK_MONOTONIC）上都不计休眠——笔记本合盖、虚拟机挂起超过两个租约，对面早已接手，这边还以为刚续过。所以 Linux 用 +CLOCK_BOOTTIME，macOS 用 clock_gettime(CLOCK_MONOTONIC)（这个计入休眠），只看它；墙钟会被 NTP 一步校时、手动改时间、虚拟机 +对时拨来拨去，取得到计入休眠的钟时不拿它判（否则单进程墙钟一跳就把自己隔离了）。两只都取不到时退回 monotonic，墙钟当后备： +两只里走得多的那只说过了一个租约就算续晚了。两只钟对不上（jumped，差出一个租约以上）时不直接隔离：先续一次租，续租那一句按库的 +时钟核对「自己的租约还没过期」，以它为准——续上了记一条日志照常，续不上（库里自己的租约已经过期、或者登记不见了）才隔离。 +虚拟机被宿主挂起时客户机的 CLOCK_BOOTTIME 也可能停着：那时靠续租那一句（最迟一拍心跳之后）发现过期。 Fenced 是 BaseException：节点、工具里常见的 except Exception 拦不住它，一路退到执行的最外层。没起协调（测试、脚本）时什么都不查。 """ from __future__ import annotations +import asyncio import contextvars import logging +import sys import time from collections.abc import Callable from typing import Any @@ -27,29 +39,55 @@ class Fenced(BaseException): """本进程续不上租约了（或者这段执行属于上一次隔离之前）：不许再写。""" +def _sleep_aware_clock() -> tuple[Callable[[], float], bool]: + """计入休眠的单调钟，和它是不是真的计入休眠。""" + if sys.platform.startswith("linux") and hasattr(time, "CLOCK_BOOTTIME"): + return (lambda: time.clock_gettime(time.CLOCK_BOOTTIME)), True + if sys.platform == "darwin" and hasattr(time, "CLOCK_MONOTONIC"): + return (lambda: time.clock_gettime(time.CLOCK_MONOTONIC)), True + return time.monotonic, False + + +_clock, SLEEP_AWARE = _sleep_aware_clock() +#: 墙钟（单独一个名字，测试里拨它） +_wall: Callable[[], float] = time.time + _generation = 0 _renewed: float | None = None +_renewed_wall: float | None = None _active = False _handlers: list[Callable[[], Any]] = [] _bound: contextvars.ContextVar[int | None] = contextvars.ContextVar("agentlab_fence_generation", default=None) -def activate(renewed_at: float | None = None) -> None: +def now() -> float: + """计入休眠的单调钟的读数（本进程里判时长都用它：cluster/leader 判还当不当班也是）。""" + return _clock() + + +def activate() -> None: """起协调、续上了第一次租约时（PostgreSQL；SQLite 只有一个进程，不用隔离）。""" - global _active, _renewed + global _active _active = True - _renewed = time.monotonic() if renewed_at is None else renewed_at + renewed() def deactivate() -> None: - global _active, _renewed - _active, _renewed = False, None + global _active, _renewed, _renewed_wall + _active, _renewed, _renewed_wall = False, None, None -def renewed(at: float | None = None) -> None: - """续上了一次租约。""" - global _renewed - _renewed = time.monotonic() if at is None else at +def renewed(*, ago: float = 0.0) -> None: + """续上了一次租约（ago：其实是多少秒以前续上的，测试里装「卡了一阵」用）。""" + global _renewed, _renewed_wall + _renewed, _renewed_wall = _clock() - ago, _wall() - ago + + +def _since() -> tuple[float, float] | None: + """离上次续上：(计入休眠的钟走了多少, 墙钟走了多少)。""" + if _renewed is None or _renewed_wall is None: + return None + return _clock() - _renewed, _wall() - _renewed_wall def generation() -> int: @@ -61,32 +99,69 @@ def bind() -> None: _bound.set(_generation) -def stale(now: float | None = None) -> bool: - """离上次续上租约是不是已经超过一个租约时长。""" - if not _active or _renewed is None: +def stale() -> bool: + """离上次续上租约是不是已经超过一个租约时长：取得到计入休眠的钟时只看它；取不到时它（monotonic）和墙钟里走得多的那只说了算。 + 隔离之后、新身份还没登记上的那一段也算（pending）。""" + if not _active: + return False + since = _since() + if since is None: + return True + from app.cluster.workers import lease_s + + return (since[0] if SLEEP_AWARE else max(since)) > lease_s() + + +def pending() -> bool: + """隔离过、换的新身份还没登记上。""" + return _active and _renewed is None + + +async def ready(timeout: float | None = None) -> None: + """新开始的一段执行（运行、后台任务）在 bind() 之前等一等：隔离之后、新身份还没登记上时，它认领的行在别的进程眼里属主 + 不在，会被当成没人管接手（复核 R-b9）。最多等一个租约时长（timeout），等不到抛 Fenced——不开始，那几行由别的进程照「属主 + 不在了」处理（cluster/service 那边同时会放弃这个一直没登记上的身份）。""" + if not pending(): + return + from app.cluster.workers import lease_s + + deadline = _clock() + (lease_s() if timeout is None else timeout) + while pending(): + if _clock() >= deadline: + raise Fenced() + await asyncio.sleep(0.05) + + +def jumped() -> bool: + """两只钟差出一个租约以上（墙钟往前或往后跳了，或者单调钟没计入休眠）。续租时问：是的话不直接隔离，以续租那一句按库的时钟 + 核对的结果为准（cluster/service.renew）。""" + if not _active: + return False + since = _since() + if since is None: return False from app.cluster.workers import lease_s - return (time.monotonic() if now is None else now) - _renewed > lease_s() + return abs(since[0] - since[1]) > lease_s() def fenced_here() -> bool: - """当前这段执行还能不能写：世代不对（隔离之前开始的）或者租约续晚了。""" + """当前这段执行还能不能写：世代不对（隔离之前开始的）或者租约续晚了。只说绑定了世代的执行（bind 过的）。""" if not _active: return False bound = _bound.get() - return (bound is not None and bound != _generation) or stale() + return bound is not None and (bound != _generation or stale()) def check() -> None: - """写之前问一句：不许写就抛 Fenced（续晚了的话顺手隔离）。""" + """写之前问一句。续晚了就隔离（谁先问到都一样）；抛 Fenced 只对绑定了世代的执行（运行的一段执行、后台任务：bind 过的）—— + 别的调用者（HTTP 处理、调度那一拍、进程间请求）不是认领了行在跑的执行，抛给它们只会把调度循环之类的打停（复核 R-b9）。""" if not _active: return - if stale(): + if _renewed is not None and stale(): trip("执行要写的时候发现租约已经续晚了") - raise Fenced() bound = _bound.get() - if bound is not None and bound != _generation: + if bound is not None and (bound != _generation or stale()): raise Fenced() @@ -97,11 +172,11 @@ def on_trip(handler: Callable[[], Any]) -> None: def trip(why: str) -> bool: """隔离：世代加一，叫各处停掉手上的执行（cluster/service 再换身份重新登记）。已经在这一世代隔离过的不重复。交回这一次有没有动手。""" - global _generation, _renewed + global _generation, _renewed, _renewed_wall if not _active or _renewed is None: return False _generation += 1 - _renewed = None + _renewed = _renewed_wall = None logger.error("本进程的租约续晚了（%s）：当作认领的运行、问数据、后台任务可能已被别的进程接手，停掉手上的执行、不写终态，" "换一个身份重新登记", why) for handler in list(_handlers): @@ -112,5 +187,5 @@ def trip(why: str) -> bool: return True -__all__ = ["Fenced", "activate", "bind", "check", "deactivate", "fenced_here", "generation", "on_trip", "renewed", - "stale", "trip"] +__all__ = ["Fenced", "SLEEP_AWARE", "activate", "bind", "check", "deactivate", "fenced_here", "generation", "jumped", "now", + "on_trip", "pending", "ready", "renewed", "stale", "trip"] diff --git a/backend/app/cluster/leader.py b/backend/app/cluster/leader.py index fe169232..b532b11e 100644 --- a/backend/app/cluster/leader.py +++ b/backend/app/cluster/leader.py @@ -7,6 +7,9 @@ 两个进程短暂地都以为自己当班（续租那一拍库慢了、本地还没过期）时，重复不了：定时认领本来就是条件 UPDATE、触发记录 (任务, 时刻, 种类) 唯一，事件派发按投递记录的唯一约束，待办的升级、核实各自也是条件更新。当班只是让它们平常不撞。 本地另留一截余量（MARGIN）：租约快到期还没续上时，自己先当作不当班了。 + +租约按库的时钟写、按库的时钟判（复核 R-b7，同 cluster/workers）：机器之间墙钟差多少都不会让两边都以为租约归自己。本地判「还当 +不当班」不看墙钟，看计入休眠的单调钟（cluster/fence.now）：从发出续租那一刻起算一个租约时长，库里的租约只会更晚到期。 """ from __future__ import annotations @@ -15,9 +18,9 @@ from datetime import datetime, timedelta from sqlalchemy import case, or_, update -from app.cluster import identity -from app.cluster.workers import lease_s -from app.db.base import SessionLocal, utcnow +from app.cluster import fence, identity +from app.cluster.workers import db_now, lease_s +from app.db.base import SessionLocal logger = logging.getLogger(__name__) @@ -29,33 +32,42 @@ MARGIN_S = 2.0 class Leadership: def __init__(self, name: str) -> None: self.name = name + #: 库里记的到期时刻（库的时钟；给人看、交班用） self.until: datetime | None = None + #: 本地判还当不当班用的到期时刻（cluster/fence.now 的读数） + self.local_until: float | None = None self.epoch = 0 def single(self) -> None: """只有本进程（SQLite）：直接算当班，不写库。""" self.until = datetime.max.replace(tzinfo=None) + self.local_until = float("inf") - def held(self, now: datetime | None = None) -> bool: + def drop(self) -> None: + """本地当作不当班了（隔离之后换身份时：库里那份租约属于旧身份）。""" + self.until = self.local_until = None + + def held(self) -> bool: """本地看：还在当班（租约离到期还有余量）。""" - if self.until is None: + if self.until is None or self.local_until is None: return False if self.until.tzinfo is None: return True # single()：只有本进程 # 余量不超过租约的四分之一：租约调得很短（验收脚本）时也留得出当班的时间 margin = min(MARGIN_S, lease_s() / 4) - return (now or utcnow()) + timedelta(seconds=margin) < self.until + return fence.now() + margin < self.local_until - async def claim(self, now: datetime | None = None) -> bool: + async def claim(self) -> bool: """接班或续租。交回这一拍之后是不是当班。""" from app.db.models_cluster import ClusterLease from app.db.portable import insert_ignore - now = now or utcnow() - until = now + timedelta(seconds=lease_s()) + started = fence.now() me = identity.owner() mine = ClusterLease.owner == me async with SessionLocal() as session: + now = await db_now(session) + until = now + timedelta(seconds=lease_s()) done = await session.execute( update(ClusterLease).where( ClusterLease.name == self.name, @@ -74,15 +86,15 @@ class Leadership: .on_conflict_do_nothing(index_elements=["name"])) epoch = 1 if inserted.rowcount == 1 else None await session.commit() - was = self.held(now) + was = self.held() if epoch is None: if was or self.until is not None: logger.warning("调度当班交给了别的进程（本进程 %s 不再当班）", identity.describe(me)) - self.until = None + self.drop() return False if not was and self.epoch != epoch: logger.info("本进程（%s）接手调度当班（第 %d 班）", identity.describe(me), epoch) - self.until, self.epoch = until, int(epoch) + self.until, self.local_until, self.epoch = until, started + lease_s(), int(epoch) return True async def release(self) -> None: @@ -92,7 +104,7 @@ class Leadership: if self.until is None: return single = self.until.tzinfo is None - self.until = None + self.drop() if single: return async with SessionLocal() as session: diff --git a/backend/app/cluster/quota.py b/backend/app/cluster/quota.py index 9a726966..bb811ab8 100644 --- a/backend/app/cluster/quota.py +++ b/backend/app/cluster/quota.py @@ -29,13 +29,15 @@ def _key(name: str, day: str) -> str: async def take(name: str, limit: int, *, seed: Callable[[Any], Awaitable[int]] | None = None, - day: str | None = None) -> bool: - """占一份额度：这一天占过的还不到 limit 就加一、交回 True，到了交回 False（什么都不改）。seed(session)：这一天第一次占时 - 已经用掉了几份（按库里的记录数）。""" + day: str | None = None) -> str | None: + """占一份额度：这一天占过的还不到 limit 就加一、交回占在了哪一天（本地日期），到了交回 None（什么都不改）。seed(session)： + 这一天第一次占时已经用掉了几份（按库里的记录数）。退回时把交回的日期给 give_back（复核 R-b11：23:59:59 占、00:00:01 发起 + 失败，要退回前一天那一份，不是第二天的）。""" from app.db.models_cluster import RateCounter as Row from app.db.portable import insert_ignore - key = _key(name, day or today()) + day = day or today() + key = _key(name, day) async with SessionLocal() as session: if await session.scalar(select(Row.key).where(Row.key == key)) is None: used = int(await seed(session)) if seed is not None else 0 @@ -45,11 +47,11 @@ async def take(name: str, limit: int, *, seed: Callable[[Any], Awaitable[int]] | done = await session.execute(update(Row).where(Row.key == key, Row.count < int(limit)) .values(count=Row.count + 1).execution_options(synchronize_session=False)) await session.commit() - return done.rowcount == 1 + return day if done.rowcount == 1 else None async def give_back(name: str, *, day: str | None = None) -> None: - """退回一份（占了额度、却没发起成）。""" + """退回一份（占了额度、却没发起成）。day 给 take 交回的那个日期（不给按今天，只有测试这么用）。""" from app.db.models_cluster import RateCounter as Row key = _key(name, day or today()) diff --git a/backend/app/cluster/requests.py b/backend/app/cluster/requests.py index 34fb7e8f..2711d4d6 100644 --- a/backend/app/cluster/requests.py +++ b/backend/app/cluster/requests.py @@ -94,9 +94,11 @@ _running: set[asyncio.Future[Any]] = set() async def _handle(fn: Callable[..., Awaitable[Any]], kind: str, target: str, target_id: str, actor: str | None, payload: dict[str, Any]) -> None: + from app.cluster.fence import Fenced + try: await fn(target_id, actor, payload) - except Exception: # noqa: BLE001 + except (Exception, Fenced): # noqa: BLE001 - Fenced 是 BaseException，不接住就成了没人取的任务异常 logger.exception("办进程间请求 %s/%s %s 出错", kind, target, target_id) diff --git a/backend/app/cluster/service.py b/backend/app/cluster/service.py index ff9cf7b3..2d40ed08 100644 --- a/backend/app/cluster/service.py +++ b/backend/app/cluster/service.py @@ -65,6 +65,11 @@ class _Service: #: SQLite：守卫已经确认只有本进程——登记只写一次、不续租、本进程直接当班、接手只在启动时做（不为协调多写库） self.single = False self._lease_conn: Any = None + self.lease_wake: asyncio.Event | None = None + #: 隔离之后、新身份还没登记上：从什么时候起（cluster/fence.now 的读数） + self.pending_since: float | None = None + #: 隔离掉（或者一直没登记上、放弃了）的旧身份：新身份登记上之后收拾它们留下的登记、名额、请求、当班租约 + self.retired: list[str] = [] async def start(self) -> None: """登记、开始收发通知、起心跳。SQLite 时先取单进程的守卫（拿不到抛 guard.SingleProcessOnly）。""" @@ -94,6 +99,7 @@ class _Service: await self._lead() fence.activate() fence.on_trip(self._on_fenced) + self.lease_wake = asyncio.Event() self.lease_task = asyncio.create_task(self._lease_loop(), name="agentlab-cluster-lease") self.task = asyncio.create_task(self._loop(), name="agentlab-cluster-heartbeat") logger.info("工作进程 %s 已登记（%s，租约 %g 秒、心跳 %g 秒；现在另有 %d 个进程）", identity.describe(identity.owner()), @@ -123,23 +129,28 @@ class _Service: return self._lease_conn async def renew(self) -> bool: - """续一拍租约（在专用连接上，不走连接池：池子被占满时取连接要等，等得比租约还长）。续晚了（离上次续上超过一个租约时长） - 就不续这个身份了：隔离（cluster/fence.py），换身份重新登记。交回续没续上。""" - from datetime import timedelta - + """续一拍租约（在专用连接上，不走连接池：池子被占满时取连接要等，等得比租约还长）。按库的时钟续、按库的时钟核对（复核 + R-b7）：库里自己的租约已经过期（或者登记不见了），别的进程可能已经接手，不再续这个身份——隔离（cluster/fence.py），换身份 + 重新登记。计入休眠的钟说续晚了，同样隔离；两只钟差出一个租约（墙钟跳了）只记一条日志，以库的核对为准（复核 R-b14）。 + 交回续没续上。""" from app.cluster import fence - from app.db.base import utcnow if fence.stale(): fence.trip("续租约的时候发现已经续晚了") return False - now = utcnow() - until = now + timedelta(seconds=workers.lease_s()) + # 两只钟对不上（墙钟跳了，或者单调钟没计入休眠）：不直接隔离，以下面这一句按库的时钟核对的结果为准（复核 R-b14） + jumped = fence.jumped() try: conn = await self._lease_conn_open() - done = await asyncio.wait_for(conn.execute( - "UPDATE cluster_workers SET heartbeat_at = $2, lease_until = $3 WHERE id = $1", - identity.owner(), now.replace(tzinfo=None), until.replace(tzinfo=None)), timeout=workers.heartbeat_s()) + renewed = await asyncio.wait_for(conn.fetchval( + "UPDATE cluster_workers SET heartbeat_at = timezone('UTC', clock_timestamp()), " + "lease_until = timezone('UTC', clock_timestamp()) + make_interval(secs => $2) " + "WHERE id = $1 AND lease_until >= timezone('UTC', clock_timestamp()) RETURNING id", + identity.owner(), float(workers.lease_s())), timeout=workers.heartbeat_s()) + if renewed is None: + expired = await asyncio.wait_for(conn.fetchval( + "SELECT lease_until < timezone('UTC', clock_timestamp()) FROM cluster_workers WHERE id = $1", + identity.owner()), timeout=workers.heartbeat_s()) except Exception: # noqa: BLE001 - 库一时连不上：这一拍没续上，下一拍再试（续晚了的话下一拍隔离） logger.warning("续租约没续上，下一拍再试", exc_info=True) with contextlib.suppress(Exception): @@ -147,34 +158,115 @@ class _Service: await self._lease_conn.close() self._lease_conn = None return False - if done.endswith(" 0"): - logger.warning("本进程的登记不见了（可能被当成没了），重新登记") - await workers.register() + if renewed is None: + why = "按库的时钟本进程的租约已经过期" if expired else "本进程的登记不见了（被当成不在了）" + fence.trip(why + ("；本机的墙钟和单调钟也差出了一个租约以上" if jumped else "")) + return False + if jumped: + logger.warning("本机的墙钟和单调钟差出一个租约以上（墙钟跳了：NTP 校时、手动改时间、虚拟机对时），按库的时钟本进程的租约" + "还没过期：照常续上，不隔离") fence.renewed() return True async def _lease_loop(self) -> None: + """续租约；隔离之后改成重新登记新身份。哪一下出了错（库一抖、版本守卫拦下）都不让这个任务退出——它一退，之后就再也没有 + 隔离保护了（复核 R-b9）：照实记日志，退避着重试（最长一个租约时长）。""" + from app.cluster import fence + + delay = workers.heartbeat_s() while True: - await self.renew() - await asyncio.sleep(workers.heartbeat_s()) + try: + if self.pending_since is not None: + await self.register_again() + else: + await self.renew() + delay = workers.heartbeat_s() + except asyncio.CancelledError: + raise + except (Exception, fence.Fenced): # noqa: BLE001 + delay = min(max(delay * 2, workers.heartbeat_s()), workers.lease_s()) + logger.exception("%s出错，%g 秒后再试", "隔离之后重新登记" if self.pending_since is not None else "续租约", delay) + if self.lease_wake is None: + await asyncio.sleep(delay) + continue + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(self.lease_wake.wait(), timeout=delay) + self.lease_wake.clear() def _on_fenced(self) -> None: - """隔离了：手上的执行由各处自己停（runner、job_runtime 登记的处理），这里换身份重新登记。""" - asyncio.ensure_future(self._reidentify()) - - async def _reidentify(self) -> None: + """隔离了：手上的执行由各处自己停（runner、job_runtime 登记的处理）；这里马上换身份（之后新建的行记新身份），叫续租的任务 + 去登记它。""" from app.cluster import fence old = identity.owner() + self.retired.append(old) new = identity.regenerate() - leader.scheduler.until = None + self.pending_since = fence.now() + leader.scheduler.drop() + logger.warning("隔离：换成新身份 %s（旧身份 %s），登记上之前新开始的执行先等着", new, old) + if self.lease_wake is not None: + self.lease_wake.set() + + async def register_again(self) -> None: + """隔离之后登记新身份。一个租约时长都没登记上（库连不上、版本守卫拦着）就放弃这个身份、再换一个：这期间等不到的执行已经 + 放弃开始（fence.ready），它们的行记着这个身份——它永远不登记，别的进程就照「属主不在了」处理那几行。""" + from app.cluster import fence + + assert self.pending_since is not None + if fence.now() - self.pending_since > workers.lease_s(): + gave_up = identity.owner() + self.retired.append(gave_up) + identity.regenerate() + self.pending_since = fence.now() + logger.error("新身份 %s 一个租约时长都没登记上，放弃它、换成 %s 再试", gave_up, identity.owner()) + await workers.register() + retired, self.retired = self.retired, [] + self.pending_since = None + fence.renewed() + logger.warning("隔离之后换了身份重新登记：%s → %s（之前认领的行由当班进程按「属主不在了」接手）", + "、".join(retired), identity.owner()) try: - await workers.register() - fence.renewed() - logger.warning("隔离之后换了身份重新登记：%s → %s（之前认领的行由当班进程按「属主不在了」接手）", - identity.describe(old), identity.describe(new)) - except Exception: # noqa: BLE001 - logger.exception("隔离之后重新登记失败，下一拍再试") + await self.tidy_retired(retired) + except Exception: # noqa: BLE001 - 收拾不了的等旧身份的租约过期、按「属主不在了」处理 + logger.exception("收拾旧身份留下的登记、名额、请求时出错（等它们的租约过期再按属主不在了处理）") + + async def tidy_retired(self, retired: list[str]) -> dict[str, int]: + """新身份登记上之后，收拾旧身份留下的（复核 R-b10）：它的执行都已停下、写不了了，不必等它的租约过期——删掉它的登记（别的 + 进程马上按「属主不在了」接手它认领的运行、任务）、它占着的运行名额和同源队的位置，交出它拿着的当班租约；发给它的请求转到 + 新身份名下照办（停止请求对已经不在本进程执行的运行照停止的规矩办：属主不在了就接手、放弃）。""" + from sqlalchemy import delete, update + + from app.db.base import SessionLocal + from app.db.models_cluster import ClusterLease, ClusterRequest, ClusterWorker, LaneTicket, RunSlot + + if not retired: + return {} + me = identity.owner() + async with SessionLocal() as session: + out = { + "workers": (await session.execute(delete(ClusterWorker).where(ClusterWorker.id.in_(retired)))).rowcount, + "slots": (await session.execute(delete(RunSlot).where(RunSlot.owner.in_(retired)))).rowcount, + "lanes": (await session.execute(delete(LaneTicket).where(LaneTicket.owner.in_(retired)))).rowcount, + "requests": (await session.execute(update(ClusterRequest).where(ClusterRequest.owner.in_(retired)) + .values(owner=me).execution_options(synchronize_session=False))).rowcount, + "leases": (await session.execute(update(ClusterLease).where(ClusterLease.owner.in_(retired)) + .values(owner=None, lease_until=None) + .execution_options(synchronize_session=False))).rowcount, + } + await session.commit() + out = {k: int(v or 0) for k, v in out.items()} + for old in retired: + signals.publish("gone", old) + if out["slots"]: + signals.publish("slot") + if out["lanes"]: + signals.publish("lane") + if out["requests"]: + signals.publish("q", me) + self.reap_now = True + logger.warning("收拾旧身份 %s：删了登记 %d 条、运行名额 %d 个、同源队的位置 %d 个，%d 条发给它的请求转给 %s，交出当班租约 %d 份", + "、".join(retired), out["workers"], out["slots"], out["lanes"], out["requests"], me, out["leases"]) + return out # ---------------- 当班、心跳的其余活 ---------------- @@ -198,8 +290,9 @@ class _Service: return now_leading async def beat(self) -> None: - """心跳的一拍（续租约另在 _lease_loop；测试里直接调）。SQLite 只做本进程自己攒着的活。""" - if self.single: + """心跳的一拍（续租约另在 _lease_loop；测试里直接调）。SQLite 只做本进程自己攒着的活；隔离之后新身份还没登记上时也只做 + 这些（不争当班、不接手）。""" + if self.single or self.pending_since is not None: await self._own_work() return live = await workers.liveness() @@ -249,6 +342,8 @@ class _Service: return out async def _loop(self) -> None: + from app.cluster import fence + assert self.wake is not None while True: with contextlib.suppress(asyncio.TimeoutError): @@ -258,7 +353,7 @@ class _Service: await self.beat() except asyncio.CancelledError: raise - except Exception: # noqa: BLE001 - 一拍出错（库一时不通）不该让心跳停下 + except (Exception, fence.Fenced): # noqa: BLE001 - 一拍出错（库一时不通）不该让心跳停下 logger.exception("心跳这一拍出错，下一拍照常进行") async def _light_loop(self) -> None: @@ -279,6 +374,7 @@ class _Service: with contextlib.suppress(asyncio.CancelledError, Exception): await task self.task = self.lease_task = None + self.pending_since, self.retired = None, [] fence.deactivate() with contextlib.suppress(Exception): await leader.scheduler.release() diff --git a/backend/app/cluster/slots.py b/backend/app/cluster/slots.py index d363dedb..48878c6d 100644 --- a/backend/app/cluster/slots.py +++ b/backend/app/cluster/slots.py @@ -196,6 +196,8 @@ class _Holder: woke: asyncio.Event = field(default_factory=asyncio.Event) #: 还在排队等名额（发下来之前） queued: bool = True + #: 要名额时本进程的身份（名额行记的属主）：放的时候按它删——隔离之后换了身份也删得到（复核 R-b10） + owner: str = "" _holders: dict[str, _Holder] = {} @@ -341,27 +343,30 @@ async def _granted(run_id: str) -> bool | None: return await session.scalar(select(RunSlot.granted).where(RunSlot.run_id == run_id)) -async def _request(run_id: str, klass: str, priority: int) -> None: +async def _request(run_id: str, klass: str, priority: int) -> str: + """排进名额的队。交回这一行记的属主（本进程此刻的身份）。""" + me = identity.owner() if _in_memory(): - _mem[run_id] = _MemRow(run_id, klass, priority, utcnow(), identity.owner()) - return + _mem[run_id] = _MemRow(run_id, klass, priority, utcnow(), me) + return me from app.db.models_cluster import RunSlot async with SessionLocal() as session: await session.execute(delete(RunSlot).where(RunSlot.run_id == run_id)) - session.add(RunSlot(run_id=run_id, klass=klass, priority=priority, owner=identity.owner(), granted=False, + session.add(RunSlot(run_id=run_id, klass=klass, priority=priority, owner=me, granted=False, requested_at=utcnow())) await session.commit() + return me -async def _release(run_id: str) -> None: +async def _release(run_id: str, owner: str) -> None: if _in_memory(): _mem.pop(run_id, None) return from app.db.models_cluster import RunSlot async with SessionLocal() as session: - await session.execute(delete(RunSlot).where(RunSlot.run_id == run_id, RunSlot.owner == identity.owner())) + await session.execute(delete(RunSlot).where(RunSlot.run_id == run_id, RunSlot.owner == owner)) await session.commit() @@ -375,7 +380,7 @@ async def hold(run_id: str, origin_kind: str | None) -> AsyncIterator[str]: _holders[run_id] = holder poll = POLL_S.get(settings.db_dialect, 5.0) try: - await _request(run_id, klass, priority) + holder.owner = await _request(run_id, klass, priority) while True: holder.woke.clear() await grant() @@ -384,7 +389,7 @@ async def hold(run_id: str, origin_kind: str | None) -> AsyncIterator[str]: holder.queued = False break if got is None: - await _request(run_id, klass, priority) + holder.owner = await _request(run_id, klass, priority) continue with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for(holder.woke.wait(), timeout=poll) @@ -393,7 +398,7 @@ async def hold(run_id: str, origin_kind: str | None) -> AsyncIterator[str]: if _holders.get(run_id) is holder: _holders.pop(run_id, None) try: - await asyncio.shield(_release(run_id)) + await asyncio.shield(_release(run_id, holder.owner or identity.owner())) except Exception: # noqa: BLE001 - 放不掉的名额下次发的时候按「本进程里已经没人用」收回 logger.exception("运行 %s 的名额没放掉，下次发名额时收回", run_id) wake_local() diff --git a/backend/app/cluster/version.py b/backend/app/cluster/version.py index 7020bd5a..be3df40b 100644 --- a/backend/app/cluster/version.py +++ b/backend/app/cluster/version.py @@ -4,7 +4,8 @@ 规矩（认领、租约、接手的说法），两边各按各的写同一批行。以前只在 README 里写「升级要先停掉全部进程」，现在代码里拦住。 - 代码版本 code_version：包版本 + 构建号（AGENTLAB_BUILD，Docker 镜像构建时写进去）或 git 提交 + 源码指纹（app/ 下全部 .py - 的内容哈希）。比的是源码指纹：同一份代码一定相同，改了一行就不同，取不到 git、没有构建号时也区分得开。 + 的内容，加上依赖锁文件 requirements.lock 的内容；Docker 镜像里也留着这个文件）。比的是源码指纹：同一份代码、同一套依赖一定相同， + 改了一行、升了一个依赖就不同，取不到 git、没有构建号时也区分得开。锁文件取不到（不是从仓库或镜像装的）时换成几个关键包的版本。 - 库结构版本 schema_version：模型里各表各列（名字、类型、可空）加迁移清单的哈希。 怎么拦：PostgreSQL 时，补列之前先看一眼（preflight）——还有活着的、版本不同的进程就不启动；它的心跳停着（多半是强杀之后留下的、 @@ -33,6 +34,30 @@ class VersionMismatch(RuntimeError): """库里还有活着的、版本不同的进程。""" +#: 依赖锁文件（仓库里 backend/requirements.lock；Docker 镜像里 /app/backend/requirements.lock） +LOCK_FILE = _APP_DIR.parent / "requirements.lock" +#: 锁文件取不到时拿这几个包的版本代替（升级时最容易带出行为差别的那几个） +KEY_PACKAGES = ("fastapi", "pydantic", "sqlalchemy", "asyncpg", "psycopg", "langgraph", "langgraph-checkpoint", + "langgraph-checkpoint-postgres", "langchain-core", "starlette", "uvicorn") + + +def dependency_marker() -> bytes: + """依赖这一份进指纹的样子：锁文件的内容；取不到就是几个关键包的版本。""" + try: + return b"lock:" + LOCK_FILE.read_bytes() + except OSError: + pass + from importlib.metadata import PackageNotFoundError, version + + parts = [] + for name in KEY_PACKAGES: + try: + parts.append(f"{name}=={version(name)}") + except PackageNotFoundError: + parts.append(f"{name}==") + return ("packages:" + ";".join(parts)).encode() + + @lru_cache(maxsize=1) def source_fingerprint() -> str: digest = hashlib.sha256() @@ -44,6 +69,8 @@ def source_fingerprint() -> str: digest.update(path.read_bytes()) except OSError: continue + # 依赖升级（锁文件变了）也算版本不同：代码一行没改、换了一版 SQLAlchemy / LangGraph 的进程同样不许混进来（复核 R-b13） + digest.update(b"\0deps\0" + dependency_marker()) return digest.hexdigest()[:12] @@ -115,10 +142,13 @@ async def _peers(conn: Any) -> list[dict[str, Any]] | None: return [dict(r) for r in rows] -def _describe(peer: dict[str, Any]) -> str: +def _describe(peer: dict[str, Any], beating: set[str] | None = None) -> str: code = peer.get("code_version") or "（升级前的版本，没有记版本）" schema = peer.get("schema_version") or "未知" - return f"{identity.describe(peer['id'])}（代码 {code}，库结构 {schema}）" + state = "" + if beating is not None: + state = "，心跳在走" if peer["id"] in beating else "，心跳停着、等租约过期" + return f"{identity.describe(peer['id'])}（代码 {code}，库结构 {schema}{state}）" async def preflight(*, observe_s: float | None = None) -> None: @@ -129,7 +159,6 @@ async def preflight(*, observe_s: float | None = None) -> None: if settings.db_dialect != "postgresql": return from app.cluster.workers import heartbeat_s, lease_s - from app.db.base import utcnow observe = 2 * heartbeat_s() + 1.0 if observe_s is None else observe_s loop = asyncio.get_running_loop() @@ -137,16 +166,17 @@ async def preflight(*, observe_s: float | None = None) -> None: first_beats: dict[str, Any] = {} warned = False while True: - mismatched = await _mismatched_alive(utcnow()) + mismatched = await _mismatched_alive() if not mismatched: return for peer in mismatched: first_beats.setdefault(peer["id"], peer.get("heartbeat_at")) - # 心跳还在走：真活着，照实报错。心跳停着的（强杀、断电留下的登记）等它们的租约过期 - beating = [p for p in mismatched if p.get("heartbeat_at") != first_beats.get(p["id"])] + # 心跳还在走：真活着，照实报错。心跳停着的（强杀、断电留下的登记）等它们的租约过期。报错时列出全部活着的、版本不同的 + # 进程（复核 R-b13：只点名心跳在走的那个，人停掉它再起又会被另一个拦下） + beating = {p["id"] for p in mismatched if p.get("heartbeat_at") != first_beats.get(p["id"])} waited = loop.time() - started if beating or waited >= lease_s() * 2 + observe: - raise VersionMismatch(_message(beating or mismatched)) + raise VersionMismatch(_message(mismatched, beating=beating)) if waited >= observe and not warned: warned = True logger.warning("库里登记着 %d 个版本不同的进程，心跳停着（多半是强杀之后留下的）：等它们的租约过期再启动", @@ -154,14 +184,19 @@ async def preflight(*, observe_s: float | None = None) -> None: await asyncio.sleep(min(0.5, max(0.1, heartbeat_s() / 2))) -async def _mismatched_alive(now: Any) -> list[dict[str, Any]]: +async def _mismatched_alive(now: Any = None) -> list[dict[str, Any]]: + """租约没过期、版本不同的别的进程。now 不给就按库的时钟（租约是按库的时钟写的，cluster/workers）。""" + from sqlalchemy import text + from app.db.base import engine async with engine.connect() as conn: peers = await _peers(conn) + if peers and now is None: + now = (await conn.execute(text("SELECT timezone('UTC', clock_timestamp())"))).scalar() if not peers: return [] - naive = now.replace(tzinfo=None) + naive = _naive(now) alive = [p for p in peers if p.get("lease_until") is not None and _naive(p["lease_until"]) > naive] return [p for p in alive if not same(p.get("code_version"), p.get("schema_version"))] @@ -170,15 +205,16 @@ def _naive(moment: Any) -> Any: return moment.replace(tzinfo=None) if getattr(moment, "tzinfo", None) else moment -def _message(peers: list[dict[str, Any]]) -> str: - listed = "；".join(_describe(p) for p in peers) +def _message(peers: list[dict[str, Any]], *, beating: set[str] | None = None) -> str: + listed = "；".join(_describe(p, beating) for p in peers) return (f"库里还有 {len(peers)} 个版本不同的 AgentLab 进程活着：{listed}。本进程是代码 {code_version()}、库结构 " f"{schema_version()}。不支持新旧版本同时运行：先停掉这些进程再启动新版本（升级步骤见 README「多进程」的「升级」）") def check_rows(rows: list[Any], now: Any) -> None: - """登记那一下（cluster/workers.register，持着事务级咨询锁）再查一遍：别的活着的进程里有版本不同的就抛 VersionMismatch。""" - naive = now.replace(tzinfo=None) + """登记那一下（cluster/workers.register，持着事务级咨询锁）再查一遍：别的活着的进程里有版本不同的就抛 VersionMismatch。 + now 是库的时钟。""" + naive = _naive(now) bad = [{"id": r.id, "code_version": r.code_version, "schema_version": r.schema_version} for r in rows if r.lease_until is not None and _naive(r.lease_until) > naive and r.id != identity.owner() and not same(r.code_version, r.schema_version)] diff --git a/backend/app/cluster/workers.py b/backend/app/cluster/workers.py index adbd6ef2..b71b0a2b 100644 --- a/backend/app/cluster/workers.py +++ b/backend/app/cluster/workers.py @@ -6,17 +6,19 @@ Liveness.alive(owner, local=…) —— 属主是别的进程：看它的租约；属主是自己：看这一行在本进程里是不是真的在跑（local）， 同一个进程里前一次起停留下的（测试里一个进程起停多次服务）照样算不在了；没有属主（升级前的行）算不在了。 -时间按应用的 utcnow（各台机器要对时，NTP）：租约以十秒计，几百毫秒的偏差无碍。 +租约用库的时钟（复核 R-b7）：PostgreSQL 上续租、登记写的是 lease_until = 库的 now() + 租约，判活比的也是库的 now()——租约由 +各台机器续、由别的机器判，用同一只钟，机器之间墙钟差多少都不起作用（cluster/leader 的当班租约同理）。SQLite 只有一个进程，用本机 +的钟。别处按「停了两个租约时长」判的时间（待办、动作的 updated_at）仍是各进程按本机时钟写的：各台机器还是要对时（NTP）。 """ from __future__ import annotations import logging import os from dataclasses import dataclass -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone from typing import Any -from sqlalchemy import delete, select, update +from sqlalchemy import DateTime, delete, func, select, update from app.cluster import identity from app.core.config import settings @@ -38,6 +40,26 @@ def heartbeat_s() -> float: return float(settings.worker_heartbeat_s) +def _postgres(session: Any) -> bool: + return session.get_bind().dialect.name == "postgresql" + + +def clock(session: Any) -> Any: + """判租约用的「现在」（写进 UTCDateTime 列、跟它比的那种：UTC、不带时区）。PostgreSQL 是库的时钟（SQL 表达式），SQLite + 是本机的（只有一个进程）。""" + if _postgres(session): + return func.timezone("UTC", func.clock_timestamp(), type_=DateTime()) + return utcnow() + + +async def db_now(session: Any) -> datetime: + """库的时钟此刻的读数（带时区的 UTC）：要在 Python 里拿它算、比的时候用（登记、当班）。""" + if _postgres(session): + value = await session.scalar(select(func.timezone("UTC", func.clock_timestamp()))) + return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value.astimezone(timezone.utc) + return utcnow() + + @dataclass(frozen=True) class Liveness: """某一刻还活着的别的进程。""" @@ -60,15 +82,14 @@ class Liveness: async def liveness(now: datetime | None = None) -> Liveness: - """现在还活着（租约没过期）的别的进程。""" + """现在还活着（租约没过期）的别的进程。now 只给测试用（不给就是库的时钟）。""" from app.db.models_cluster import ClusterWorker - now = now or utcnow() me = identity.owner() # 租约过期之后再多等一个租约时长才算不在了：卡住了但没死的进程（SIGSTOP、虚拟机暂停、库一时连不上）自己在一个租约时长时就 # 停手（cluster/fence.py），两边中间留着一段，不会它还在写、这边已经接手 - floor = now - timedelta(seconds=lease_s()) async with SessionLocal() as session: + floor = (now if now is not None else clock(session)) - timedelta(seconds=lease_s()) ids = (await session.execute(select(ClusterWorker.id).where( ClusterWorker.lease_until > floor, ClusterWorker.id != me))).scalars().all() return Liveness(frozenset(ids)) @@ -93,8 +114,8 @@ async def register(*, exclusive: bool = False) -> None: from app.cluster import version me = identity.owner() - now = utcnow() async with SessionLocal() as session: + now = await db_now(session) if exclusive: await session.execute(delete(ClusterWorker).where(ClusterWorker.id != me)) elif session.get_bind().dialect.name == "postgresql": @@ -102,6 +123,7 @@ async def register(*, exclusive: bool = False) -> None: from sqlalchemy import text await session.execute(text("SELECT pg_advisory_xact_lock(:k)"), {"k": _REGISTER_KEY}) + now = await db_now(session) # 等锁可能等了一会儿：按拿到锁之后的库时钟 others = (await session.execute(select(ClusterWorker).where(ClusterWorker.id != me))).scalars().all() version.check_rows(list(others), now) row = await session.get(ClusterWorker, me) @@ -115,12 +137,14 @@ async def register(*, exclusive: bool = False) -> None: async def renew() -> bool: - """续一拍租约。自己那一行不见了（别的进程以为我没了、清掉了）返回 False，调用方重新登记。""" + """续一拍租约（走连接池的那一种；服务里续租在专用连接上，见 cluster/service.renew）。自己那一行不见了、或者按库的时钟已经 + 过期了返回 False（不续）。""" from app.db.models_cluster import ClusterWorker - now = utcnow() async with SessionLocal() as session: - done = await session.execute(update(ClusterWorker).where(ClusterWorker.id == identity.owner()).values( + now = clock(session) + done = await session.execute(update(ClusterWorker).where( + ClusterWorker.id == identity.owner(), ClusterWorker.lease_until >= now).values( heartbeat_at=now, lease_until=now + timedelta(seconds=lease_s())) .execution_options(synchronize_session=False)) await session.commit() @@ -139,8 +163,8 @@ async def prune(now: datetime | None = None) -> int: """删掉早就过期的登记（当班进程做）：租约过期超过十倍租约时长的。它们认领的行早已被接手。""" from app.db.models_cluster import ClusterWorker - floor = (now or utcnow()) - timedelta(seconds=10 * lease_s()) async with SessionLocal() as session: + floor = (now if now is not None else clock(session)) - timedelta(seconds=10 * lease_s()) done = await session.execute(delete(ClusterWorker).where(ClusterWorker.lease_until < floor)) await session.commit() return int(done.rowcount or 0) @@ -150,8 +174,8 @@ async def listing() -> list[dict[str, Any]]: """现在登记着的进程（系统信息页、验收脚本看）。""" from app.db.models_cluster import ClusterWorker - now = utcnow() async with SessionLocal() as session: + now = await db_now(session) rows = (await session.execute(select(ClusterWorker).order_by(ClusterWorker.started_at, ClusterWorker.id))).scalars().all() return [{"id": r.id, "host": r.host, "pid": r.pid, "alive": r.lease_until > now, "me": r.id == identity.owner(), "code_version": r.code_version, "schema_version": r.schema_version, @@ -159,5 +183,5 @@ async def listing() -> list[dict[str, Any]]: "heartbeat_at": r.heartbeat_at.isoformat() if r.heartbeat_at else None} for r in rows] -__all__ = ["Liveness", "alive", "heartbeat_s", "lease_s", "listing", "liveness", "prune", "register", "renew", - "unregister"] +__all__ = ["Liveness", "alive", "clock", "db_now", "heartbeat_s", "lease_s", "listing", "liveness", "prune", "register", + "renew", "unregister"] diff --git a/backend/app/core/job_runtime.py b/backend/app/core/job_runtime.py index b050ee83..15775ee7 100644 --- a/backend/app/core/job_runtime.py +++ b/backend/app/core/job_runtime.py @@ -543,6 +543,11 @@ class JobRun: async def run(self) -> None: from app.cluster import fence + try: + await fence.ready() # 隔离之后新身份还没登记上时先等一等（cluster/fence.ready） + except fence.Fenced: + logger.error("%s %s：本进程隔离之后一直没重新登记上，不开始（由别的进程按「属主不在了」处理）", self.kind, self.job_id) + return fence.bind() try: try: diff --git a/backend/app/db/base.py b/backend/app/db/base.py index 85434064..5369bd15 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -12,6 +12,7 @@ from sqlalchemy import JSON, BigInteger, Column, DateTime, Integer, MetaData, ev from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column from sqlalchemy.pool import NullPool +from sqlalchemy.engine import make_url from sqlalchemy.types import TypeDecorator from app.core.config import settings @@ -300,11 +301,40 @@ async def dispose_session_engine() -> None: def check_pgbouncer() -> None: - """AGENTLAB_DB_PGBOUNCER=1 却没给直连地址：拒绝启动（事务模式下 LISTEN、会话级咨询锁不成立，跨进程的锁会锁不住、放不掉）。""" - if settings.db_pgbouncer and settings.db_dialect == "postgresql" and not settings.database_direct_url: + """AGENTLAB_DB_PGBOUNCER=1 时直连地址要真的是直连：没给、或者给的和应用库地址指着同一个主机端口（还是那个 PgBouncer），都拒绝 + 启动（事务模式下 LISTEN、会话级咨询锁不成立，跨进程的锁会锁不住、放不掉；复核 R-b12）。""" + if not (settings.db_pgbouncer and settings.db_dialect == "postgresql"): + return + if not settings.database_direct_url: raise RuntimeError("AGENTLAB_DB_PGBOUNCER=1（应用库地址是 PgBouncer 的事务模式）时要另给 AGENTLAB_DATABASE_DIRECT_URL：" "进程间通知（LISTEN）、跨进程的锁（会话级咨询锁）、断点库要直连 PostgreSQL（或 PgBouncer 的会话模式）。" - "见 README「用 PostgreSQL」里的三种部署") + "见 README「多进程」里的「直连与 PgBouncer」") + if _endpoint(settings.database_direct_url) == _endpoint(settings.database_url): + from app.db.urls import describe + + raise RuntimeError(f"AGENTLAB_DATABASE_DIRECT_URL 和 AGENTLAB_DATABASE_URL 指着同一个地址（{describe(settings.database_url)}）：" + "AGENTLAB_DB_PGBOUNCER=1 说应用库地址是 PgBouncer 的事务模式，直连地址要绕过它、直接连 PostgreSQL" + "（或者 PgBouncer 的会话模式那个端口）。见 README「多进程」里的「直连与 PgBouncer」") + + +def _endpoint(url: str) -> tuple[str, int]: + """地址连的是哪个主机端口（端口不写按 5432）。主机名规范化：localhost、127.x.x.x、::1 都算本机回环；IP 地址按地址比 + （::ffff:10.0.0.5 和 10.0.0.5 相同）；解析不成 IP 的按字面（小写）比，不查 DNS。""" + import ipaddress + + parsed = make_url(url) + host = (parsed.host or "").strip().strip("[]").lower().rstrip(".") + try: + addr = ipaddress.ip_address(host) + except ValueError: + addr = None + if isinstance(addr, ipaddress.IPv6Address) and addr.ipv4_mapped is not None: + addr = addr.ipv4_mapped + if host == "localhost" or host.endswith(".localhost") or (addr is not None and addr.is_loopback): + host = "loopback" + elif addr is not None: + host = str(addr) + return host, int(parsed.port or 5432) async def get_session() -> AsyncIterator[AsyncSession]: @@ -584,10 +614,19 @@ def _migrate(conn) -> None: # pragma: no cover - 同步回调 _backfill_not_null(conn, columns) +#: 存量库里 run_events 撞了号、这一次没建唯一索引时记下的结果（settings 表里；cluster: 开头的不是用户设置，api/settings 不列） +SEQ_DUPLICATES_KEY = "cluster:run_events_seq_duplicates" +#: 记下的撞号最多几处（之后启动只核对这几处还在不在） +SEQ_DUPLICATES_KEPT = 200 + + def _unique_run_event_seq(conn: Any, columns: dict[str, set[str]]) -> None: # pragma: no cover - 同步回调 """run_events 的 (run_id, seq) 补唯一索引（ux_run_events_run_seq，换掉原来不唯一的 ix_run_events_run_seq）。存量库里可能已经 有撞了号的（多进程之前的进程内计数、按需裁判的老写法）：先数一数，有就照实记日志、这一次不建（照旧用不唯一的索引），不让启动 - 失败——那几次运行的封存核对本来就对不上，删哪一条都不对，要人看。""" + 失败——那几次运行的封存核对本来就对不上，删哪一条都不对，要人看。 + + 数一次就记下结果（settings 表的 SEQ_DUPLICATES_KEY：撞了号的那几处），之后启动不再整表 GROUP BY（复核 R-b13）：只按索引核对 + 记下的那几处还撞不撞，还撞着就一句话带过；都清掉了才再整表数一遍、建索引、删掉这条记录。""" from sqlalchemy import inspect, text if "run_events" not in columns: @@ -595,16 +634,76 @@ def _unique_run_event_seq(conn: Any, columns: dict[str, set[str]]) -> None: # p names = {ix["name"] for ix in inspect(conn).get_indexes("run_events")} if "ux_run_events_run_seq" in names: return - dupes = conn.execute(text("SELECT run_id, seq, COUNT(*) FROM run_events GROUP BY run_id, seq HAVING COUNT(*) > 1 " - "LIMIT 20")).all() + known = _known_seq_duplicates(conn, columns) + if known: + still = _still_duplicated(conn, known["pairs"]) + if still: + logger.info("run_events 里还有撞了号的事件（%s 记下的 %d 处里还剩 %d 处）：(run_id, seq) 的唯一索引照旧不建，" + "处理掉之后下次启动会再核对一遍", known.get("found_at") or "之前", len(known["pairs"]), still) + return + dupes = conn.execute(text("SELECT run_id, seq FROM run_events GROUP BY run_id, seq HAVING COUNT(*) > 1 LIMIT :n"), + {"n": SEQ_DUPLICATES_KEPT + 1}).all() if dupes: runs = sorted({str(r[0]) for r in dupes}) - logger.warning("run_events 里有撞了号的事件（%d 处以上，涉及运行 %s）：这一次不建 (run_id, seq) 的唯一索引。" - "这几次运行的封存核对对不上，请人工核对后处理", len(dupes), "、".join(runs[:5])) + logger.warning("run_events 里有撞了号的事件（%d 处%s，涉及运行 %s）：这一次不建 (run_id, seq) 的唯一索引。" + "这几次运行的封存核对对不上，请人工核对后处理（之后启动只核对这几处，不再整表查）", + min(len(dupes), SEQ_DUPLICATES_KEPT), "以上" if len(dupes) > SEQ_DUPLICATES_KEPT else "", + "、".join(runs[:5])) + _remember_seq_duplicates(conn, [[str(r[0]), int(r[1])] for r in dupes[:SEQ_DUPLICATES_KEPT]], + more=len(dupes) > SEQ_DUPLICATES_KEPT) return conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS ux_run_events_run_seq ON run_events (run_id, seq)")) if "ix_run_events_run_seq" in names: conn.execute(text("DROP INDEX IF EXISTS ix_run_events_run_seq")) + if known is not None: + _forget_seq_duplicates(conn) + logger.info("run_events 里的撞号都处理掉了：补上了 (run_id, seq) 的唯一索引") + + +def _known_seq_duplicates(conn: Any, columns: dict[str, set[str]]) -> dict[str, Any] | None: # pragma: no cover + if "settings" not in columns: + return None + from sqlalchemy import select + + from app.db.models import Setting + + row = conn.execute(select(Setting.value).where(Setting.key == SEQ_DUPLICATES_KEY)).first() + value = row[0] if row else None + if not isinstance(value, dict) or not isinstance(value.get("pairs"), list): + return None + return value + + +def _still_duplicated(conn: Any, pairs: list[Any]) -> int: # pragma: no cover + """记下的那几处里还撞着的有几处（每处按 (run_id, seq) 索引查一次）。""" + from sqlalchemy import func, select + + from app.db.models import RunEvent + + left = 0 + for pair in pairs: + if not (isinstance(pair, list) and len(pair) == 2): + continue + n = conn.execute(select(func.count()).select_from(RunEvent).where( + RunEvent.run_id == str(pair[0]), RunEvent.seq == int(pair[1]))).scalar() + left += int((n or 0) > 1) + return left + + +def _remember_seq_duplicates(conn: Any, pairs: list[list[Any]], *, more: bool) -> None: # pragma: no cover + from app.db.models import Setting + + _forget_seq_duplicates(conn) + conn.execute(Setting.__table__.insert().values( + key=SEQ_DUPLICATES_KEY, value={"pairs": pairs, "more": more, "found_at": utcnow().isoformat(timespec="seconds")})) + + +def _forget_seq_duplicates(conn: Any) -> None: # pragma: no cover + from sqlalchemy import delete + + from app.db.models import Setting + + conn.execute(delete(Setting).where(Setting.key == SEQ_DUPLICATES_KEY)) #: init_db 在 PostgreSQL 上拿的会话级咨询锁的号：几个进程同时启动时，建表和迁移一个一个来（同一个号别处别用） diff --git a/backend/app/db/to_postgres.py b/backend/app/db/to_postgres.py index b861871d..6e94cec6 100644 --- a/backend/app/db/to_postgres.py +++ b/backend/app/db/to_postgres.py @@ -481,13 +481,23 @@ async def copy_checkpoints(snapshot_ckpt: Path, conninfo: str, *, dry_run: bool, # -------------------------------------------------------------------------- +def checkpoint_conninfo(target_url: str, *, checkpoint_url: str | None = None, direct_url: str | None = None) -> str: + """断点搬到哪儿（psycopg 的连接串）：给了断点库地址用它；否则进目标库——给了直连地址就走直连（复核 R-b12：目标地址是 + PgBouncer 的事务模式时，psycopg 的预备语句、建表时取的会话级锁都会出问题）。""" + if checkpoint_url: + return normalize_checkpoint_url(checkpoint_url) + return libpq_conninfo(normalize_database_url(direct_url) if direct_url else normalize_database_url(target_url)) + + async def copy_to_postgres(source: str | Path, target_url: str, *, checkpoint_url: str | None = None, - dry_run: bool = False, assume_stopped: bool = False, batch: int = DEFAULT_BATCH, - work_dir: str | Path | None = None, progress: Callable[[str], None] = print) -> Report: + direct_url: str | None = None, dry_run: bool = False, assume_stopped: bool = False, + batch: int = DEFAULT_BATCH, work_dir: str | Path | None = None, + progress: Callable[[str], None] = print) -> Report: + """direct_url：目标库的直连地址（目标地址是 PgBouncer 时给；断点走它）。""" from app.db.base import create_schema, make_engine target_url = normalize_database_url(target_url) - conninfo = normalize_checkpoint_url(checkpoint_url) if checkpoint_url else libpq_conninfo(target_url) + conninfo = checkpoint_conninfo(target_url, checkpoint_url=checkpoint_url, direct_url=direct_url) db, ckpt = resolve_source(source) if (wal := source_in_use(db, ckpt)) and not assume_stopped: raise CopyRefused(f"{wal} 还在：服务看起来还开着这个库。先停服务再导入；确认服务已停（上次是强杀，留下了 " @@ -567,6 +577,8 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("source", help="数据目录，或者 agentlab.db 的路径（断点库 checkpoints.db 在同一个目录里找）") parser.add_argument("target", nargs="?", help="目标地址 postgresql+asyncpg://…；不给就用 AGENTLAB_DATABASE_URL") parser.add_argument("--checkpoint-url", help="断点放到另一个 PostgreSQL 库；不给就用 AGENTLAB_CHECKPOINT_URL，再没有就进目标库") + parser.add_argument("--direct-url", help="目标库的直连地址（目标地址是 PgBouncer 时给，断点走它）；不给、目标又是 " + "AGENTLAB_DATABASE_URL 时用 AGENTLAB_DATABASE_DIRECT_URL") parser.add_argument("--dry-run", action="store_true", help="只核对：数来源和目标各有多少行，不建表、不写") parser.add_argument("--assume-stopped", action="store_true", help="不查服务停没停（来源旁边留着强杀后的 -wal 时用）") @@ -576,9 +588,16 @@ def main(argv: list[str] | None = None) -> int: target = args.target or settings.database_url if not target: parser.error("没有目标地址：给第二个参数，或者设 AGENTLAB_DATABASE_URL") + direct = args.direct_url or None + if direct is None and settings.database_direct_url: + try: + same = normalize_database_url(target) == settings.database_url + except ValueError: + same = False + direct = settings.database_direct_url if same else None try: report = asyncio.run(copy_to_postgres( - args.source, target, checkpoint_url=args.checkpoint_url or settings.checkpoint_url or None, + args.source, target, checkpoint_url=args.checkpoint_url or settings.checkpoint_url or None, direct_url=direct, dry_run=args.dry_run, assume_stopped=args.assume_stopped, batch=max(args.batch, 1), work_dir=args.work_dir)) except (CopyRefused, ValueError) as e: diff --git a/backend/app/engine/runner.py b/backend/app/engine/runner.py index 6599ca5b..3688318a 100644 --- a/backend/app/engine/runner.py +++ b/backend/app/engine/runner.py @@ -58,6 +58,9 @@ REJECTED_CONTINUE = "被驳回的运行不能继续运行：改好内容后重 REJECTED_RESUME = "这次运行已被驳回，没有待处理的审批，无法恢复：改好内容后重新发起" #: 服务重启时还在排队的运行改成失败时写的话（run.error 和补发的 run.failed 是同一句） QUEUED_AT_RESTART = "服务重启时这次运行仍在排队，没有可恢复的进度，请重新发起运行" +#: 本进程隔离（续不上租约）之后没能重新登记，这一段执行没有开始（cluster/fence.ready 等不到） +NOT_REREGISTERED_QUEUED = "本进程隔离后没能重新登记，这次运行没有开始执行，请重新发起运行" +NOT_REREGISTERED_RESUMABLE = "本进程隔离后没能重新登记，这次运行没有接着执行，已挂起，可从断点继续运行" class ContinueRefused(ValueError): @@ -794,7 +797,12 @@ class RunManager: # 等待从 run.interrupted 算起——以前执行一直计到 finally，收尾那一截执行和等待两边都算， # active + wait 比墙钟还长（机器一忙就差出几十毫秒） paused: float | None = None - # 这段执行属于现在这一世代（cluster/fence.py）：本进程隔离之后，它不再写 + # 这段执行属于现在这一世代（cluster/fence.py）：本进程隔离之后，它不再写。隔离之后新身份还没登记上时先等一等 + try: + await fence.ready() + except Fenced: + await self._not_reregistered(run_id) + return fence.bind() mine = identity.owner() #: 被隔离了（续不上租约）：停手，不写终态（这次运行由接手的进程按「属主不在了」处理） @@ -1001,6 +1009,45 @@ class RunManager: finally: await slot.__aexit__(None, None, None) + async def _not_reregistered(self, run_id: str) -> None: + """隔离之后新身份一直没登记上、这一段执行不开始了（fence.ready 等不到，复核 R-b15）：照实收尾，不留给别的进程按「服务 + 重启时还在排队」去标。排队中的（一步都没执行过）收成失败；已经执行过、停在断点上又要接着执行的收成中断（断点完好，可以 + 继续运行）。条件更新：状态、属主还是刚才看到的才改（别的进程先接手了就不动）。""" + from app.cluster.claims import same + + logger.error("运行 %s：本进程隔离之后一直没重新登记上，这一段执行不开始", run_id) + try: + event: RunEventModel | None = None + async with SessionLocal() as session: + seen = (await session.execute(select(Run.status, Run.owner).where(Run.id == run_id))).first() + if seen is None or seen.status not in ("queued", "running"): + return + if seen.status == "queued": + values: dict[str, Any] = {"status": "failed", "error": NOT_REREGISTERED_QUEUED} + else: + values = {"status": "interrupted", "error": NOT_REREGISTERED_RESUMABLE} + claimed = await session.execute(update(Run).where( + Run.id == run_id, Run.status == seen.status, same(Run.owner, seen.owner)) + .values(owner=None, **values).execution_options(synchronize_session=False)) + if claimed.rowcount != 1: + await session.rollback() + return + run = await session.get(Run, run_id) + bus.set_seq(run_id, await _seq_floor(session, run_id, run.last_seq if run else 0)) + event = (self._event(run_id, EventType.RUN_FAILED, data={"error": NOT_REREGISTERED_QUEUED}) + if seen.status == "queued" else + self._event(run_id, EventType.LOG, data={"level": "warn", "code": "not_reregistered", + "message": NOT_REREGISTERED_RESUMABLE})) + await _stage(session, event, run) + await session.commit() + await bus.publish(event) + await bus.close(run_id) + except Exception: # noqa: BLE001 - 收不了尾的留给别的进程按「属主不在了」处理 + logger.exception("运行 %s：收尾时出错，留给别的进程按属主不在了处理", run_id) + finally: + self._tasks.pop(run_id, None) + self._stopped_by.pop(run_id, None) + async def _stopped_in_queue(self, run_id: str) -> None: """还在等并发名额时被取消。关停：什么都不写（还是排队中 / 刚认领的样子），照「属主不在了」的规矩由下次启动或别的 进程接手（排队的标失败、认领了的标中断）；人点了停止：收成 cancelled，同执行中停止（run.cancelled 带上是谁停的、封存）。""" @@ -1786,15 +1833,19 @@ run_manager = RunManager() def _fenced_saver(saver: BaseCheckpointSaver) -> BaseCheckpointSaver: - """断点库的写也要过 cluster/fence：本进程续不上租约之后，手上的执行不再写断点（接手的进程会从断点接着跑，两边都写就分叉了）。""" + """断点库的写也要过 cluster/fence：本进程续不上租约之后，手上的执行不再写断点（接手的进程会从断点接着跑，两边都写就分叉了）。 + + 包成协程、在协程里查（复核 R-b8）：LangGraph 在节点的完成回调里调 aput_writes 拿协程交给后台任务——同步包装在回调里就抛了， + 每次隔离 asyncio 都记一条「Exception in callback … FuturesDict.on_done」；在协程里抛，异常留在写断点的那个任务里，由 LangGraph + 照常收走。""" for name in ("aput", "aput_writes"): original = getattr(saver, name, None) if original is None: continue - def guarded(*args: Any, __original: Any = original, **kwargs: Any) -> Any: + async def guarded(*args: Any, __original: Any = original, **kwargs: Any) -> Any: fence.check() - return __original(*args, **kwargs) + return await __original(*args, **kwargs) setattr(saver, name, guarded) return saver @@ -1816,8 +1867,9 @@ def _register_requests() -> None: @requests.handle("stop", "run") async def _stop(run_id: str, actor: str | None, _payload: dict[str, Any]) -> None: - if run_manager.is_active(run_id): - await run_manager.cancel(run_id, actor=actor) + # 不在本进程执行的也照停止的规矩办（cancel 自己看属主）：在别的活着的进程里执行就转给它，属主不在了就接手、放弃，停在 + # 断点上的收成 cancelled，已经结束的什么都不做。隔离之后转到新身份名下的停止请求就是这样落下的（复核 R-b10） + await run_manager.cancel(run_id, actor=actor) @requests.handle("note", "run") async def _note(run_id: str, _actor: str | None, payload: dict[str, Any]) -> None: diff --git a/backend/app/ledger/issue.py b/backend/app/ledger/issue.py index e2bddd08..ff742d6c 100644 --- a/backend/app/ledger/issue.py +++ b/backend/app/ledger/issue.py @@ -379,17 +379,18 @@ async def issue(run_id: str, workflow_id: str, version: int, *, principal: Any, from app.cluster import quota limit = int((await ledger_settings(session))["issue_daily_limit"]) - if not await quota.take(QUOTA, limit, seed=issued_today): + day = await quota.take(QUOTA, limit, seed=issued_today) + if day is None: raise IssueRefused(429, OVER_LIMIT.format(limit=limit), DAILY_LIMIT) try: launched = await launch_run(session, workflow_id=workflow_id, run_class="formal", version=int(version), input_payload=plan.inputs, actor=actor, origin_kind=ORIGIN, origin=origin, expect_hash=plan.expect_hash) except LaunchRefused as e: - await quota.give_back(QUOTA) + await quota.give_back(QUOTA, day=day) raise IssueRefused(e.status_code, e.message, e.code or REFUSED) from e except BaseException: - await quota.give_back(QUOTA) + await quota.give_back(QUOTA, day=day) raise return launched, True diff --git a/backend/app/scheduler/loop.py b/backend/app/scheduler/loop.py index 7a7fba52..6c2a9f14 100644 --- a/backend/app/scheduler/loop.py +++ b/backend/app/scheduler/loop.py @@ -201,6 +201,7 @@ class _Runner: async def _loop(self, interval: float, stopping: asyncio.Event) -> None: from app.cluster import leading + from app.cluster.fence import Fenced from app.scheduler import webhook self._poked = asyncio.Event() @@ -218,7 +219,7 @@ class _Runner: self.notify = asyncio.create_task(_notify(webhook), name="agentlab-scheduler-notify") except asyncio.CancelledError: raise - except Exception: # noqa: BLE001 - 这一拍出错不该让循环停下 + except (Exception, Fenced): # noqa: BLE001 - 这一拍出错不该让循环停下（Fenced 是 BaseException：复核 R-b9） logger.exception("定时任务调度这一拍出错，下一拍照常进行") await self._wait(stopping, interval) diff --git a/backend/tests/test_cluster.py b/backend/tests/test_cluster.py index 89c708cd..8b7aea5d 100644 --- a/backend/tests/test_cluster.py +++ b/backend/tests/test_cluster.py @@ -7,6 +7,7 @@ import asyncio import os import subprocess import sys +import time from datetime import timedelta from pathlib import Path @@ -337,6 +338,63 @@ async def test_stopping_a_run_whose_owner_is_gone_takes_it_over_and_abandons_it( assert row.status == "cancelled" and row.end_reason == "stopped" +@pytest.mark.skipif(app_settings.db_dialect != "postgresql", reason="SQLite 只有一个进程，名额在进程里（复核 R-b3）") +async def test_a_slot_is_released_under_the_identity_it_was_taken_with(monkeypatch): + """复核 R-b10：拿着名额的时候隔离换了身份，放名额照样删得到那一行（按要名额时的身份删），不用等旧身份过期才收回。""" + _policy(monkeypatch, 2, interactive=(2, 40)) + monkeypatch.setattr(identity, "OWNER", identity.OWNER) + async with slots.hold("rb10-run", "chat"): + identity.regenerate() + async with SessionLocal() as session: + left = list((await session.execute(select(RunSlot).where(RunSlot.run_id == "rb10-run"))).scalars()) + assert left == [] + + +async def test_registering_again_tidies_up_what_the_retired_identity_left(monkeypatch): + """复核 R-b10：隔离之后新身份登记上，旧身份的登记、名额、同源队的位置删掉，当班租约交出，发给它的请求转到新身份名下照办—— + 停止请求对停在断点上的那次运行照停止的规矩收成 cancelled。""" + from app.cluster import requests, version + from app.cluster import service as service_mod + from app.cluster.fence import now as fence_now + + monkeypatch.setattr(identity, "OWNER", f"{identity.host()}:{os.getpid()}:retired0") # 测试完换回原来的身份 + old = identity.owner() + now = utcnow() + async with SessionLocal() as session: + session.add(ClusterWorker(id=old, host="this-host", pid=1, started_at=now, heartbeat_at=now, + lease_until=now + timedelta(seconds=60), code_version=version.code_version(), + schema_version=version.schema_version())) + session.add(RunSlot(run_id="retired-slot", klass="interactive", priority=40, owner=old, granted=True, + requested_at=now)) + session.add(LaneTicket(source_id="src-retired", kind="profile", job_id="retired-job", owner=old, granted=True)) + session.add(ClusterLease(name="test-retired", owner=old, lease_until=now + timedelta(seconds=60), epoch=1)) + run = Run(status="interrupted", owner=old, graph={}, input={}) + session.add(run) + await session.flush() + session.add(ClusterRequest(owner=old, kind="stop", target="run", target_id=run.id, actor="张工", payload={})) + await session.commit() + svc = service_mod._Service() + svc.retired = [old] + identity.regenerate() + svc.pending_since = fence_now() + await svc.register_again() + new = identity.owner() + async with SessionLocal() as session: + ids = set((await session.execute(select(ClusterWorker.id))).scalars()) + slot_rows = list((await session.execute(select(RunSlot).where(RunSlot.owner == old))).scalars()) + lane_rows = list((await session.execute(select(LaneTicket).where(LaneTicket.owner == old))).scalars()) + lease = await session.get(ClusterLease, "test-retired") + sent = list((await session.execute(select(ClusterRequest.owner).where(ClusterRequest.target_id == run.id))).scalars()) + assert old not in ids and new in ids, "旧身份的登记删掉：别的进程马上按属主不在了接手" + assert slot_rows == [] and lane_rows == [] and lease.owner is None + assert sent == [new], "发给旧身份的请求转到新身份名下" + await requests.drain() + await requests.settle() + async with SessionLocal() as session: + row = await session.get(Run, run.id) + assert row.status == "cancelled" and row.end_reason == "stopped", "停止请求照停止的规矩落下了" + + async def test_leadership_is_handed_over_when_the_lease_runs_out(): me = leader_mod.Leadership("test-lead") assert await me.claim() and me.held() @@ -355,6 +413,87 @@ async def test_leadership_is_handed_over_when_the_lease_runs_out(): assert (await leader_mod.current("test-lead"))["owner"] is None +def test_the_fence_judges_by_the_sleep_aware_clock_and_falls_back_to_the_wall(monkeypatch): + """复核 R-b7、R-b14：取得到计入休眠的钟时只看它——它走过两个租约（合盖、挂起）算续晚了，墙钟往前、往后跳一小时不算（只算 + 两只钟对不上，jumped）；取不到计入休眠的钟时，monotonic 和墙钟里走得多的那只说了算（墙钟当后备）。""" + from app.cluster import fence + + clocks = {"mono": 1000.0, "wall": 5000.0} + monkeypatch.setattr(fence, "_clock", lambda: clocks["mono"]) + monkeypatch.setattr(fence, "_wall", lambda: clocks["wall"]) + monkeypatch.setattr(fence, "SLEEP_AWARE", True) + fence.activate() + try: + assert not fence.stale() and not fence.jumped() + for jump in (3600.0, -3600.0): # NTP 一步校时、手动改时间 + fence.renewed() + clocks["wall"] += jump + assert not fence.stale() and fence.jumped(), "只是墙钟跳了：不算续晚了" + fence.renewed() + clocks["mono"] += 2 * workers.lease_s() # 计入休眠的钟走过两个租约（合盖） + assert fence.stale() + fence.renewed() + clocks["mono"] += workers.lease_s() / 2 + clocks["wall"] += workers.lease_s() / 2 + assert not fence.stale() and not fence.jumped() + + monkeypatch.setattr(fence, "SLEEP_AWARE", False) # 取不到计入休眠的钟：monotonic 不计休眠，墙钟当后备 + fence.renewed() + clocks["wall"] += 2 * workers.lease_s() + assert fence.stale() + finally: + fence.deactivate() + + +@pytest.mark.skipif(app_settings.db_dialect != "postgresql", reason="SQLite 只有一个进程，租约用本机的钟") +async def test_leases_use_the_database_clock_not_the_local_wall_clock(monkeypatch): + """复核 R-b7：本机墙钟快了一个小时——别的进程照样算活着、它拿着的当班租约照样抢不走；按库的时钟自己的租约已经过期时，续租 + 不再续这个身份，而是隔离。""" + from app.cluster import fence + from app.cluster.service import service + from app.db.base import utcnow as real_utcnow + import app.db.base as db_base + + await _peer(PEER) + lead = leader_mod.Leadership("test-skew") + async with SessionLocal() as session: + session.add(ClusterLease(name="test-skew", owner=PEER, lease_until=real_utcnow() + timedelta(seconds=60), epoch=1)) + await session.commit() + skewed = lambda: real_utcnow() + timedelta(hours=1) # noqa: E731 + monkeypatch.setattr(workers, "utcnow", skewed) + monkeypatch.setattr(db_base, "utcnow", skewed) + assert PEER in (await workers.liveness()).others, "按库的时钟它的租约没过期" + assert not await lead.claim(), "当班租约按库的时钟判：抢不走" + async with SessionLocal() as session: + await session.execute(delete(ClusterWorker).where(ClusterWorker.id == PEER)) # 它没记版本：登记时会被版本守卫拦下 + await session.execute(delete(ClusterLease).where(ClusterLease.name == "test-skew")) + await session.commit() + + await workers.register() + fence.activate() + before = fence.generation() + try: + assert await service.renew(), "按库的时钟没过期：续上" + monkeypatch.setattr(fence, "SLEEP_AWARE", True) + for jump in (3600.0, -3600.0): # 单进程、墙钟往前往后跳一小时（复核 R-b14） + monkeypatch.setattr(fence, "_wall", lambda jump=jump: time.time() + jump) + assert not fence.stale() and fence.jumped() + fence.check() + assert await service.renew() and fence.generation() == before, "库的时钟说租约还在：续上、不隔离" + assert not fence.jumped(), "续上之后两只钟重新对齐" + monkeypatch.setattr(fence, "_wall", time.time) + async with SessionLocal() as session: + row = await session.get(ClusterWorker, identity.owner()) + row.lease_until = real_utcnow() - timedelta(seconds=1) # 库的时钟已经过了它的租约，本机的两只钟都没觉得 + await session.commit() + assert not await service.renew() and fence.generation() == before + 1, "过期了不续这个身份：隔离" + finally: + fence.deactivate() + if service._lease_conn is not None: + await service._lease_conn.close() + service._lease_conn = None + + async def test_the_daily_cap_is_not_overrun_by_concurrent_charges(): from app.engine.spend import Ledger @@ -405,12 +544,27 @@ async def test_a_daily_quota_is_taken_up_to_the_limit_by_concurrent_callers(): return 2 # 今天已经做过 2 次 got = await asyncio.gather(*(quota.take(name, 5, seed=seed) for _ in range(10))) - assert sum(got) == 3 and await quota.used(name) == 5 + assert sum(1 for g in got if g) == 3 and await quota.used(name) == 5 assert not await quota.take(name, 5) await quota.give_back(name) assert await quota.take(name, 5) and not await quota.take(name, 5) +async def test_a_quota_is_given_back_to_the_day_it_was_taken(monkeypatch): + """复核 R-b11：23:59:59 占、00:00:01 发起失败——退回前一天那一份，第二天的不动。""" + from app.cluster import quota + + name = f"test-quota-{new_id()[:8]}" + days = iter(["2026-10-18", "2026-10-19", "2026-10-19", "2026-10-19"]) + monkeypatch.setattr(quota, "today", lambda: next(days)) + day = await quota.take(name, 5) + assert day == "2026-10-18" + assert await quota.take(name, 5) == "2026-10-19" # 过了零点又占了一份 + await quota.give_back(name, day=day) # 前一天占的那一份发起失败 + assert await quota.used(name, day="2026-10-18") == 0 + assert await quota.used(name, day="2026-10-19") == 1, "第二天的那一份不动" + + def test_a_daily_quota_holds_across_two_processes(): """两个进程同时各占 15 次、上限 10：合计正好 10 次。""" name = f"test-quota-{new_id()[:8]}" @@ -419,7 +573,7 @@ def test_a_daily_quota_holds_across_two_processes(): "from app.cluster import quota\n" "async def main():\n" f" got = await asyncio.gather(*(quota.take({name!r}, 10) for _ in range(15)))\n" - " print(sum(got))\n" + " print(sum(1 for g in got if g))\n" "asyncio.run(main())\n") env = {**os.environ, "PYTHONPATH": str(backend), "AGENTLAB_DB_POOL_SIZE": "5"} procs = [subprocess.Popen([sys.executable, "-c", script], cwd=str(backend), env=env, stdout=subprocess.PIPE, @@ -504,6 +658,109 @@ async def test_preflight_refuses_a_beating_peer_and_waits_out_a_stale_one(): await workers.register() +@pytest.mark.skipif(app_settings.db_dialect != "postgresql", reason="preflight 只在 PostgreSQL 上看") +async def test_preflight_names_every_live_process_of_another_version(): + """复核 R-b13：一个心跳在走、一个心跳停着——报错两个都列出来（只点名先看到心跳的那个，人停掉它再起还会被另一个拦下）。""" + from app.cluster import version + + await _peer_version(PEER, "0.0.9+old+aaaaaaaaaaaa") + await _peer_version(GHOST, "0.0.8+older+bbbbbbbbbbbb") + + async def beat() -> None: + for _ in range(15): + await asyncio.sleep(0.1) + async with SessionLocal() as session: + row = await session.get(ClusterWorker, PEER) + row.heartbeat_at = utcnow() + await session.commit() + + beating = asyncio.create_task(beat()) + with pytest.raises(version.VersionMismatch) as e: + await version.preflight(observe_s=0.5) + await beating + text = str(e.value) + assert "2 个版本不同" in text and "aaaaaaaaaaaa" in text and "bbbbbbbbbbbb" in text + assert "心跳在走" in text and "心跳停着" in text + + +def test_the_fingerprint_changes_with_the_dependency_lock(monkeypatch, tmp_path): + """复核 R-b13：依赖锁文件变了（升了依赖）源码指纹跟着变；锁文件取不到时换成几个关键包的版本。""" + from app.cluster import version + + lock = tmp_path / "requirements.lock" + monkeypatch.setattr(version, "LOCK_FILE", lock) + + def fingerprint() -> str: + version.source_fingerprint.cache_clear() + return version.source_fingerprint() + + try: + lock.write_text("sqlalchemy==2.0.40\n") + first = fingerprint() + assert fingerprint() == first + lock.write_text("sqlalchemy==2.0.41\n") + assert fingerprint() != first, "升了一个依赖：指纹不同" + lock.unlink() + assert version.dependency_marker().startswith(b"packages:") and b"sqlalchemy==" in version.dependency_marker() + assert fingerprint() not in (first,) + finally: + version.source_fingerprint.cache_clear() + + +async def test_a_known_seq_collision_is_not_rescanned_on_every_start(tmp_path, caplog): + """复核 R-b13：存量库里有撞了号的事件——第一次启动整表数一遍、记下结果、不建唯一索引；之后启动只核对记下的那几处，不再整表 + GROUP BY；撞号处理掉之后再数一遍、建索引、删掉记录。用一个单独的 SQLite 库（两种库的测试里都跑）。""" + import logging + + from sqlalchemy import event, inspect, text + + from app.db import base as db_base + from app.db.models import Setting + + target = db_base.make_engine(f"sqlite+aiosqlite:///{tmp_path / 'seq.db'}") + scans: list[str] = [] + event.listen(target.sync_engine, "before_cursor_execute", + lambda conn, cursor, statement, *a: scans.append(statement) if "GROUP BY run_id, seq" in statement else None) + try: + await db_base.create_schema(target) + from sqlalchemy.ext.asyncio import async_sessionmaker + + async with target.begin() as conn: + await conn.execute(text("DROP INDEX ux_run_events_run_seq")) + await conn.execute(text("CREATE INDEX ix_run_events_run_seq ON run_events (run_id, seq)")) + async with async_sessionmaker(target, expire_on_commit=False)() as session: + session.add(Run(id="r1", status="succeeded", graph={}, input={})) + await session.flush() + session.add_all([RunEvent(id=eid, run_id="r1", seq=seq, type="log", ts=0, data={}) + for eid, seq in (("e1", 1), ("e2", 1), ("e3", 2))]) + await session.commit() + + async def start() -> set[str]: + async with target.begin() as conn: + def run(sync_conn): + cols = {t: {c["name"] for c in inspect(sync_conn).get_columns(t)} for t in ("run_events", "settings")} + db_base._unique_run_event_seq(sync_conn, cols) + return {ix["name"] for ix in inspect(sync_conn).get_indexes("run_events")} + return await conn.run_sync(run) + + caplog.set_level(logging.INFO, logger="app.db.base") + assert "ux_run_events_run_seq" not in await start() and len(scans) == 1 + async with target.connect() as conn: + marker = (await conn.execute(text(f"SELECT value FROM settings WHERE key = '{db_base.SEQ_DUPLICATES_KEY}'"))).scalar() + assert '"r1"' in str(marker) and "1" in str(marker) + assert "ux_run_events_run_seq" not in await start() and len(scans) == 1, "第二次启动不再整表 GROUP BY" + assert any("还有撞了号的事件" in r.getMessage() for r in caplog.records) + async with target.begin() as conn: + await conn.execute(text("DELETE FROM run_events WHERE id = 'e2'")) + assert "ux_run_events_run_seq" in await start() and len(scans) == 2, "处理掉之后再数一遍、建上索引" + async with target.connect() as conn: + left = (await conn.execute(text(f"SELECT COUNT(*) FROM settings WHERE key = '{db_base.SEQ_DUPLICATES_KEY}'"))).scalar() + assert left == 0 + assert Setting.__tablename__ == "settings" + finally: + await target.dispose() + + @pytest.mark.sqlite_only async def test_an_idle_single_process_writes_nothing_for_coordination(monkeypatch): """复核 R-b3：SQLite（只有一个进程）空闲时协调不写库——不续租约、不写当班、不周期接手；要名额、放名额也不落库。用另一条连接 diff --git a/backend/tests/test_cluster_fence.py b/backend/tests/test_cluster_fence.py index f963e9c6..f3146d5f 100644 --- a/backend/tests/test_cluster_fence.py +++ b/backend/tests/test_cluster_fence.py @@ -4,7 +4,8 @@ from __future__ import annotations import asyncio -import time +import contextlib +import os import pytest from sqlalchemy import func, select @@ -99,7 +100,7 @@ async def _stall_and_take_over(slow) -> str: await asyncio.wait_for(slow["entered"].wait(), timeout=15) await _checkpointed_before_slow(run.id) before = fence.generation() - fence.renewed(time.monotonic() - lease_s() - 1) # 卡了一个多租约时长 + fence.renewed(ago=lease_s() + 1) # 卡了一个多租约时长 await _taken_over(run.id) slow["gate"].set() # 缓过来，手上的执行接着往下走 for _ in range(200): @@ -111,8 +112,17 @@ async def _stall_and_take_over(slow) -> str: return run.id -async def test_a_stalled_executor_does_not_overwrite_the_takeover(slow): +async def test_a_stalled_executor_does_not_overwrite_the_takeover(slow, caplog): + import gc + import logging + + caplog.set_level(logging.ERROR, logger="asyncio") run_id = await _stall_and_take_over(slow) + await asyncio.sleep(0.1) + gc.collect() + await asyncio.sleep(0) + noisy = [r.getMessage() for r in caplog.records if r.name == "asyncio"] + assert noisy == [], f"隔离时 asyncio 不该记回调异常、没人取的任务异常（复核 R-b8）：{noisy}" async with SessionLocal() as session: row = await session.get(Run, run_id) assert row.status == "interrupted" and row.owner is None, "接手之后的样子没被盖掉" @@ -139,6 +149,50 @@ async def test_a_run_continued_elsewhere_finishes_once_and_the_stalled_side_stay await _no_repeated_seq(run_id) +@pytest.mark.parametrize("jump", [3600.0, -3600.0]) +async def test_a_wall_clock_jump_does_not_fence_a_running_run(slow, monkeypatch, jump): + """复核 R-b14：执行到一半本机墙钟往前 / 往后跳一小时（NTP 一步校时、手动改时间）——计入休眠的钟没问题，不隔离，运行照常跑完。""" + import time as _time + + monkeypatch.setattr(fence, "SLEEP_AWARE", True) + run = await run_manager.start(graph=GRAPH, input_payload={"question": "q"}) + await asyncio.wait_for(slow["entered"].wait(), timeout=15) + before = fence.generation() + monkeypatch.setattr(fence, "_wall", lambda: _time.time() + jump) + slow["gate"].set() + for _ in range(300): + async with SessionLocal() as session: + row = await session.get(Run, run.id) + if row.status in ("succeeded", "failed", "interrupted") and not run_manager.is_active(run.id): + break + await asyncio.sleep(0.05) + assert row.status == "succeeded", row.error + assert fence.generation() == before, "没有隔离" + + +async def test_a_run_that_cannot_start_after_fencing_is_failed_with_the_real_reason(monkeypatch): + """复核 R-b15：隔离之后新身份一直没登记上，新发起的运行等不到（fence.ready）——这一行照实收成失败（写明是隔离后没能重新 + 登记，不是「服务重启时还在排队」），run.failed 记下来，执行任务也从 run_manager 里摘掉。""" + from app.cluster import workers + from app.engine.runner import NOT_REREGISTERED_QUEUED + + monkeypatch.setattr(workers, "lease_s", lambda: 0.3) + fence.trip("测试") # 隔离了，新身份还没登记上 + run = await run_manager.start(graph=GRAPH, input_payload={"question": "q"}) + for _ in range(100): + if run.id not in run_manager._tasks: + break + await asyncio.sleep(0.05) + assert run.id not in run_manager._tasks, "执行任务摘掉了" + async with SessionLocal() as session: + row = await session.get(Run, run.id) + assert row.status == "failed" and row.error == NOT_REREGISTERED_QUEUED and row.owner is None + events = await _events(run.id) + assert [e.type for e in events if e.type == "run.failed"] == ["run.failed"] + assert events[-1].data["error"] == NOT_REREGISTERED_QUEUED + fence.renewed() + + async def test_the_final_write_checks_the_owner(): async with SessionLocal() as session: run = Run(status="running", owner="peer-host:1:abcd", graph={}, input={}) @@ -159,3 +213,108 @@ async def test_a_repeated_seq_is_refused_by_the_database(): RunEvent(id=new_id(), run_id=run.id, seq=1, type="log", ts=0, data={})]) with pytest.raises(IntegrityError): await session.commit() + + +# ---------------- 复核 R-b9：谁会被打停、谁不会 ---------------- + + +async def test_only_executions_bound_to_a_generation_are_stopped_by_check(): + """续晚之后第一个问 check() 的不一定是运行本身：没绑定世代的调用者（HTTP 处理、调度那一拍）只隔离、不抛；绑定了的抛。 + 隔离之后新身份还没登记上，新开始的执行在 ready() 上等，等不到抛 Fenced；登记上了就放行。""" + gate = asyncio.Event() + seen: list[str] = [] + + async def bound_writer() -> None: + fence.bind() + await gate.wait() + try: + fence.check() + seen.append("wrote") + except fence.Fenced: + seen.append("fenced") + + task = asyncio.create_task(bound_writer()) + await asyncio.sleep(0) + before = fence.generation() + fence.renewed(ago=lease_s() + 1) + fence.check() # 没绑定世代：只隔离 + assert fence.generation() == before + 1 and fence.pending() + gate.set() + await task + assert seen == ["fenced"], "绑定在隔离之前的那一段执行：不许再写" + + with pytest.raises(fence.Fenced): + await fence.ready(timeout=0.1) # 新身份还没登记上：新开始的执行等不到就不开始 + waiting = asyncio.create_task(fence.ready(timeout=5)) + await asyncio.sleep(0.1) + assert not waiting.done() + fence.renewed() # 登记上了 + await asyncio.wait_for(waiting, timeout=2) + fence.bind() + fence.check() # 新世代照常 + + +async def test_the_scheduler_loop_survives_a_fenced_tick(monkeypatch): + from app.scheduler import loop as sloop + + calls = {"n": 0} + + async def fire_due(now=None): + calls["n"] += 1 + if calls["n"] == 1: + raise fence.Fenced() + return [] + + monkeypatch.setattr(sloop, "fire_due", fire_due) + await sloop.runner.start(interval=0.02) + try: + for _ in range(200): + if calls["n"] >= 3: + break + await asyncio.sleep(0.02) + assert calls["n"] >= 3 and sloop.runner.running(), "被 Fenced 打到的那一拍之后照常往下" + finally: + await sloop.runner.stop() + + +async def test_the_lease_task_keeps_going_when_registering_again_fails(monkeypatch): + """隔离之后重新登记：库一抖、版本守卫拦下都不让续租的任务退出，退避着重试；登记上了照常续租。""" + from app.cluster import identity, service as service_mod, workers + from app.cluster.version import VersionMismatch + + monkeypatch.setattr(identity, "OWNER", f"{identity.host()}:{os.getpid()}:retired1") # 测试完换回原来的身份 + monkeypatch.setattr(workers, "heartbeat_s", lambda: 0.01) + monkeypatch.setattr(workers, "lease_s", lambda: 5.0) + svc = service_mod._Service() + tries = {"register": 0, "renew": 0} + + async def register(*, exclusive: bool = False) -> None: # noqa: ARG001 + tries["register"] += 1 + if tries["register"] == 1: + raise VersionMismatch("库里还有版本不同的进程活着") + if tries["register"] == 2: + raise OSError("库一时连不上") + + async def renew() -> bool: + tries["renew"] += 1 + return True + + monkeypatch.setattr(workers, "register", register) + monkeypatch.setattr(svc, "renew", renew) + old = identity.owner() + fence.trip("测试") + svc._on_fenced() + assert identity.owner() != old and svc.pending_since is not None + task = asyncio.create_task(svc._lease_loop()) + try: + for _ in range(300): + if tries["renew"] >= 2: + break + await asyncio.sleep(0.01) + assert not task.done(), "续租的任务还在" + assert tries["register"] == 3 and svc.pending_since is None and not fence.pending() + assert tries["renew"] >= 2, "登记上之后照常续租" + finally: + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task diff --git a/backend/tests/test_db_direct.py b/backend/tests/test_db_direct.py index 552a0c40..13c52dd0 100644 --- a/backend/tests/test_db_direct.py +++ b/backend/tests/test_db_direct.py @@ -67,3 +67,67 @@ def test_transaction_mode_without_a_direct_url_refuses_to_start(monkeypatch): monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None).database_direct_url) base.check_pgbouncer() + + +def test_transaction_mode_with_a_direct_url_that_is_the_same_endpoint_refuses_to_start(monkeypatch): + """复核 R-b12：直连地址和应用库地址规范化之后连的是同一个主机端口（还是那个 PgBouncer）——同样拒绝，说清楚。""" + from app.db import base + + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "db_pgbouncer", True) + for same in (APP, "postgresql+asyncpg://other:pw@PGBOUNCER.internal:6432/agentlab?ssl=require"): + monkeypatch.setattr(settings, "database_direct_url", + Settings(database_url=APP, database_direct_url=same, _env_file=None).database_direct_url) + with pytest.raises(RuntimeError, match="指着同一个地址"): + base.check_pgbouncer() + monkeypatch.setattr(settings, "database_direct_url", Settings( + database_url=APP, database_direct_url="postgresql://agentlab:secret@pgbouncer.internal:5432/agentlab", + _env_file=None).database_direct_url) + base.check_pgbouncer() # 同一台机器上的另一个端口（直连 PostgreSQL 或会话模式）：放行 + + +def test_copying_checkpoints_goes_through_the_direct_url(monkeypatch, tmp_path): + """复核 R-b12：copy-to-postgres 搬断点缺省走直连地址（设了的话）；单独给了断点库地址就用它。""" + from app.db import to_postgres + + assert "db.internal:5432" in to_postgres.checkpoint_conninfo(APP, direct_url=DIRECT) + assert "pgbouncer.internal:6432" in to_postgres.checkpoint_conninfo(APP) + assert "ckpt.internal" in to_postgres.checkpoint_conninfo(APP, checkpoint_url="postgresql://u:p@ckpt.internal/c", + direct_url=DIRECT) + seen: dict = {} + + async def fake_copy(source, target, **kw): + seen.update(kw, target=target) + raise to_postgres.CopyRefused("测试：到这里为止") + + monkeypatch.setattr(to_postgres, "copy_to_postgres", fake_copy) + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "database_direct_url", + Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None).database_direct_url) + monkeypatch.setattr(settings, "checkpoint_url", "") + assert to_postgres.main([str(tmp_path)]) == 2 + assert seen["direct_url"] == settings.database_direct_url, "目标就是应用库：断点走它的直连地址" + assert to_postgres.main([str(tmp_path), "postgresql://u:p@elsewhere.internal:5432/x"]) == 2 + assert seen["direct_url"] is None, "目标是别的库：不套用应用库的直连地址" + + +def test_loopback_spellings_count_as_the_same_endpoint(monkeypatch): + """R-b12 的尾巴：localhost、127.0.0.1、::1 是同一台机器——同一个端口时照样拒绝；IP 按地址比，解析不成 IP 的按字面比。""" + from app.db import base + + def norm(url: str) -> str: + return Settings(database_url=url, _env_file=None).database_url + + monkeypatch.setattr(settings, "db_pgbouncer", True) + for app_url, direct in (("postgresql://u:p@localhost:6432/db", "postgresql://u:p@127.0.0.1:6432/db"), + ("postgresql://u:p@127.0.0.1:6432/db", "postgresql://u:p@[::1]:6432/db"), + ("postgresql://u:p@10.0.0.5:6432/db", "postgresql://u:p@[::ffff:10.0.0.5]:6432/db")): + monkeypatch.setattr(settings, "database_url", norm(app_url)) + monkeypatch.setattr(settings, "database_direct_url", norm(direct)) + with pytest.raises(RuntimeError, match="指着同一个地址"): + base.check_pgbouncer() + monkeypatch.setattr(settings, "database_url", norm("postgresql://u:p@localhost:6432/db")) + monkeypatch.setattr(settings, "database_direct_url", norm("postgresql://u:p@127.0.0.1:5432/db")) + base.check_pgbouncer() # 本机另一个端口：放行 + monkeypatch.setattr(settings, "database_direct_url", norm("postgresql://u:p@db.internal:6432/db")) + base.check_pgbouncer() # 名字不同、不查 DNS：按字面算不同