Migrate compute service to rqueue

.env.example

@@ -14,3 +14,9 @@ APP_DB_PASSWORD=
 
 # Browser-reachable MinIO endpoint for output download URLs.
 MINIO_PUBLIC_ENDPOINT=localhost:9000
+
+# Task queue. The queue name and the schema rqueue owns; both reach the
+# migration step and every process that reads or writes the queue, so a value
+# set here applies everywhere or nowhere.
+COMPUTE_QUEUE=simulations
+COMPUTE_QUEUE_SCHEMA=task_queue

ARCHITECTURE.md

@@ -16,7 +16,7 @@ SvelteKit web app
   |
   v
 FastAPI compute service
-  |-- compute.jobs and the Procrastinate queue
+  |-- compute.jobs and the rqueue task_queue schema
   |-- worker -> tsdhn engine -> MinIO
 ```
 
@@ -48,7 +48,7 @@ job ID, a compute-service selector, or output storage keys.
 The compute service owns:
 
 - the API used by the web server;
-- `compute.jobs` and the Procrastinate queue tables;
+- `compute.jobs` and the `task_queue` schema rqueue owns;
 - the internal compute job ID;
 - job progress, retry state, and worker heartbeats;
 - simulation work directories and checkpoints;
@@ -93,8 +93,8 @@ submission_error
 created_at
 ```
 
-The compute service creates and writes `compute.jobs` and the Procrastinate
-queue tables. `compute.jobs.simulation_id` links a compute job to the web
+The compute service creates and writes `compute.jobs` and the `task_queue`
+schema rqueue owns. `compute.jobs.simulation_id` links a compute job to the web
 simulation. The value is unique because repeating a submission must return the
 same compute job.
 
@@ -191,3 +191,30 @@ The engine records enough state to continue valid completed work after a
 retry. The compute service keeps the job in `compute.jobs` and reports the
 final failure when retries are exhausted. Deployment settings determine retry
 limits, worker recovery, storage, and cleanup.
+
+A worker asked to stop gracefully stops taking new work and gives what it is
+already running a short grace period. A simulation runs far longer than that, so
+it is cancelled and its claim is handed straight back for another worker, which
+resumes from the checkpoints in the job's work directory.
+
+A worker that dies mid-run loses its lease, and the queue reclaims the job
+without asking the dead worker anything. When that happens on the job's last
+attempt the queue records the failure by itself, so no running code is left to
+update `compute.jobs`. The worker process therefore reconciles: it periodically
+finds jobs the queue has finished that `compute.jobs` still shows as running,
+and marks them failed with an error saying the status was reconciled rather
+than reported by the run. A job that reported its own outcome is never
+overwritten.
+
+Because the simulation runs on a thread that the service cannot stop on demand,
+a job records which attempt currently owns it, and every write an attempt makes
+is accepted only if that attempt still owns the job and the job is not already
+finished. A late write from an attempt that has been replaced, or from one whose
+job has already been reconciled, is refused rather than applied.
+
+The same reasoning covers the job's working directory, which holds the
+checkpoints a retry resumes from. An attempt holds an exclusive claim on that
+directory for as long as its simulation is actually running, so a replacement
+never reads checkpoints another attempt is still writing. If the worker process
+dies the claim is released with it, and the replacement resumes normally; if the
+previous attempt is still running, the replacement waits and retries instead.

docker-compose.yml

@@ -44,13 +44,14 @@ services:
         condition: service_healthy
     environment:
       COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn
+      COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue}
       APP_DB_ROLE: ${APP_DB_ROLE:-tsdhn_app}
       APP_DB_PASSWORD: ${APP_DB_PASSWORD:?set APP_DB_PASSWORD in .env}
     command:
       [
         "sh",
         "-lc",
-        "uv run --no-dev tsdhn-compute-migrate && uv run --no-dev tsdhn-procrastinate-migrate",
+        "uv run --no-dev tsdhn-compute-migrate && uv run --no-dev rqueue --database-url \"$$COMPUTE_DATABASE_URL\" --schema \"$$COMPUTE_QUEUE_SCHEMA\" migrate",
       ]
     restart: "no"
 
@@ -69,6 +70,8 @@ services:
         condition: service_healthy
     environment:
       COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn
+      COMPUTE_QUEUE: ${COMPUTE_QUEUE:-simulations}
+      COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue}
       MINIO_ENDPOINT: minio:9000
       MINIO_PUBLIC_ENDPOINT: ${MINIO_PUBLIC_ENDPOINT:-localhost:9000}
       MINIO_ACCESS_KEY: ${MINIO_ACCESS_KEY:-minioadmin}
@@ -94,6 +97,8 @@ services:
         condition: service_healthy
     environment:
       COMPUTE_DATABASE_URL: postgresql://tsdhn:tsdhn@postgres:5432/tsdhn
+      COMPUTE_QUEUE: ${COMPUTE_QUEUE:-simulations}
+      COMPUTE_QUEUE_SCHEMA: ${COMPUTE_QUEUE_SCHEMA:-task_queue}
       MINIO_ENDPOINT: minio:9000
       MINIO_ACCESS_KEY: ${MINIO_ACCESS_KEY:-minioadmin}
       MINIO_SECRET_KEY: ${MINIO_SECRET_KEY:-minioadmin}

mise.toml

@@ -155,7 +155,7 @@ description = "Apply compute, queue, and web migrations to local PostgreSQL"
 depends = ["db:start"]
 run = [
   "uv run tsdhn-compute-migrate",
-  "uv run tsdhn-procrastinate-migrate",
+  "uv run rqueue --database-url \"${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn}\" --schema \"${COMPUTE_QUEUE_SCHEMA:-task_queue}\" migrate",
   "DATABASE_URL=\"${COMPUTE_DATABASE_URL:-postgresql://tsdhn:tsdhn@127.0.0.1:5432/tsdhn}\" bun --filter web db:migrate",
   "uv run tsdhn-web-grants",
 ]

packages/api/api/core/db.py

@@ -1,84 +1,188 @@
-"""Database connections for API reads, workers, and job notifications."""
+"""The process-wide asyncpg pool, shared by the API and the worker.
+
+Both processes open one pool in `open_pool(...)` and take a connection from it
+per statement (or per transaction). Nothing holds a connection across a
+simulation run: the only long-lived borrow in the system is the one
+`rqueue.Worker` makes to LISTEN on its wake channel, which is why
+`settings.worker_pool_size()` sizes the worker's pool the way it does.
+"""
 
-import threading
 import uuid
-from collections.abc import Iterator
-from contextlib import contextmanager
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from typing import Any
 
-import psycopg
-from psycopg.rows import dict_row
-from psycopg_pool import ConnectionPool
+import asyncpg
 
 from api.core.errors import TransientInfraError
-from api.core.settings import (
-    COMPUTE_DATABASE_URL,
-    DB_POOL_MAX_SIZE,
-    DB_POOL_MIN_SIZE,
-)
+from api.core.settings import COMPUTE_DATABASE_URL
 
 __all__ = [
     "CONNECT_TIMEOUT",
     "JobRow",
+    "acquire",
     "close_pool",
     "connect",
     "get_pool",
+    "is_transient",
     "notify_channel",
-    "pooled",
+    "open_pool",
+    "transient_connection_errors",
 ]
 
 JobRow = dict[str, Any]
 CONNECT_TIMEOUT = 2
 
-_pool: ConnectionPool[psycopg.Connection[JobRow]] | None = None
-_pool_lock = threading.Lock()
-
-
-def _new_pool() -> ConnectionPool[psycopg.Connection[JobRow]]:
-    return ConnectionPool(
-        COMPUTE_DATABASE_URL,
-        min_size=DB_POOL_MIN_SIZE,
-        max_size=DB_POOL_MAX_SIZE,
-        kwargs={"row_factory": dict_row, "connect_timeout": CONNECT_TIMEOUT},
-        open=False,
-    )
-
-
-def get_pool() -> ConnectionPool[psycopg.Connection[JobRow]]:
-    """Return the process-wide pool, creating it on first use."""
+_pool: asyncpg.Pool | None = None
+
+# Client-side failures: the server was never reached, or the socket died.
+_CLIENT_ERRORS = (ConnectionError, OSError, TimeoutError)
+
+# Server-side failures are classified by SQLSTATE *class* rather than by
+# exception type. An enumerated tuple of asyncpg classes is the wrong basis:
+# it looked complete and was not (`TooManyConnectionsError` is an
+# `InsufficientResourcesError`, not a `PostgresConnectionError`, so pool
+# exhaustion read as permanent), and it grows by one entry per incident.
+#
+#   08  connection_exception     the connection broke or was refused
+#   53  insufficient_resources   too many connections, out of memory, disk full
+#   57  operator_intervention    admin shutdown, crash shutdown, cannot connect now
+#   40  transaction_rollback     serialization failure, deadlock detected
+#
+# The asymmetry justifies the width. Classifying a transient failure as
+# permanent throws away a simulation that would have succeeded -- tens of
+# minutes of compute, and an operator retry to get it back. Classifying a
+# permanent failure as transient costs two extra attempts and roughly 45s of
+# backoff before it fails anyway with the same error. So the few permanent
+# members inside these classes (57P04 database_dropped, say) are worth carrying
+# to catch every transient one.
+#
+# Deliberately *not* included: 22 (data exception), 23 (integrity constraint),
+# 42 (syntax and access rules). Those are application bugs, and a retry there
+# re-runs a whole simulation to reach the identical failure.
+_TRANSIENT_SQLSTATE_CLASSES = frozenset({"08", "40", "53", "57"})
+
+
+def is_transient(exc: BaseException) -> bool:
+    """Whether this failure could plausibly succeed on a later attempt."""
+    if isinstance(exc, _CLIENT_ERRORS):
+        return True
+    sqlstate = getattr(exc, "sqlstate", None)
+    return isinstance(sqlstate, str) and sqlstate[:2] in _TRANSIENT_SQLSTATE_CLASSES
+
+
+async def open_pool(*, min_size: int, max_size: int) -> asyncpg.Pool:
+    """Create the process-wide pool. Idempotent within one process."""
     global _pool
-    with _pool_lock:
-        if _pool is None or _pool.closed:
-            _pool = _new_pool()
-            _pool.open(wait=False)
-        return _pool
+    if _pool is None:
+        _pool = await asyncpg.create_pool(
+            COMPUTE_DATABASE_URL,
+            min_size=min_size,
+            max_size=max_size,
+            timeout=CONNECT_TIMEOUT,
+        )
+    return _pool
 
 
-def close_pool() -> None:
+async def close_pool() -> None:
     global _pool
-    with _pool_lock:
-        if _pool is not None and not _pool.closed:
-            _pool.close()
+    if _pool is not None:
+        await _pool.close()
         _pool = None
 
 
-@contextmanager
-def pooled() -> Iterator[psycopg.Connection[JobRow]]:
-    """Provide a short-lived connection for API queries."""
-    with get_pool().connection() as conn:
-        yield conn
+def get_pool() -> asyncpg.Pool:
+    """Return the open pool, or explain that startup did not open one."""
+    if _pool is None:
+        raise RuntimeError("database pool is not open; call db.open_pool() first")
+    return _pool
+
+
+@asynccontextmanager
+async def transient_connection_errors(
+    connection: asyncpg.Connection,
+) -> AsyncIterator[None]:
+    """Translate a connection lost mid-statement into `TransientInfraError`.
+
+    `acquire()` covers only the borrow, which is not where a database restart
+    under a running simulation actually lands. asyncpg reports a connection
+    that dies around a statement two different ways, and only one of them was
+    already covered:
+
+    - the backend dies with the statement in flight -> `ConnectionDoesNotExist`,
+      SQLSTATE 08003, which `is_transient` recognises;
+    - the backend is already gone when the statement is issued -> a bare
+      `InterfaceError("connection is closed")`, which is not a connection
+      error class at all.
+
+    `InterfaceError` is also what asyncpg raises for a genuine programming
+    error ("the server expects 2 arguments for this query, 1 was passed"), so
+    it is translated only when the connection really is gone. That distinction
+    is load bearing: `TRANSIENT_RETRY` retries `TransientInfraError`, so
+    mistranslating a bug would spend a job's whole retry budget re-running
+    something that cannot succeed. A programming error leaves the connection
+    open and propagates unchanged.
+    """
+    try:
+        yield
+    except asyncpg.InterfaceError as e:
+        if not _connection_is_gone(connection):
+            raise
+        raise TransientInfraError("database unavailable") from e
+    except Exception as e:
+        if not is_transient(e):
+            raise
+        raise TransientInfraError("database unavailable") from e
+
+
+def _connection_is_gone(connection: asyncpg.Connection) -> bool:
+    """Whether this connection can still be used for anything at all.
+
+    Not simply `is_closed()`. When a pooled connection's backend dies, the
+    pool terminates the connection and takes the proxy back, after which
+    *every* method on that proxy raises `InterfaceError("cannot call
+    Connection.is_closed(): connection has been released back to the pool")` --
+    the question included. A proxy that cannot answer is certainly gone,
+    whereas a connection that merely rejected a malformed call answers
+    `False` and keeps working.
+    """
+    try:
+        return bool(connection.is_closed())
+    except asyncpg.InterfaceError:
+        return True
+
+
+@asynccontextmanager
+async def acquire() -> AsyncIterator[asyncpg.Connection]:
+    """Borrow a pooled connection for one statement or one transaction."""
+    pool = get_pool()
+    try:
+        connection = await pool.acquire()
+    except Exception as e:
+        # Pool exhaustion arrives here as TooManyConnectionsError (53300), which
+        # is every bit as retryable as a refused socket.
+        if not is_transient(e):
+            raise
+        raise TransientInfraError("database unavailable") from e
+    try:
+        yield connection
+    finally:
+        await pool.release(connection)
+
 
+async def connect() -> asyncpg.Connection:
+    """Open a connection outside the pool.
 
-def connect() -> psycopg.Connection[JobRow]:
-    """Open a dedicated worker connection."""
+    The SSE endpoint holds one of these for the lifetime of a stream, which can
+    be half an hour. Taking that from the request pool would let a handful of
+    watching browsers starve every other route.
+    """
     try:
-        return psycopg.connect(
-            COMPUTE_DATABASE_URL,
-            connect_timeout=CONNECT_TIMEOUT,
-            row_factory=dict_row,
-        )
-    except psycopg.OperationalError as e:
-        raise TransientInfraError("database connection failed") from e
+        return await asyncpg.connect(COMPUTE_DATABASE_URL, timeout=CONNECT_TIMEOUT)
+    except Exception as e:
+        if not is_transient(e):
+            raise
+        raise TransientInfraError("database unavailable") from e
 
 
 def notify_channel(simulation_id: uuid.UUID) -> str: