From ad21d635738d8c1491ef8d4502e9fc4a77746fe2 Mon Sep 17 00:00:00 2001 From: David <20960328+totallynotdavid@users.noreply.github.com> Date: Sun, 6 Sep 2026 14:23:04 +0000 Subject: [PATCH 1/2] Release claims after cancellation Cancelled workers can outlive their async waiter. Release claims acquired by background threads so later attempts are not blocked. --- packages/api/api/core/tasks.py | 80 +++++++++++++++++++++++++------- packages/api/tests/test_tasks.py | 64 +++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 18 deletions(-) diff --git a/packages/api/api/core/tasks.py b/packages/api/api/core/tasks.py index d44a2ef..112addc 100644 --- a/packages/api/api/core/tasks.py +++ b/packages/api/api/core/tasks.py @@ -179,6 +179,10 @@ class AbandonedAttempt(Exception): # plus the grace, not by the interval. RECONCILE_INTERVAL_SECONDS = 60.0 +# Keep cancellation cleanup tasks strongly reachable until they have released +# the claim returned by their background thread. +_CLAIM_DRAIN_TASKS: set[asyncio.Task[None]] = set() + def decode_payload(payload: Any) -> uuid.UUID: """Turn a queued payload into the compute job id it names. @@ -270,21 +274,63 @@ def claim_workspace(work_dir: Path, attempt: int) -> tuple[WorkspaceClaim, bool] lock_path.parent.mkdir(parents=True, exist_ok=True) fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o600) try: - fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) - except OSError as e: - os.close(fd) - # Transient on purpose: the previous attempt's thread unwinds at its - # next progress write, which the database fence refuses, so the - # workspace frees itself. rqueue's backoff is the wait. - raise TransientInfraError( - f"simulation workspace {work_dir.name} is still held by an earlier attempt" - ) from e - os.ftruncate(fd, 0) - os.write(fd, f"attempt {attempt}\n".encode()) - # Evaluated under the lock: without it, "are there checkpoints to resume - # from" is a question about a directory someone else may be halfway - # through writing. - return WorkspaceClaim(fd), work_dir.exists() + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except OSError as e: + # Transient on purpose: the previous attempt's thread unwinds at + # its next progress write, which the database fence refuses, so + # the workspace frees itself. rqueue's backoff is the wait. + raise TransientInfraError( + f"simulation workspace {work_dir.name} is still held by an " + "earlier attempt" + ) from e + os.ftruncate(fd, 0) + os.write(fd, f"attempt {attempt}\n".encode()) + # Evaluated under the lock: without it, "are there checkpoints to + # resume from" is a question about a directory someone else may be + # halfway through writing. + return WorkspaceClaim(fd), work_dir.exists() + except BaseException: + # Once the descriptor exists, this function owns closing it until the + # claim has been handed to the caller. This also covers failures after + # flock succeeds, such as ftruncate/write, and constructor failures. + with contextlib.suppress(OSError): + os.close(fd) + raise + + +async def _release_unreceived_claim( + claim_future: asyncio.Future[tuple[WorkspaceClaim, bool]], +) -> None: + """Release a claim whose thread completed after its waiter was cancelled.""" + try: + claim, _resume = await claim_future + except BaseException: + return + claim.release() + claim.release() + + +async def _claim_workspace_safely( + work_dir: Path, attempt: int +) -> tuple[WorkspaceClaim, bool]: + """Claim a workspace without losing a result to cancellation. + + `asyncio.to_thread` cannot stop the thread that has already acquired the + lock. Shielding its task lets that thread finish, while the detached drain + takes ownership of a returned claim if cancellation wins the handoff + before this coroutine receives it. + """ + claim_future = asyncio.create_task( + asyncio.to_thread(claim_workspace, work_dir, attempt) + ) + try: + return await asyncio.shield(claim_future) + except asyncio.CancelledError: + drain = asyncio.create_task(_release_unreceived_claim(claim_future)) + _CLAIM_DRAIN_TASKS.add(drain) + drain.add_done_callback(_CLAIM_DRAIN_TASKS.discard) + raise def remove_workspace(work_dir: Path) -> bool: @@ -402,9 +448,7 @@ def run_kernel(claim: WorkspaceClaim, resume: bool) -> Any: # rqueue installs a ThreadPoolExecutor sized from the worker's # concurrency as the loop's default executor, so these hops are capacity # bounded without building an executor here. - held, resume = await asyncio.to_thread( - claim_workspace, work_dir, context.attempt - ) + held, resume = await _claim_workspace_safely(work_dir, context.attempt) claim = held result = await asyncio.to_thread(run_kernel, held, resume) async with db.acquire() as conn: diff --git a/packages/api/tests/test_tasks.py b/packages/api/tests/test_tasks.py index ea58edd..78199df 100644 --- a/packages/api/tests/test_tasks.py +++ b/packages/api/tests/test_tasks.py @@ -527,6 +527,70 @@ def test_releasing_a_claim_more_than_twice_is_harmless(tmp_path: Path) -> None: claim.release() +@pytest.mark.parametrize("operation", ["ftruncate", "write"]) +def test_claim_closes_fd_when_lock_setup_fails( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, operation: str +) -> None: + work_dir = tmp_path / "sim" + original = getattr(tasks_module.os, operation) + + def fail(*_args: Any) -> None: + raise OSError(f"{operation} failed") + + monkeypatch.setattr(tasks_module.os, operation, fail) + with pytest.raises(OSError, match=f"{operation} failed"): + tasks.claim_workspace(work_dir, 1) + + # If the descriptor leaked, this second claim would still be refused by + # the flock even though the first call never returned a claim. + monkeypatch.setattr(tasks_module.os, operation, original) + claim, _resume = tasks.claim_workspace(work_dir, 2) + claim.release() + claim.release() + + +@pytest.mark.asyncio +async def test_cancelled_claim_handoff_releases_a_claim_it_never_received( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + work_dir = tmp_path / "sim" + claimed = threading.Event() + release = threading.Event() + original = tasks.claim_workspace + + def claim_then_pause( + current_work_dir: Path, attempt: int + ) -> tuple[tasks.WorkspaceClaim, bool]: + result = original(current_work_dir, attempt) + claimed.set() + release.wait(5) + return result + + monkeypatch.setattr(tasks_module, "claim_workspace", claim_then_pause) + task = asyncio.create_task(tasks_module._claim_workspace_safely(work_dir, 1)) + assert await asyncio.to_thread(claimed.wait, 5) + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # The background thread returns the already-acquired claim after the + # cancellation. The helper's detached drain owns and releases both shares. + monkeypatch.setattr(tasks_module, "claim_workspace", original) + release.set() + for _ in range(100): + try: + next_claim, _resume = tasks_module.claim_workspace(work_dir, 2) + except TransientInfraError: + await asyncio.sleep(0.01) + else: + next_claim.release() + next_claim.release() + break + else: + pytest.fail("cancelled claim handoff kept the workspace locked") + + @pytest.mark.asyncio async def test_a_cancelled_attempt_keeps_holding_its_workspace( worker_job: tuple[uuid.UUID, uuid.UUID, Path], From f13f33d02b0bfdc9dc54113d864a83efd7b68a91 Mon Sep 17 00:00:00 2001 From: David <20960328+totallynotdavid@users.noreply.github.com> Date: Sun, 6 Sep 2026 14:24:57 +0000 Subject: [PATCH 2/2] Clean up completed workspaces Completed jobs can leave checkpoints behind after workers finish. Include all terminal jobs in retention cleanup so stale workspaces do not accumulate. --- packages/api/api/core/repository.py | 7 ++--- packages/api/api/core/tasks.py | 4 +-- .../api/tests/test_repository_integration.py | 14 +++++++++- packages/api/tests/test_tasks.py | 27 +++++++++++++++++++ 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/packages/api/api/core/repository.py b/packages/api/api/core/repository.py index 4b1100e..f85298c 100644 --- a/packages/api/api/core/repository.py +++ b/packages/api/api/core/repository.py @@ -250,14 +250,15 @@ async def get_current_step(conn: asyncpg.Connection, job_uuid: uuid.UUID) -> str async def list_abandoned_work_dirs(cutoff: datetime) -> list[str]: - """Return failed job workspaces older than the retention cutoff.""" + """Return terminal job workspaces older than the retention cutoff.""" async with acquire() as conn: rows = await conn.fetch( """ SELECT simulation_id FROM compute.jobs - WHERE status = $1 AND finished_at IS NOT NULL AND finished_at < $2 + WHERE status = ANY($1::text[]) + AND finished_at IS NOT NULL AND finished_at < $2 """, - JobStatus.FAILED.value, + [JobStatus.FAILED.value, JobStatus.COMPLETED.value], cutoff, ) return [str(row["simulation_id"]) for row in rows] diff --git a/packages/api/api/core/tasks.py b/packages/api/api/core/tasks.py index 112addc..9d6122b 100644 --- a/packages/api/api/core/tasks.py +++ b/packages/api/api/core/tasks.py @@ -155,7 +155,7 @@ class AbandonedAttempt(Exception): # starts without resuming, which would take the lock with it. WORKSPACE_LOCK_SUFFIX = ".lock" -# Keep failed workspaces for local inspection and manual recovery. +# Keep terminal workspaces for local inspection and manual recovery. WORK_DIR_TTL = timedelta(hours=24) SWEEP_INTERVAL_SECONDS = 3600.0 @@ -584,7 +584,7 @@ async def _record_failure( async def sweep_abandoned_work_dirs() -> None: - """Delete the workspaces of jobs that failed longer than WORK_DIR_TTL ago.""" + """Delete terminal-job workspaces older than WORK_DIR_TTL.""" cutoff = datetime.now().astimezone() - WORK_DIR_TTL for simulation_id in await repository.list_abandoned_work_dirs(cutoff): await asyncio.to_thread(remove_workspace, JOBS_DIR / simulation_id) diff --git a/packages/api/tests/test_repository_integration.py b/packages/api/tests/test_repository_integration.py index 6962bfd..7f8a914 100644 --- a/packages/api/tests/test_repository_integration.py +++ b/packages/api/tests/test_repository_integration.py @@ -276,13 +276,14 @@ async def test_a_notification_is_only_delivered_once_its_update_commits( @pytest.mark.asyncio -async def test_list_abandoned_work_dirs_returns_only_old_failed_jobs( +async def test_list_abandoned_work_dirs_returns_old_terminal_jobs( queue: Queue, ) -> None: simulation_id_text, compute_job_id = await _create_job() simulation_id = uuid.UUID(simulation_id_text) fresh_simulation_id_text, fresh_job_id = await _create_job() fresh_simulation_id = uuid.UUID(fresh_simulation_id_text) + completed_simulation_id_text, completed_job_id = await _create_job() cutoff = datetime.now(UTC) - timedelta(hours=24) async with db.acquire() as conn: @@ -305,9 +306,20 @@ async def test_list_abandoned_work_dirs_returns_only_old_failed_jobs( cutoff - timedelta(minutes=1), compute_job_id, ) + await conn.execute( + """ + UPDATE compute.jobs + SET status = $1, finished_at = $2 + WHERE id = $3 + """, + "completed", + cutoff - timedelta(minutes=1), + completed_job_id, + ) abandoned = await repository.list_abandoned_work_dirs(cutoff) assert simulation_id_text in abandoned + assert completed_simulation_id_text in abandoned assert fresh_simulation_id_text not in abandoned diff --git a/packages/api/tests/test_tasks.py b/packages/api/tests/test_tasks.py index 78199df..3f91b96 100644 --- a/packages/api/tests/test_tasks.py +++ b/packages/api/tests/test_tasks.py @@ -685,3 +685,30 @@ def test_a_held_workspace_is_left_for_the_next_sweep(tmp_path: Path) -> None: claim.release() assert tasks.remove_workspace(work_dir) is True assert not work_dir.exists() + + +@pytest.mark.asyncio +async def test_sweep_retries_a_completed_workspace_after_lock_race( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + simulation_id = str(uuid.uuid4()) + work_dir = tmp_path / simulation_id + work_dir.mkdir() + claim, _resume = tasks.claim_workspace(work_dir, 1) + + async def list_terminal_work_dirs(_cutoff: Any) -> list[str]: + # The repository query includes completed jobs as well as failed ones. + return [simulation_id] + + monkeypatch.setattr( + tasks_module.repository, "list_abandoned_work_dirs", list_terminal_work_dirs + ) + monkeypatch.setattr(tasks_module, "JOBS_DIR", tmp_path) + + await tasks.sweep_abandoned_work_dirs() + assert work_dir.exists() + + claim.release() + claim.release() + await tasks.sweep_abandoned_work_dirs() + assert not work_dir.exists()