diff --git a/compose.yml b/compose.yml index 4a80742..25ce27b 100644 --- a/compose.yml +++ b/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 diff --git a/shepherd_utils/config.py b/shepherd_utils/config.py index 78e9907..d369ca6 100644 --- a/shepherd_utils/config.py +++ b/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 diff --git a/shepherd_utils/db.py b/shepherd_utils/db.py index dfdff43..6260ba8 100644 --- a/shepherd_utils/db.py +++ b/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,10 +185,26 @@ 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() @@ -188,14 +212,23 @@ async def apply_schema_upgrades() -> None: 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: diff --git a/tests/unit/test_db_postgres.py b/tests/unit/test_db_postgres.py index 9de7764..2f646da 100644 --- a/tests/unit/test_db_postgres.py +++ b/tests/unit/test_db_postgres.py @@ -298,10 +298,26 @@ 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") ) @@ -309,6 +325,23 @@ async def test_initialize_db_survives_schema_upgrade_failure(mocker): 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 -----------------------------------------------------