Stop the whole fleet stampeding one advisory lock at boot

compose.yml

@@ -26,7 +26,14 @@ services:
     container_name: shepherd_db
     environment:
       POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-supersecretpassw0rd}
-    command: postgres -c max_connections=200
+    # Connection budget: every container holds its own psycopg pool, so the
+    # fleet ceiling is (containers x pool max size). 22 workers/monitor at the
+    # default max of 10, plus shepherd_server at 30, is 250 -- which overran
+    # the old 200 and surfaced as "sorry, too many clients already" under load
+    # even though the comments in this file and shepherd_utils/config.py both
+    # say the sum must stay under it. 300 restores the headroom the docs claim.
+    # Adding a worker service? Re-do this arithmetic.
+    command: postgres -c max_connections=300
     build:
       context: .
       dockerfile: shepherd_db/Dockerfile

shepherd_utils/config.py

@@ -55,7 +55,7 @@ class Settings(BaseSettings):
     postgres_pool_timeout: float = 5.0
     # Per-process Postgres pool bounds. Every container (server + each worker)
     # holds its own pool, so the fleet-wide ceiling is (number of containers x
-    # max size) and must stay under Postgres's max_connections (200 in
+    # max size) and must stay under Postgres's max_connections (300 in
     # compose.yml). The server fields all client HTTP traffic (sync-query
     # status polling, /callback lookups) and is the component that exhausts
     # its pool first under load, so compose.yml/Helm give it a larger pool via

shepherd_utils/db.py

@@ -163,9 +163,17 @@ def decompress_zstd(blob: bytes) -> bytes:
 # predates a schema addition never pick it up from there; re-running these
 # here upgrades them in place. Everything in this list must be safe to re-run
 # and effectively free once already applied.
+# ``(index name, DDL)``. The name is what the pre-flight check below looks for
+# in the catalog, so it must match the index the DDL creates.
 _SCHEMA_UPGRADES = (
-    "CREATE INDEX IF NOT EXISTS idx_callbacks_callback_id ON callbacks (callback_id)",
-    "CREATE INDEX IF NOT EXISTS idx_callbacks_query_id ON callbacks (query_id)",
+    (
+        "idx_callbacks_callback_id",
+        "CREATE INDEX IF NOT EXISTS idx_callbacks_callback_id ON callbacks (callback_id)",
+    ),
+    (
+        "idx_callbacks_query_id",
+        "CREATE INDEX IF NOT EXISTS idx_callbacks_query_id ON callbacks (query_id)",
+    ),
 )
 
 # Arbitrary-but-fixed advisory lock id serializing the upgrades across the
@@ -177,25 +185,50 @@ def decompress_zstd(blob: bytes) -> bytes:
 async def apply_schema_upgrades() -> None:
     """Bring an existing database up to date with init_db.sql additions."""
     async with pool.connection(settings.postgres_pool_timeout) as conn:
+        # Pre-flight catalog check, deliberately OUTSIDE the advisory lock.
+        # Every container in the stack runs this at boot, and they all boot the
+        # instant Postgres reports healthy, so taking the lock unconditionally
+        # made ~23 containers queue up on a single lock for work that is a
+        # no-op on any volume created since these indexes landed in
+        # init_db.sql. Whoever lost that queue blew ``postgres_pool_timeout``
+        # and logged a PoolTimeout traceback on an otherwise healthy startup.
+        # The check is a single indexed catalog read and does not serialize, so
+        # the common "already applied" case now costs one query and no lock.
+        cursor = await conn.execute(
+            "SELECT count(*) FROM pg_class WHERE relkind = 'i' AND relname = ANY(%s)",
+            ([name for name, _ in _SCHEMA_UPGRADES],),
+        )
+        row = await cursor.fetchone()
+        if row is not None and row[0] == len(_SCHEMA_UPGRADES):
+            return
         await conn.execute(
             "SELECT pg_advisory_xact_lock(%s)", (_SCHEMA_UPGRADE_LOCK_ID,)
         )
-        for ddl in _SCHEMA_UPGRADES:
+        for _, ddl in _SCHEMA_UPGRADES:
             await conn.execute(ddl)
         await conn.commit()
 
 
 async def initialize_db() -> None:
     """Open connection and create db."""
     await pool.open()
-    try:
-        await apply_schema_upgrades()
-    except Exception:
-        # A failed upgrade must never keep a worker from starting: the schema
-        # additions are performance aids, and the janitor/next boot retries.
-        logging.getLogger("shepherd.db").warning(
-            "Failed to apply startup schema upgrades", exc_info=True
-        )
+    for attempt in range(PG_RETRIES):
+        try:
+            await apply_schema_upgrades()
+            return
+        except Exception:
+            # Retry with the same backoff the query paths use. Workers boot the
+            # moment the DB reports healthy, so a Postgres crash-restart (or
+            # any blip in the first seconds) otherwise burned the single
+            # attempt and every container logged a traceback at once.
+            if attempt == PG_RETRIES - 1:
+                break
+            await asyncio.sleep(0.1 * (2**attempt))
+    # A failed upgrade must never keep a worker from starting: the schema
+    # additions are performance aids, and the janitor/next boot retries.
+    logging.getLogger("shepherd.db").warning(
+        "Failed to apply startup schema upgrades", exc_info=True
+    )
 
 
 async def shutdown_db() -> None:

tests/unit/test_db_postgres.py

@@ -298,17 +298,50 @@ async def test_initialize_db_applies_callback_indexes(mocker):
     assert mock_conn.commit.called
 
 
+@pytest.mark.asyncio
+async def test_initialize_db_skips_upgrades_when_indexes_present(mocker):
+    """The pre-flight catalog check short-circuits the common case: when every
+    upgrade index already exists, no advisory lock is taken and no DDL runs, so
+    a whole fleet booting at once doesn't queue on one lock for a no-op."""
+    mock_conn, mock_pool = _install_pool_mock(
+        mocker, cursor_fetchone=(len(db._SCHEMA_UPGRADES),)
+    )
+    await db.initialize_db()
+    assert mock_pool.open.called
+    executed = " ".join(str(c.args[0]) for c in mock_conn.execute.call_args_list)
+    assert "pg_advisory_xact_lock" not in executed
+    assert "CREATE INDEX" not in executed
+
+
 @pytest.mark.asyncio
 async def test_initialize_db_survives_schema_upgrade_failure(mocker):
     """A failed upgrade (e.g. transient connection error) is logged, not
     raised -- workers must still start."""
+    mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock)
     mock_conn, mock_pool = _install_pool_mock(
         mocker, raise_on_execute=OperationalError("pg not ready")
     )
     await db.initialize_db()
     assert mock_pool.open.called
 
 
+@pytest.mark.asyncio
+async def test_initialize_db_retries_schema_upgrades(mocker):
+    """A blip on the first attempt (the DB is still coming up) is retried with
+    backoff rather than burning the one attempt and logging a traceback."""
+    sleep = mocker.patch("asyncio.sleep", new_callable=mocker.AsyncMock)
+    apply = mocker.patch.object(
+        db,
+        "apply_schema_upgrades",
+        new_callable=mocker.AsyncMock,
+        side_effect=[OperationalError("pg in recovery"), None],
+    )
+    mocker.patch.object(db, "pool", AsyncMock(spec=AsyncConnectionPool))
+    await db.initialize_db()
+    assert apply.await_count == 2
+    assert sleep.await_count == 1
+
+
 # --- check_connection -----------------------------------------------------