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/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/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/tests/test_cluster.py b/backend/tests/test_cluster.py index dcaa42eb..ffe61b21 100644 --- a/backend/tests/test_cluster.py +++ b/backend/tests/test_cluster.py @@ -526,12 +526,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]}" @@ -540,7 +555,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,