mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 01:11:59 +00:00
复核 R-b11(跨零点退回额度退到第二天那一行):quota.take 交回占在了哪一天(本地日期),占不到交回 None;出具、按需视图发起失败时按这个日期退回——23:59:59 占、00:00:01 失败,退的是前一天那一份,第二天的不动。补测试:跨零点占了两份、退回前一天那一份,前一天归零、第二天那份不动
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user