cap/picv workspace claim leak

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]

packages/api/api/core/tasks.py

@@ -155,7 +155,7 @@
 # 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
 
@@ -179,6 +179,10 @@
 # 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 @@
     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 @@
         # 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:
@@ -540,7 +584,7 @@
 
 
 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)

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
 
 

packages/api/tests/test_tasks.py

@@ -527,6 +527,70 @@
         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],
@@ -621,3 +685,30 @@
     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()