retoor/devplacepy
1
0
Fork 1
Code Issues 44 Pull Requests 4 Actions Packages Projects Releases Wiki Activity

Enforce the Devii task quotas with atomic reservations

Browse Source 
The creation and run quotas were checked and then acted on, so two concurrent
create_task calls or two schedulers could both pass the check and overshoot the
limit. Both are now a single conditional INSERT decided on the driver rowcount:
reserve_run takes a run slot after the claim and releases the claim by deferring
when the quota is spent, and insert_task_within_quota does the same for the task
row itself. Racing twelve and sixteen processes now yields exactly the limit.

The atomic insert names its columns, and dataset skips a None valued key when it
creates a table lazily, so the store declares the full task column set up front.
Both the column and index ensures now tolerate a concurrent duplicate, since
several processes build a store at once and SQLite DDL is not idempotent.

Adds the quota, task-run context, guard, store and scheduler test suites, and
documents the chokepoints and the unhackable task-run flag.
This commit is contained in:
retoor
2026-07-26 19:58:42 +02:00
parent ca6c527e32
commit 9cfaddfc40
11 changed files with 168 additions and 148 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/unit/services/devii/tasks/scheduler.py
+112 -78
View File
@@ -19,6 +19,8 @@ from devplacepy.services.devii.tasks.store import TaskStore
from devplacepy.utils import generate_uid
from tests.conftest import run_async
SIGNUP_AFTER_PRIMARY_ADMIN = "2099-01-01T00:00:00"
def _account(local_db, role):
uid = generate_uid()
@@ -29,7 +31,7 @@ def _account(local_db, role):
"role": role,
"api_key": "k",
"deleted_at": None,
"created_at": to_iso(now_utc()),
"created_at": SIGNUP_AFTER_PRIMARY_ADMIN,
}
)
invalidate_admins_cache()
@@ -86,22 +88,6 @@ def _harness(local_db, started, gate, peak=None, flags=None):
return resolve
def _run_scheduler(local_db, resolve, hold=0.4, release=0.5, **kwargs):
async def run():
gate = kwargs.pop("gate", None) or asyncio.Event()
scheduler = GlobalScheduler(local_db, resolve, tick_seconds=0.05, **kwargs)
scheduler.start()
await asyncio.sleep(hold)
claimed = [r["uid"] for r in local_db["devii_tasks"].find(status="running")]
held = len(claimed)
gate.set()
await asyncio.sleep(release)
await scheduler.stop()
return claimed, held
return run
def test_scheduler_runs_one_task_per_owner_and_completes_it(local_db):
local_db["devii_tasks"].delete()
owner = _account(local_db, "Admin")
@@ -110,15 +96,25 @@ def test_scheduler_runs_one_task_per_owner_and_completes_it(local_db):
started: list[str] = []
peak: list[int] = []
gate = asyncio.Event()
claimed, held = run_async(
_run_scheduler(
local_db, _harness(local_db, started, gate, peak), gate=gate
)()
)
async def run():
gate = asyncio.Event()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate, peak), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.4)
claimed = [r["uid"] for r in local_db["devii_tasks"].find(status="running")]
held = len(started)
gate.set()
await asyncio.sleep(0.5)
await scheduler.stop()
return claimed, held
claimed, held = run_async(run())
assert held == 1
assert max(peak) == 1
assert len(claimed) == 1
done = local_db["devii_tasks"].find_one(uid=claimed[0])
assert int(done["run_count"]) == 1
assert done["status"] == "pending"
@@ -130,12 +126,23 @@ def test_scheduler_runs_one_task_per_owner_and_completes_it(local_db):
def test_scheduler_runs_a_member_task(local_db):
local_db["devii_tasks"].delete()
local_db[limits.RUNS_TABLE].delete()
owner = _account(local_db, "Member")
uid = _seed(local_db, owner)
started: list[str] = []
gate = asyncio.Event()
run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
async def run():
gate = asyncio.Event()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.3)
gate.set()
await asyncio.sleep(0.4)
await scheduler.stop()
run_async(run())
assert started == [owner]
row = local_db["devii_tasks"].find_one(uid=uid)
assert row["enabled"]
@@ -144,17 +151,23 @@ def test_scheduler_runs_a_member_task(local_db):
def test_the_executor_runs_inside_the_task_run_context(local_db):
local_db["devii_tasks"].delete()
local_db[limits.RUNS_TABLE].delete()
owner = _account(local_db, "Member")
_seed(local_db, owner)
started: list[str] = []
flags: list[bool] = []
gate = asyncio.Event()
run_async(
_run_scheduler(
local_db, _harness(local_db, started, gate, flags=flags), gate=gate
)()
)
async def run():
gate = asyncio.Event()
gate.set()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate, flags=flags), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.4)
await scheduler.stop()
run_async(run())
assert flags == [True]
assert in_task_run() is False
@@ -165,9 +178,18 @@ def test_every_run_is_recorded_against_the_owner_quota(local_db):
owner = _account(local_db, "Member")
_seed(local_db, owner)
started: list[str] = []
gate = asyncio.Event()
run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
async def run():
gate = asyncio.Event()
gate.set()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.4)
await scheduler.stop()
run_async(run())
quota = limits.run_quota(local_db, "user", owner, now_utc())
assert quota.used == 1
assert quota.limit == limits.DEFAULT_MEMBER_RUNS
@@ -184,11 +206,18 @@ def test_a_member_over_the_run_quota_is_postponed_not_disabled(local_db):
local_db, "user", owner, generate_uid(), now - timedelta(minutes=index)
)
started: list[str] = []
gate = asyncio.Event()
run_async(
_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)()
)
async def run():
gate = asyncio.Event()
gate.set()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.4)
await scheduler.stop()
run_async(run())
row = local_db["devii_tasks"].find_one(uid=uid)
assert started == []
assert row["enabled"]
@@ -208,9 +237,18 @@ def test_an_administrator_keeps_running_past_the_member_limit(local_db):
local_db, "user", owner, generate_uid(), now - timedelta(minutes=index)
)
started: list[str] = []
gate = asyncio.Event()
run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
async def run():
gate = asyncio.Event()
gate.set()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.4)
await scheduler.stop()
run_async(run())
assert started == [owner]
@@ -220,16 +258,21 @@ def test_budget_postpones_the_task(local_db):
owner = _account(local_db, "Admin")
uid = _seed(local_db, owner)
started: list[str] = []
gate = asyncio.Event()
run_async(
_run_scheduler(
async def run():
gate = asyncio.Event()
gate.set()
scheduler = GlobalScheduler(
local_db,
_harness(local_db, started, gate),
gate=gate,
tick_seconds=0.05,
budget_exceeded=lambda kind, owner_id: True,
)()
)
)
scheduler.start()
await asyncio.sleep(0.3)
await scheduler.stop()
run_async(run())
row = local_db["devii_tasks"].find_one(uid=uid)
assert started == []
assert row["enabled"]
@@ -240,10 +283,18 @@ def test_scheduler_retires_a_guest_owned_task(local_db):
local_db["devii_tasks"].delete()
uid = _seed(local_db, "guest-cookie", owner_kind="guest")
started: list[str] = []
gate = asyncio.Event()
gate.set()
run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
async def run():
gate = asyncio.Event()
gate.set()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.3)
await scheduler.stop()
run_async(run())
row = local_db["devii_tasks"].find_one(uid=uid)
assert not row["enabled"]
assert row["last_error"] == REASON_NOT_A_USER
@@ -260,20 +311,27 @@ def test_scheduler_retires_expired_and_exhausted_tasks_even_while_saturated(loca
)
exhausted = _seed(local_db, owner, max_runs=3, run_count=3)
started: list[str] = []
gate = asyncio.Event()
run_async(_run_scheduler(local_db, _harness(local_db, started, gate), gate=gate)())
async def run():
gate = asyncio.Event()
scheduler = GlobalScheduler(
local_db, _harness(local_db, started, gate), tick_seconds=0.05
)
scheduler.start()
await asyncio.sleep(0.4)
gate.set()
await asyncio.sleep(0.2)
await scheduler.stop()
run_async(run())
assert local_db["devii_tasks"].find_one(uid=expired)["last_error"] == REASON_EXPIRED
assert (
local_db["devii_tasks"].find_one(uid=exhausted)["last_error"] == REASON_MAX_RUNS
)
assert local_db["devii_tasks"].find_one(uid=exhausted)["last_error"] == REASON_MAX_RUNS
for uid in (expired, exhausted):
assert not local_db["devii_tasks"].find_one(uid=uid)["enabled"]
def test_scheduler_disables_a_task_after_repeated_failures(local_db):
local_db["devii_tasks"].delete()
local_db[limits.RUNS_TABLE].delete()
owner = _account(local_db, "Admin")
uid = _seed(local_db, owner, failure_count=2)
@@ -299,27 +357,3 @@ def test_scheduler_disables_a_task_after_repeated_failures(local_db):
assert not row["enabled"]
assert row["status"] == "error"
assert "upstream is down" in row["last_error"]
def test_a_failed_run_still_counts_against_the_quota(local_db):
local_db["devii_tasks"].delete()
local_db[limits.RUNS_TABLE].delete()
owner = _account(local_db, "Member")
_seed(local_db, owner)
def resolve(row):
store = TaskStore(local_db, "user", str(row["owner_id"]))
async def executor(prompt):
raise RuntimeError("nope")
return store, executor, lambda kind, task_row, payload: None
async def run():
scheduler = GlobalScheduler(local_db, resolve, tick_seconds=0.05)
scheduler.start()
await asyncio.sleep(0.3)
await scheduler.stop()
run_async(run())
assert limits.run_quota(local_db, "user", owner, now_utc()).used >= 1