mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:11:59 +00:00
出具、按需视图的每日上限改成先占额度再发起(用户拍板):新增 cluster/quota——额度记在库里(rate_counters 按名字 + 本地日期一行),占一份是一句带条件的 UPDATE(没到上限才加一),两种库上都是原子的,这一天第一次占按库里已有的记录起算,发起被拒或出错就退回。出具去掉全站一把锁,只留同一(运行、工作流、版本)的去重锁,几份不同的出具可以同时发起;按需视图的锁缩回按缓存键(视图 + 参数)去重,每日上限改用这个应用的额度。补测试:20 个协程同时占正好到上限、按已有记录起算、退回再占;两个进程同时占正好到上限;三份不同的答案同时出具正好到上限、多出来的回 429
This commit is contained in:
+24
-14
@@ -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
|
||||
|
||||
@@ -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"]
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user