From 01af07b193b68d0a39cf43ae9631c0302da2b671 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 19 Oct 2026 12:57:26 +0800 Subject: [PATCH] =?UTF-8?q?=E5=87=BA=E5=85=B7=E3=80=81=E6=8C=89=E9=9C=80?= =?UTF-8?q?=E8=A7=86=E5=9B=BE=E7=9A=84=E6=AF=8F=E6=97=A5=E4=B8=8A=E9=99=90?= =?UTF-8?q?=E6=94=B9=E6=88=90=E5=85=88=E5=8D=A0=E9=A2=9D=E5=BA=A6=E5=86=8D?= =?UTF-8?q?=E5=8F=91=E8=B5=B7=EF=BC=88=E7=94=A8=E6=88=B7=E6=8B=8D=E6=9D=BF?= =?UTF-8?q?=EF=BC=89=EF=BC=9A=E6=96=B0=E5=A2=9E=20cluster/quota=E2=80=94?= =?UTF-8?q?=E2=80=94=E9=A2=9D=E5=BA=A6=E8=AE=B0=E5=9C=A8=E5=BA=93=E9=87=8C?= =?UTF-8?q?=EF=BC=88rate=5Fcounters=20=E6=8C=89=E5=90=8D=E5=AD=97=20+=20?= =?UTF-8?q?=E6=9C=AC=E5=9C=B0=E6=97=A5=E6=9C=9F=E4=B8=80=E8=A1=8C=EF=BC=89?= =?UTF-8?q?=EF=BC=8C=E5=8D=A0=E4=B8=80=E4=BB=BD=E6=98=AF=E4=B8=80=E5=8F=A5?= =?UTF-8?q?=E5=B8=A6=E6=9D=A1=E4=BB=B6=E7=9A=84=20UPDATE=EF=BC=88=E6=B2=A1?= =?UTF-8?q?=E5=88=B0=E4=B8=8A=E9=99=90=E6=89=8D=E5=8A=A0=E4=B8=80=EF=BC=89?= =?UTF-8?q?=EF=BC=8C=E4=B8=A4=E7=A7=8D=E5=BA=93=E4=B8=8A=E9=83=BD=E6=98=AF?= =?UTF-8?q?=E5=8E=9F=E5=AD=90=E7=9A=84=EF=BC=8C=E8=BF=99=E4=B8=80=E5=A4=A9?= =?UTF-8?q?=E7=AC=AC=E4=B8=80=E6=AC=A1=E5=8D=A0=E6=8C=89=E5=BA=93=E9=87=8C?= =?UTF-8?q?=E5=B7=B2=E6=9C=89=E7=9A=84=E8=AE=B0=E5=BD=95=E8=B5=B7=E7=AE=97?= =?UTF-8?q?=EF=BC=8C=E5=8F=91=E8=B5=B7=E8=A2=AB=E6=8B=92=E6=88=96=E5=87=BA?= =?UTF-8?q?=E9=94=99=E5=B0=B1=E9=80=80=E5=9B=9E=E3=80=82=E5=87=BA=E5=85=B7?= =?UTF-8?q?=E5=8E=BB=E6=8E=89=E5=85=A8=E7=AB=99=E4=B8=80=E6=8A=8A=E9=94=81?= =?UTF-8?q?=EF=BC=8C=E5=8F=AA=E7=95=99=E5=90=8C=E4=B8=80=EF=BC=88=E8=BF=90?= =?UTF-8?q?=E8=A1=8C=E3=80=81=E5=B7=A5=E4=BD=9C=E6=B5=81=E3=80=81=E7=89=88?= =?UTF-8?q?=E6=9C=AC=EF=BC=89=E7=9A=84=E5=8E=BB=E9=87=8D=E9=94=81=EF=BC=8C?= =?UTF-8?q?=E5=87=A0=E4=BB=BD=E4=B8=8D=E5=90=8C=E7=9A=84=E5=87=BA=E5=85=B7?= =?UTF-8?q?=E5=8F=AF=E4=BB=A5=E5=90=8C=E6=97=B6=E5=8F=91=E8=B5=B7=EF=BC=9B?= =?UTF-8?q?=E6=8C=89=E9=9C=80=E8=A7=86=E5=9B=BE=E7=9A=84=E9=94=81=E7=BC=A9?= =?UTF-8?q?=E5=9B=9E=E6=8C=89=E7=BC=93=E5=AD=98=E9=94=AE=EF=BC=88=E8=A7=86?= =?UTF-8?q?=E5=9B=BE=20+=20=E5=8F=82=E6=95=B0=EF=BC=89=E5=8E=BB=E9=87=8D?= =?UTF-8?q?=EF=BC=8C=E6=AF=8F=E6=97=A5=E4=B8=8A=E9=99=90=E6=94=B9=E7=94=A8?= =?UTF-8?q?=E8=BF=99=E4=B8=AA=E5=BA=94=E7=94=A8=E7=9A=84=E9=A2=9D=E5=BA=A6?= =?UTF-8?q?=E3=80=82=E8=A1=A5=E6=B5=8B=E8=AF=95=EF=BC=9A20=20=E4=B8=AA?= =?UTF-8?q?=E5=8D=8F=E7=A8=8B=E5=90=8C=E6=97=B6=E5=8D=A0=E6=AD=A3=E5=A5=BD?= =?UTF-8?q?=E5=88=B0=E4=B8=8A=E9=99=90=E3=80=81=E6=8C=89=E5=B7=B2=E6=9C=89?= =?UTF-8?q?=E8=AE=B0=E5=BD=95=E8=B5=B7=E7=AE=97=E3=80=81=E9=80=80=E5=9B=9E?= =?UTF-8?q?=E5=86=8D=E5=8D=A0=EF=BC=9B=E4=B8=A4=E4=B8=AA=E8=BF=9B=E7=A8=8B?= =?UTF-8?q?=E5=90=8C=E6=97=B6=E5=8D=A0=E6=AD=A3=E5=A5=BD=E5=88=B0=E4=B8=8A?= =?UTF-8?q?=E9=99=90=EF=BC=9B=E4=B8=89=E4=BB=BD=E4=B8=8D=E5=90=8C=E7=9A=84?= =?UTF-8?q?=E7=AD=94=E6=A1=88=E5=90=8C=E6=97=B6=E5=87=BA=E5=85=B7=E6=AD=A3?= =?UTF-8?q?=E5=A5=BD=E5=88=B0=E4=B8=8A=E9=99=90=E3=80=81=E5=A4=9A=E5=87=BA?= =?UTF-8?q?=E6=9D=A5=E7=9A=84=E5=9B=9E=20429?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/apps/resolve.py | 38 ++++++++++------ backend/app/cluster/quota.py | 70 ++++++++++++++++++++++++++++++ backend/app/cluster/service.py | 7 +-- backend/app/ledger/issue.py | 22 ++++++++-- backend/tests/test_cluster.py | 34 +++++++++++++++ backend/tests/test_ledger_issue.py | 22 ++++++++++ 6 files changed, 173 insertions(+), 20 deletions(-) create mode 100644 backend/app/cluster/quota.py diff --git a/backend/app/apps/resolve.py b/backend/app/apps/resolve.py index d20bb73a..95bb9d58 100644 --- a/backend/app/apps/resolve.py +++ b/backend/app/apps/resolve.py @@ -388,7 +388,7 @@ _LOCKS: dict[str, asyncio.Lock] = {} def _lock(key: str) -> Any: - """按需视图的单飞锁（app/cluster/mutex，多进程之间也是一把）：按应用取，同一个键正在运行的那次接上、不另发，每日上限也数得准。""" + """按需视图的单飞锁（app/cluster/mutex，多进程之间也是一把）：按缓存键（视图 + 参数）取，同一个键正在运行的那次接上、不另发。""" from app.cluster import mutex return mutex.lock(f"app-view:{key}") @@ -471,25 +471,35 @@ async def run_view(session: AsyncSession, *, app: Any, version: AppVersion, view raise FormRejected(errors) inputs = normalize_params(form, params) key = await cache_key(session, state, inputs) - # 按应用一把锁（不按缓存键）：同一个键正在运行的那次接上不另发，每日上限的「数一数再发」几个进程、几个键之间也不越过 - async with _lock(app.id): + # 锁按缓存键（这个视图、这组参数）：管的是去重——同一个键正在运行的那次接上，不另算一遍。每日上限不靠它：发起之前在库里先占 + # 一份这个应用今天的额度（app/cluster/quota，带条件的加一），几个键、几个进程同时刷新也正好到上限为止，发起不成退回 + from app.cluster import quota + + async with _lock(key): run = await _cached(session, key) if run is not None: return run, run.status == "succeeded" - # 「每天」按服务端所在时区的自然日算 - start = datetime.now().astimezone().replace(hour=0, minute=0, second=0, microsecond=0).astimezone(timezone.utc) - used = await session.scalar(select(func.count(Run.id)).where( - Run.app_id == app.id, Run.origin_kind == "app_view", Run.created_at >= start)) limit = await daily_limit(session) - if int(used or 0) >= limit: + + async def used_today(s: AsyncSession) -> int: + # 「每天」按服务端所在时区的自然日算 + start = datetime.now().astimezone().replace(hour=0, minute=0, second=0, microsecond=0).astimezone(timezone.utc) + 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): raise ViewRefused(429, f"这个应用今天的按需运行已达到上限（{limit} 次），请明天再试，或请管理员调整上限", DAILY_LIMIT) - run = await launch_run( - session, workflow_id=state.pin["workflow_id"], run_class="formal", version=state.pin["version"], - expect_hash=state.pin.get("graph_hash"), input_payload=inputs, actor=actor, origin_kind="app_view", - app_id=app.id, app_version=version.version, - origin={"view": view_id, "spec_hash": version.spec_hash, "actor_verified": principal.verified, - "auth_mode": principal.mode, **_key_trace(principal)}) + try: + run = await launch_run( + session, workflow_id=state.pin["workflow_id"], run_class="formal", version=state.pin["version"], + expect_hash=state.pin.get("graph_hash"), input_payload=inputs, actor=actor, origin_kind="app_view", + app_id=app.id, app_version=version.version, + 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}") + raise ttl = source.get("ttl_s") if ttl is None and await reads_live(session, state.graph): ttl = TTL_DEFAULT diff --git a/backend/app/cluster/quota.py b/backend/app/cluster/quota.py new file mode 100644 index 00000000..9a726966 --- /dev/null +++ b/backend/app/cluster/quota.py @@ -0,0 +1,70 @@ +"""每日额度：先占一份额度再做事，做不成退回（ledger/issue 的出具上限、apps/resolve 的按需视图上限）。 + +**为什么不靠锁。** 以前是「数一数今天做了几次，没到上限就发起」，数和发起之间靠一把锁串行；多进程时那把锁要么只管本进程（几个 +进程一起越过上限），要么全站一把（几份出具只能一个接一个发起）。现在额度记在库里（rate_counters 里按「名字 + 本地日期」一行）， +占额度是一句带条件的 UPDATE：还没到上限才加一，否则什么都不改——两种库上这一句都是原子的，几个进程、几个协程同时来也正好 +占到上限为止，发起本身可以同时做。发起失败（被拒、出错）把那一份退回。 + +这一天第一次占的时候按库里已有的记录起算（seed：比如今天已经出具过几次），之后只看这一行。进程在占了额度、发起之前没了的， +那一份不退：宁可这一天少出一次，不会多出。 +""" +from __future__ import annotations + +from collections.abc import Awaitable, Callable +from datetime import datetime, timedelta +from typing import Any + +from sqlalchemy import select, update + +from app.db.base import SessionLocal, utcnow + + +def today() -> str: + """本地日期（每日上限按服务端所在时区的自然日算，同以前按本地零点数）。""" + return datetime.now().astimezone().date().isoformat() + + +def _key(name: str, day: str) -> str: + return f"quota:{name}:{day}" + + +async def take(name: str, limit: int, *, seed: Callable[[Any], Awaitable[int]] | None = None, + day: str | None = None) -> bool: + """占一份额度：这一天占过的还不到 limit 就加一、交回 True，到了交回 False（什么都不改）。seed(session)：这一天第一次占时 + 已经用掉了几份（按库里的记录数）。""" + from app.db.models_cluster import RateCounter as Row + from app.db.portable import insert_ignore + + key = _key(name, day or today()) + 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 + await session.execute(insert_ignore(session, Row.__table__).values( + key=key, count=used, value=0.0, at=0.0, expires_at=utcnow() + timedelta(days=3)) + .on_conflict_do_nothing(index_elements=["key"])) + 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 + + +async def give_back(name: str, *, day: str | None = None) -> None: + """退回一份（占了额度、却没发起成）。""" + from app.db.models_cluster import RateCounter as Row + + key = _key(name, day or today()) + async with SessionLocal() as session: + await session.execute(update(Row).where(Row.key == key, Row.count > 0).values(count=Row.count - 1) + .execution_options(synchronize_session=False)) + await session.commit() + + +async def used(name: str, *, day: str | None = None) -> int | None: + """这一天占了几份（没占过为 None）。""" + from app.db.models_cluster import RateCounter as Row + + async with SessionLocal() as session: + return await session.scalar(select(Row.count).where(Row.key == _key(name, day or today()))) + + +__all__ = ["give_back", "take", "today", "used"] diff --git a/backend/app/cluster/service.py b/backend/app/cluster/service.py index 424996c8..35b52029 100644 --- a/backend/app/cluster/service.py +++ b/backend/app/cluster/service.py @@ -152,10 +152,11 @@ class _Service: with contextlib.suppress(Exception): await workers.prune() await requests.prune() - if settings.db_dialect == "postgresql": - from app.cluster import caches, counters + from app.cluster import caches, counters - await counters.prune() + # 到期的计数、额度（每日额度两种库都记在库里） + await counters.prune() + if settings.db_dialect == "postgresql": await caches.prune() return out diff --git a/backend/app/ledger/issue.py b/backend/app/ledger/issue.py index e7c9181d..e2bddd08 100644 --- a/backend/app/ledger/issue.py +++ b/backend/app/ledger/issue.py @@ -329,11 +329,16 @@ _locks: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[tuple[str, st def _lock(key: tuple[str, str, int]) -> Any: - """跨进程的锁（app/cluster/mutex）：多进程时「有没有出具过、今天出了几次」的核对和发起照样一个接一个。全站一把（不按键）： - 每日上限是全站的，按键分锁时几个不同的答案同时出具会一起越过上限；出具不频繁，发起本身只是建一条运行。""" + """同一（运行、工作流、版本）的出具一个接一个（app/cluster/mutex，多进程之间也是一把）：管的是「这一次出具过没有」的去重， + 同时点了几下只发起一次。每日上限不靠它：发起之前先在库里占一份额度（app/cluster/quota），几份不同的出具可以同时发起，上限 + 照样数得准。""" from app.cluster import mutex - return mutex.lock("ledger-issue") + return mutex.lock(f"ledger-issue:{key[0]}:{key[1]}:{key[2]}") + + +#: 每日出具额度在 cluster/quota 里的名字 +QUOTA = "ledger_issue" async def _hider(run_id: str) -> Any: @@ -369,12 +374,23 @@ async def issue(run_id: str, workflow_id: str, version: int, *, principal: Any, "how": plan.how, "binding_hash": input_key(plan.binding), "match": plan.match_artifact, "auto": auto, **({"granted_by_app": plan.granted_by_app} if plan.granted_by_app else {})} actor = run.started_by if auto else principal.username + # 先占一份今天的额度（库里带条件的加一，几个进程同时出具也正好到上限为止），发起不成退回 + from app.api.settings import ledger_settings + from app.cluster import quota + + limit = int((await ledger_settings(session))["issue_daily_limit"]) + if not await quota.take(QUOTA, limit, seed=issued_today): + 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) raise IssueRefused(e.status_code, e.message, e.code or REFUSED) from e + except BaseException: + await quota.give_back(QUOTA) + raise return launched, True diff --git a/backend/tests/test_cluster.py b/backend/tests/test_cluster.py index 3fb648a9..557faf08 100644 --- a/backend/tests/test_cluster.py +++ b/backend/tests/test_cluster.py @@ -341,3 +341,37 @@ async def _until(check, timeout: float = 10.0) -> None: if loop.time() > deadline: raise AssertionError("等超时了") await asyncio.sleep(0.05) + + +async def test_a_daily_quota_is_taken_up_to_the_limit_by_concurrent_callers(): + """先占额度再发起（cluster/quota）：20 个协程同时占、上限 5，正好 5 个占到；这一天第一次占按已有的记录起算；退回一份又能占。""" + from app.cluster import quota + + name = f"test-quota-{new_id()[:8]}" + + async def seed(_session) -> int: + return 2 # 今天已经做过 2 次 + + got = await asyncio.gather(*(quota.take(name, 5, seed=seed) for _ in range(20))) + assert sum(got) == 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) + + +def test_a_daily_quota_holds_across_two_processes(): + """两个进程同时各占 15 次、上限 10：合计正好 10 次。""" + name = f"test-quota-{new_id()[:8]}" + backend = Path(__file__).resolve().parents[1] + script = ("import asyncio, sys\n" + "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" + "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, + stderr=subprocess.PIPE, text=True) for _ in range(2)] + outs = [p.communicate(timeout=120) for p in procs] + assert all(p.returncode == 0 for p in procs), [o[1][-2000:] for o in outs] + assert sum(int(o[0].strip().splitlines()[-1]) for o in outs) == 10 diff --git a/backend/tests/test_ledger_issue.py b/backend/tests/test_ledger_issue.py index 19803753..00c7f4df 100644 --- a/backend/tests/test_ledger_issue.py +++ b/backend/tests/test_ledger_issue.py @@ -297,6 +297,28 @@ async def test_users_issue_only_versions_pinned_by_apps_they_can_use(monkeypatch assert created and formal.started_by == "carol" and formal.origin["granted_by_app"] == board_id +async def test_different_answers_issued_at_once_stop_exactly_at_the_daily_limit(monkeypatch): + """几份不同的答案同时出具（不再一把全站锁一个接一个）：先在库里占今天的额度再发起，正好到上限为止，多出来的回 429。""" + wf = await S.publish(level="published") + rows = [await equivalent_run(monkeypatch) for _ in range(3)] + from app.cluster import quota + + async with SessionLocal() as session: + issued = await ledger_issue.issued_today(session) + # 今天的额度在库里记着（同一个测试进程里前面的测试也出具过）；还没记过的按今天已出具的起算 + already = await quota.used(ledger_issue.QUOTA) + already = issued if already is None else already + await set_ledger(issue_daily_limit=max(already, issued) + 2) + results = await asyncio.gather(*(ledger_issue.issue(r.id, wf, 1, principal=ADMIN) for r in rows), + return_exceptions=True) + made = [r for r in results if not isinstance(r, BaseException)] + refused = [r for r in results if isinstance(r, ledger_issue.IssueRefused)] + assert len(made) == 2 and all(created for _, created in made), results + assert len(refused) == 1 and refused[0].status_code == 429 and refused[0].code == ledger_issue.DAILY_LIMIT + async with SessionLocal() as session: + assert await ledger_issue.issued_today(session) == issued + 2 + + async def test_concurrent_requests_issue_once(monkeypatch): wf = await S.publish(level="published") row = await equivalent_run(monkeypatch)