128 lines
7.8 KiB
Python
128 lines
7.8 KiB
Python
"""Opt-in real PostgreSQL queue races, using an explicitly disposable database."""
|
|
import asyncio
|
|
import os
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from fastapi import HTTPException
|
|
from psycopg import AsyncConnection
|
|
from psycopg.conninfo import conninfo_to_dict
|
|
from psycopg.rows import dict_row
|
|
|
|
from app.infra.db.postgresql.calculation_jobs import CalculationJobRepository
|
|
|
|
|
|
def test_concurrent_submission_claims_fencing_and_project_isolation():
|
|
dsn = os.getenv("CALCULATION_QUEUE_TEST_DSN")
|
|
if not dsn:
|
|
pytest.skip("Set CALCULATION_QUEUE_TEST_DSN to a disposable PostgreSQL database")
|
|
assert conninfo_to_dict(dsn)["dbname"].startswith("tjwater_calc_test_")
|
|
|
|
async def run():
|
|
project, other, owner = uuid4(), uuid4(), uuid4()
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
migrations = Path(__file__).parents[2] / "resources/sql"
|
|
await conn.execute((migrations / "026_calculation_jobs.sql").read_text())
|
|
legacy_job = uuid4()
|
|
await conn.execute("""INSERT INTO workbench.calculation_jobs
|
|
(job_id,project_id,owner_user_id,request_key,tool,parameters,binding,timeout_seconds,status,result,updated_at)
|
|
VALUES (%s,%s,%s,'legacy','run_simulation','{}','{}',1800,'succeeded','{"preserved":true}',now()-interval '30 days')""",
|
|
(legacy_job, uuid4(), owner))
|
|
lifecycle = (migrations / "027_calculation_job_lifecycle.sql").read_text()
|
|
await conn.execute(lifecycle)
|
|
cursor = await conn.execute("SELECT result,result_expires_at,result_expires_at>now()+interval '6 days' AS grace FROM workbench.calculation_jobs WHERE job_id=%s", (legacy_job,))
|
|
first = await cursor.fetchone()
|
|
assert first["result"] == {"preserved": True} and first["grace"]
|
|
await conn.execute(lifecycle)
|
|
cursor = await conn.execute("SELECT result_expires_at FROM workbench.calculation_jobs WHERE job_id=%s", (legacy_job,))
|
|
assert (await cursor.fetchone())["result_expires_at"] == first["result_expires_at"]
|
|
|
|
async def submit(key, parameters=None, selected_project=project):
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
return await CalculationJobRepository(conn, selected_project).enqueue(
|
|
owner_user_id=owner, request_key=key, tool="run_burst_location",
|
|
parameters=parameters or {"x": 1}, binding={"native_session_id": "session"}, timeout_seconds=600)
|
|
|
|
# Four independent API connections, same logical submission.
|
|
submissions = await asyncio.gather(*(submit("same") for _ in range(4)))
|
|
assert len({str(row["job_id"]) for row in submissions}) == 1
|
|
with pytest.raises(HTTPException) as conflict:
|
|
await submit("same", {"x": 2})
|
|
assert conflict.value.status_code == 409
|
|
second = await submit("second")
|
|
foreign = await submit("foreign", selected_project=other)
|
|
|
|
async def claim(worker):
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
return await CalculationJobRepository(conn, project).claim(worker, 60)
|
|
|
|
claimed = await asyncio.gather(*(claim(f"worker-{index}") for index in range(4)))
|
|
rows = [row for row in claimed if row]
|
|
assert len(rows) == 2
|
|
assert len({str(row["job_id"]) for row in rows}) == 2
|
|
assert str(foreign["job_id"]) not in {str(row["job_id"]) for row in rows}
|
|
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
repo = CalculationJobRepository(conn, project)
|
|
assert await repo.get(foreign["job_id"], owner) is None
|
|
assert await repo.get(second["job_id"], uuid4()) is None
|
|
row = rows[0]
|
|
stale = {**row, "lease_generation": row["lease_generation"] - 1}
|
|
assert await repo.heartbeat(stale, row["lease_owner"], 60) is None
|
|
assert not await repo.finish(stale, row["lease_owner"], "succeeded", result={})
|
|
cancelled = await repo.cancel(row["job_id"], owner)
|
|
assert cancelled["status"] == "cancel_requested"
|
|
assert not await repo.finish(row, row["lease_owner"], "succeeded", result={})
|
|
assert await repo.finish(row, row["lease_owner"], "cancelled")
|
|
expired = rows[1]
|
|
await conn.execute("UPDATE workbench.calculation_jobs SET lease_expires_at=now()-interval '1 second' WHERE job_id=%s", (expired["job_id"],))
|
|
await repo.recover()
|
|
assert (await repo.get(expired["job_id"], owner))["status"] == "interrupted"
|
|
assert not await repo.finish(expired, expired["lease_owner"], "succeeded", result={})
|
|
assert await repo.claim("replacement", 60) is None
|
|
|
|
# Queued cancellation is terminal immediately, with no later claim.
|
|
queued = await repo.enqueue(owner_user_id=owner, request_key="cancel-before-start",
|
|
tool="run_simulation", parameters={}, binding={}, timeout_seconds=30)
|
|
cancelled = await repo.cancel(queued["job_id"], owner)
|
|
assert cancelled["status"] == "cancelled" and cancelled["finished_at"] is not None
|
|
assert cancelled["started_at"] is None
|
|
assert await repo.claim("replacement", 60) is None
|
|
|
|
# A polling client expires queued jobs without any worker process.
|
|
waiting = await repo.enqueue(owner_user_id=owner, request_key="expired-queue",
|
|
tool="run_simulation", parameters={}, binding={}, timeout_seconds=30)
|
|
await conn.execute("UPDATE workbench.calculation_jobs SET queue_expires_at=now()-interval '1 second' WHERE job_id=%s", (waiting["job_id"],))
|
|
status = await repo.get(waiting["job_id"], owner)
|
|
assert status["status"] == "timed_out" and status["error"]["code"] == "QUEUE_TIMEOUT"
|
|
assert status["finished_at"] is not None
|
|
assert "result" not in status
|
|
|
|
# Retention scrubs only payload, keeping deduplication and status.
|
|
completed = await repo.enqueue(owner_user_id=owner, request_key="retained-identity",
|
|
tool="run_simulation", parameters={}, binding={}, timeout_seconds=30)
|
|
await conn.execute("""UPDATE workbench.calculation_jobs SET status='succeeded',
|
|
result='{"large":"payload"}'::jsonb,finished_at=now(),
|
|
result_expires_at=now()-interval '1 second' WHERE job_id=%s""", (completed["job_id"],))
|
|
expired_result = await repo.get(completed["job_id"], owner, include_result=True)
|
|
assert expired_result["result"] is None and expired_result["status"] == "succeeded"
|
|
duplicate = await repo.enqueue(owner_user_id=owner, request_key="retained-identity",
|
|
tool="run_simulation", parameters={}, binding={}, timeout_seconds=30, max_active_jobs=0)
|
|
assert duplicate["job_id"] == completed["job_id"]
|
|
|
|
# Concurrent unique requests cannot exceed an otherwise empty project's cap.
|
|
capacity_project = uuid4()
|
|
async def bounded_submit(index):
|
|
try:
|
|
async with await AsyncConnection.connect(dsn, row_factory=dict_row) as conn:
|
|
return await CalculationJobRepository(conn, capacity_project).enqueue(
|
|
owner_user_id=owner, request_key=str(index), tool="run_simulation",
|
|
parameters={}, binding={}, timeout_seconds=30, max_active_jobs=2)
|
|
except HTTPException as exc:
|
|
return exc
|
|
bounded = await asyncio.gather(*(bounded_submit(i) for i in range(8)))
|
|
assert sum(isinstance(value, dict) for value in bounded) == 2
|
|
assert all(isinstance(value, dict) or value.status_code == 503 for value in bounded)
|
|
asyncio.run(run())
|