use anyio and fixup process handling

pyproject.toml

@@ -9,6 +9,7 @@ authors = [
 requires-python = ">=3.14"
 dependencies = [
     "aiosqlite>=0.22.1",
+    "anyio>=4.14.0",
     "fastapi>=0.138.0",
     "fastapi-jwks>=2.0.2",
     "greenlet>=3.5.2",

src/fastcache_api/lifecycle.py

@@ -0,0 +1,37 @@
+import logging
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from uuid import UUID
+
+import anyio
+import anyio.abc
+
+from .reconcile import watch_and_record
+
+logger = logging.getLogger(__name__)
+
+_task_group: anyio.abc.TaskGroup | None = None
+
+
+def schedule_exit_watch(cache_id: UUID, pid: int) -> None:
+    """Eagerly await a just-spawned cache's exit and persist its outcome.
+
+    sweep_dead_caches remains the fallback for caches whose watch task was lost
+    (e.g. api restart).
+    """
+    if _task_group is None:
+        return
+    _task_group.start_soon(watch_and_record, cache_id, pid)
+
+
+@asynccontextmanager
+async def exit_watchers() -> AsyncGenerator[None]:
+    """Hold the task group that per-cache exit-watch tasks run in."""
+    global _task_group
+    async with anyio.create_task_group() as tg:
+        _task_group = tg
+        try:
+            yield
+        finally:
+            tg.cancel_scope.cancel()
+            _task_group = None

src/fastcache_api/main.py

@@ -11,6 +11,7 @@
 
 from .config import settings
 from .db import init_db
+from .lifecycle import exit_watchers
 from .reconcile import monitor_caches, reconcile_caches
 from .routes import api_router
 
@@ -28,7 +29,8 @@ async def lifespan(application: FastAPI) -> AsyncGenerator[None]:
     await reconcile_caches()
     monitor = asyncio.create_task(monitor_caches())
     try:
-        yield
+        async with exit_watchers():
+            yield
     finally:
         monitor.cancel()
         with contextlib.suppress(asyncio.CancelledError):

src/fastcache_api/models.py

@@ -65,6 +65,7 @@ class CachePublic(BaseModel):
     transfer_id: str
     user: str
     state: CacheState
+    exit_code: int | None
     log_path: Path
     config: CacheConfig
 

src/fastcache_api/process.py

@@ -6,6 +6,8 @@
 from pathlib import Path
 from uuid import UUID
 
+import anyio
+import anyio.abc
 import psutil
 
 from .config import settings
@@ -36,7 +38,31 @@ def allocate_port_pair(in_use: set[int], start: int, end: int) -> tuple[int, int
     raise RuntimeError(f"no free cache port pair in range [{start}, {end}]")
 
 
-def start_cache(cache_id: UUID, config: CacheConfig, log_path: Path) -> CacheProcess:
+# Live anyio Process handles for children we spawned
+_processes: dict[int, anyio.abc.Process] = {}
+
+
+def exit_code(pid: int) -> int | None:
+    """Exit code of a pid we spawned. None if unknown to us or still running."""
+    proc = _processes.get(pid)
+    if proc is None:
+        return None
+    return proc.returncode
+
+
+async def wait_exit(pid: int) -> int | None:
+    """Await a pid we spawned exiting. None if unknown to us."""
+    proc = _processes.get(pid)
+    if proc is None:
+        return None
+    exit_code = await proc.wait()
+    _processes.pop(pid, None)
+    return exit_code
+
+
+async def start_cache(
+    cache_id: UUID, config: CacheConfig, log_path: Path
+) -> CacheProcess:
     run_dir = log_path.parent
     run_dir.mkdir(parents=True, exist_ok=True)
 
@@ -45,22 +71,21 @@ def start_cache(cache_id: UUID, config: CacheConfig, log_path: Path) -> CachePro
 
     log_path = log_path.resolve()
     with log_path.open("ab") as log_file:
-        proc = subprocess.Popen(
+        proc = await anyio.open_process(
             [settings.FASTCACHE_BINARY, config_path],
             cwd=run_dir,
             stdout=log_file,
             stderr=subprocess.STDOUT,
             start_new_session=True,
         )
-    # (pid, create_time) is a reuse-safe identity; fail fast if the process died
-    # immediately.
     try:
         create_time = psutil.Process(proc.pid).create_time()
     except psutil.Error as exc:
         raise RuntimeError(
             f"Cache {cache_id} (pid={proc.pid}) exited immediately after launch; "
             f"check logs at {log_path}"
         ) from exc
+    _processes[proc.pid] = proc
     logger.info("Started cache %s (pid=%d)", cache_id, proc.pid)
     return CacheProcess(pid=proc.pid, create_time=create_time, log_path=log_path)
 
@@ -75,6 +100,8 @@ def resolve_process(pid: int, create_time: float | None) -> psutil.Process | Non
         return None
     try:
         proc = psutil.Process(pid)
+        if proc.status() == psutil.STATUS_ZOMBIE:
+            return None
         if abs(proc.create_time() - create_time) <= _CREATE_TIME_TOLERANCE:
             return proc
     except psutil.Error:
@@ -86,11 +113,27 @@ def is_alive(pid: int, create_time: float | None) -> bool:
     return resolve_process(pid, create_time) is not None
 
 
-def stop_cache(pid: int, create_time: float | None, timeout: float = 5.0) -> None:
-    """SIGTERM the cache process tree, then SIGKILL after ``timeout``.
+async def stop_cache(pid: int, create_time: float | None, timeout: float = 5.0) -> None:
+    """Terminate the cache process, escalating to SIGKILL after ``timeout``.
 
+    Uses anyio Process handle if we still have one; falls back to psutil by
+    (pid, create_time) identity for orphaned/cross-restart cases.
     No-ops unless the identity matches, so a recycled pid is never killed.
     """
+    proc = _processes.get(pid)
+    if proc is not None:
+        proc.terminate()
+        with anyio.move_on_after(timeout):
+            await proc.wait()
+        if proc.returncode is None:
+            logger.warning(
+                "Cache pid=%d still alive after %.1fs; sending SIGKILL", pid, timeout
+            )
+            proc.kill()
+            await proc.wait()
+        _processes.pop(pid, None)
+        return
+
     parent = resolve_process(pid, create_time)
     if parent is None:
         logger.info(

src/fastcache_api/reconcile.py

@@ -1,12 +1,14 @@
 import asyncio
 import logging
+from uuid import UUID
 
+import anyio
 from sqlalchemy import select
 
 from .config import settings
 from .db import SessionLocal
 from .models import CacheState
-from .process import is_alive
+from .process import exit_code, is_alive, wait_exit
 from .tables import Cache
 
 logger = logging.getLogger(__name__)
@@ -16,10 +18,11 @@
 
 
 async def sweep_dead_caches() -> int:
-    """Mark non-final caches whose process is gone as failed; return how many.
+    """Reconcile non-final caches whose process is gone; return how many changed.
 
     Caches outlive the api server, so a row's state is a claim we verify against
-    the live process.
+    the live process. Exit code 0 -> completed, else (or unknown, e.g. after an
+    api restart) -> failed.
     """
     async with SessionLocal() as session:
         result = await session.execute(select(Cache).where(Cache.state.in_(_NON_FINAL)))
@@ -29,12 +32,16 @@ async def sweep_dead_caches() -> int:
         for cache in caches:
             if is_alive(cache.pid, cache.create_time):
                 continue
+            ec = exit_code(cache.pid)
+            cache.state = CacheState.completed if ec == 0 else CacheState.failed
+            cache.exit_code = ec
             logger.warning(
-                "Cache %s (pid=%d) is no longer running; marking failed",
+                "Cache %s (pid=%d) is no longer running (exit_code=%s); marking %s",
                 cache.id,
                 cache.pid,
+                ec,
+                cache.state,
             )
-            cache.state = CacheState.failed
             stale += 1
 
         if stale:
@@ -45,7 +52,30 @@ async def sweep_dead_caches() -> int:
 async def reconcile_caches() -> None:
     """One-shot sweep at startup to reconcile DB state against live processes."""
     stale = await sweep_dead_caches()
-    logger.info("Startup reconcile complete; %d cache(s) marked failed", stale)
+    logger.info("Startup reconcile complete; %d cache(s) reconciled", stale)
+
+
+async def watch_and_record(cache_id: UUID, pid: int) -> None:
+    exit_code = await wait_exit(pid)
+    if exit_code is None:
+        return
+    # we shield from app shutdown - we will lose our db connection
+    # and this won't be able to commit properly
+    with anyio.CancelScope(shield=True):
+        async with SessionLocal() as session:
+            cache = await session.get(Cache, cache_id)
+            if cache is None or CacheState(cache.state).is_final():
+                return
+            cache.state = CacheState.completed if exit_code == 0 else CacheState.failed
+            cache.exit_code = exit_code
+            await session.commit()
+    logger.info(
+        "Cache %s (pid=%d) exited (code=%s); marked %s",
+        cache_id,
+        pid,
+        exit_code,
+        cache.state,
+    )
 
 
 async def monitor_caches() -> None:

src/fastcache_api/routes/cache.py

@@ -9,6 +9,7 @@
 from ..config import settings
 from ..db import get_session
 from ..dependencies import TokenPayload, require_user
+from ..lifecycle import schedule_exit_watch
 from ..models import (
     CacheConfig,
     CachePublic,
@@ -91,7 +92,7 @@ async def create_cache(
 
     cache_id = uuid4()
     try:
-        proc = start_cache(cache_id, config, req.log_path)
+        proc = await start_cache(cache_id, config, req.log_path)
     except (FileNotFoundError, OSError, RuntimeError) as exc:
         raise HTTPException(
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
@@ -115,13 +116,14 @@ async def create_cache(
     except IntegrityError as exc:
         await session.rollback()
         # Don't leave an orphaned process behind for the rejected request.
-        stop_cache(proc.pid, proc.create_time, timeout=0)
+        await stop_cache(proc.pid, proc.create_time, timeout=0)
         raise HTTPException(
             status_code=status.HTTP_409_CONFLICT,
             detail=(f"Cache with transfer_id '{req.transfer_id}' already exists"),
         ) from exc
 
     await session.refresh(cache)
+    schedule_exit_watch(cache.id, proc.pid)
     return cache
 
 

src/fastcache_api/tables.py

@@ -50,6 +50,9 @@ class Cache(DTMixin, Base):
         default=CacheState.new,
         doc="Lifecycle state; stored as the CacheState string value",
     )
+    exit_code: Mapped[int | None] = mapped_column(
+        default=None, doc="OS exit code once the process has ended"
+    )
     log_path: Mapped[str] = mapped_column(
         doc="Absolute path to the cache process log file",
     )