diff --git a/src/fastcache_api/models.py b/src/fastcache_api/models.py index 0135923..40a3e52 100644 --- a/src/fastcache_api/models.py +++ b/src/fastcache_api/models.py @@ -50,7 +50,9 @@ def to_fastcache_json(self, indent: int = 2) -> str: class CacheRequest(BaseModel): - transfer_id: str + # The cache's dedup/lookup identity. Could be per-transfer (transfer ID) + # or per-experiment (experiment name) + key: str # Human who initiated the transfer upstream (bearer token is a shared # service identity, so attribution must travel in the request body). requested_by: str @@ -62,7 +64,7 @@ class CachePublic(BaseModel): model_config = ConfigDict(from_attributes=True) id: UUID - transfer_id: str + key: str | None user: str state: CacheState exit_code: int | None diff --git a/src/fastcache_api/reconcile.py b/src/fastcache_api/reconcile.py index 02dbfca..3aa3385 100644 --- a/src/fastcache_api/reconcile.py +++ b/src/fastcache_api/reconcile.py @@ -35,6 +35,7 @@ async def sweep_dead_caches() -> int: ec = exit_code(cache.pid) cache.state = CacheState.completed if ec == 0 else CacheState.failed cache.exit_code = ec + cache.key = None # Free the key! logger.warning( "Cache %s (pid=%d) is no longer running (exit_code=%s); marking %s", cache.id, @@ -68,6 +69,7 @@ async def watch_and_record(cache_id: UUID, pid: int) -> None: return cache.state = CacheState.completed if exit_code == 0 else CacheState.failed cache.exit_code = exit_code + cache.key = None await session.commit() logger.info( "Cache %s (pid=%d) exited (code=%s); marked %s", diff --git a/src/fastcache_api/routes/cache.py b/src/fastcache_api/routes/cache.py index e3545df..1dcf2e1 100644 --- a/src/fastcache_api/routes/cache.py +++ b/src/fastcache_api/routes/cache.py @@ -1,13 +1,20 @@ from typing import Annotated from uuid import UUID, uuid4 -from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, status +from fastapi import ( + APIRouter, + BackgroundTasks, + Depends, + HTTPException, + Response, + status, +) from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from ..config import settings -from ..db import get_session +from ..db import SessionLocal, get_session from ..dependencies import TokenPayload, require_user from ..lifecycle import schedule_exit_watch from ..models import ( @@ -20,6 +27,8 @@ from ..process import ( allocate_port_pair, canonical_hostname, + exit_code, + is_alive, ports_in_use, start_cache, stop_cache, @@ -33,6 +42,29 @@ ) +async def _find_active_by_key(session: AsyncSession, key: str) -> Cache | None: + result = await session.execute( + select(Cache).where( + Cache.key == key, + Cache.state.in_([s.value for s in CacheState if not s.is_final()]), + ) + ) + candidate = result.scalar_one_or_none() + if candidate is None: + return None + if is_alive(candidate.pid, candidate.create_time): + return candidate + # The row hasn't been reconciled yet (exit watcher/sweep haven't run), + # but the process is already gone. Finalize here and free the key so + # caller creates a fresh one instead. + ec = exit_code(candidate.pid) + candidate.state = CacheState.completed if ec == 0 else CacheState.failed + candidate.exit_code = ec + candidate.key = None + await session.commit() + return None + + @router.get("/", response_model=CachesPublic) async def get_caches( session: Annotated[AsyncSession, Depends(get_session)], @@ -60,7 +92,13 @@ async def create_cache( req: CacheRequest, user: Annotated[TokenPayload, Depends(require_user)], session: Annotated[AsyncSession, Depends(get_session)], + response: Response, ): + existing = await _find_active_by_key(session, req.key) + if existing is not None: + response.status_code = status.HTTP_200_OK + return existing + # The cache always runs as a local subprocess of this api server, so the # ZMQ URIs are published under this host's canonical (FQDN) name. hostname = canonical_hostname() @@ -101,7 +139,7 @@ async def create_cache( cache = Cache( id=cache_id, - transfer_id=req.transfer_id, + key=req.key, pid=proc.pid, create_time=proc.create_time, user=req.requested_by, @@ -115,18 +153,36 @@ async def create_cache( await session.commit() except IntegrityError as exc: await session.rollback() - # Don't leave an orphaned process behind for the rejected request. + # Lost a race with a concurrent request for the same key: don't + # leave our process running, join the winner instead. 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 + winner = await _find_active_by_key(session, req.key) + if winner is None: + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"Cache key '{req.key}' conflict", + ) from exc + response.status_code = status.HTTP_200_OK + return winner await session.refresh(cache) schedule_exit_watch(cache.id, proc.pid) + response.status_code = status.HTTP_201_CREATED return cache +async def _teardown_and_free_key( + cache_id: UUID, pid: int, create_time: float | None +) -> None: + await stop_cache(pid, create_time, settings.SHUTDOWN_GRACE_SECONDS) + # Only free the key once the process is confirmed stopped + async with SessionLocal() as session: + cache = await session.get(Cache, cache_id) + if cache is not None: + cache.key = None + await session.commit() + + @router.delete("/{cache_id}", response_model=CachePublic) async def shutdown_cache( cache_id: UUID, @@ -145,6 +201,6 @@ async def shutdown_cache( # Tear down the process tree in the background so DELETE returns promptly. background_tasks.add_task( - stop_cache, cache.pid, cache.create_time, settings.SHUTDOWN_GRACE_SECONDS + _teardown_and_free_key, cache.id, cache.pid, cache.create_time ) return cache diff --git a/src/fastcache_api/tables.py b/src/fastcache_api/tables.py index fa97719..f93b9c2 100644 --- a/src/fastcache_api/tables.py +++ b/src/fastcache_api/tables.py @@ -39,7 +39,16 @@ class Cache(DTMixin, Base): __tablename__ = "caches" id: Mapped[UUID] = mapped_column(default=uuid4, primary_key=True) - transfer_id: Mapped[str] = mapped_column(unique=True, doc="Unique transfer ID") + key: Mapped[str | None] = mapped_column( + unique=True, + default=None, + doc=( + "Dedup/lookup ID from the request (e.g. a transfer id for " + "a one-off cache, or a shared key like an experiment name for a " + "joinable one). Cleared when the cache is torn down, freeing the " + "key for reuse." + ), + ) pid: Mapped[int] = mapped_column(doc="PID of the cache process") create_time: Mapped[float | None] = mapped_column( default=None,