forked from retoor/devplacepy
Enforce the Devii task quotas with atomic reservations
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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user