diff --git a/CHANGELOG.md b/CHANGELOG.md index 720c5c6f..c1e20f16 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,6 +37,7 @@ For more information about each release including git tags and artifacts, see [R ### Fixed +- Serialize Postgres schema creation behind an advisory lock so the worker and UI API no longer race on `CREATE TYPE` for the `job_status` enum on a fresh volume, which could crash the worker on first boot ([#436](https://github.com/roostorg/osprey/pull/436) by [@vedarolap](https://github.com/vedarolap), closes [#432](https://github.com/roostorg/osprey/issues/432)) - Escape literal braces when parsing f-strings in the engine ([#347](https://github.com/roostorg/osprey/pull/347) by [@haileyok](https://github.com/haileyok)) - Tolerate malformed URI escapes in `EntityWithPopover` UI component ([#377](https://github.com/roostorg/osprey/pull/377) by [@julietshen](https://github.com/julietshen)) - Add retention limits to Kafka topics to prevent unbounded disk growth ([#249](https://github.com/roostorg/osprey/pull/249) by [@VINODvoid](https://github.com/VINODvoid)) diff --git a/osprey_worker/src/osprey/worker/lib/storage/postgres.py b/osprey_worker/src/osprey/worker/lib/storage/postgres.py index b0f67045..ca386f6e 100644 --- a/osprey_worker/src/osprey/worker/lib/storage/postgres.py +++ b/osprey_worker/src/osprey/worker/lib/storage/postgres.py @@ -19,6 +19,10 @@ metadata = MetaData() Model = declarative_base(name='Model', metadata=metadata) +# Arbitrary key for the Postgres advisory lock guarding schema creation (see init_from_config). +# Any process taking this lock is guaranteed to be alone while it runs metadata.create_all(). +_SCHEMA_CREATE_LOCK_KEY = 87271 + if TYPE_CHECKING: SessionMaker = sessionmaker[Session] # type: ignore[type-var] else: @@ -36,6 +40,24 @@ def _get_or_init_session(database: str) -> SessionMaker: return sessions[database] +def create_schema(engine: sqlalchemy.engine.Engine) -> None: + """Create all tables/types defined in `metadata` against `engine`. + + Multiple processes (e.g. the worker and the UI API) can call this at the same time against + a fresh database. SQLAlchemy's enum creation is check-then-create rather than atomic, so on + a fresh volume both processes can see a type as missing and both issue CREATE TYPE, and the + loser crashes with a UniqueViolation. Take a Postgres advisory lock first so only one process + creates the schema at a time; by the time any other process acquires the lock, create_all's + own existence checks make it a no-op. + """ + with engine.connect() as connection: + connection.execute(sqlalchemy.text('SELECT pg_advisory_lock(:key)'), {'key': _SCHEMA_CREATE_LOCK_KEY}) + try: + metadata.create_all(engine) + finally: + connection.execute(sqlalchemy.text('SELECT pg_advisory_unlock(:key)'), {'key': _SCHEMA_CREATE_LOCK_KEY}) + + def init_from_config(database: str) -> None: def _init(config: Config) -> None: if not config['POSTGRES_HOSTS'].get(database): @@ -59,8 +81,7 @@ def _init(config: Config) -> None: temporary_ability_token, ) - # Create all tables defined in the metadata - metadata.create_all(new_engine) + create_schema(new_engine) CONFIG.instance().register_configuration_callback(_init) diff --git a/osprey_worker/src/osprey/worker/lib/storage/tests/test_postgres.py b/osprey_worker/src/osprey/worker/lib/storage/tests/test_postgres.py new file mode 100644 index 00000000..ee0e86ec --- /dev/null +++ b/osprey_worker/src/osprey/worker/lib/storage/tests/test_postgres.py @@ -0,0 +1,57 @@ +import threading + +import sqlalchemy +from osprey.worker.lib.singletons import CONFIG +from osprey.worker.lib.storage import postgres +from psycopg2.errors import DuplicateDatabase, InvalidCatalogName +from sqlalchemy.exc import ProgrammingError +from sqlalchemy_utils import create_database, drop_database + + +def test_create_schema_survives_concurrent_callers_on_a_fresh_database(): + """Regression test for issue #432: on a fresh database, two processes (e.g. the worker and + the UI API) both calling create_schema() at startup used to be able to race on `CREATE TYPE` + for the job_status enum, since SQLAlchemy's enum creation is check-then-create rather than + atomic. Both would see the type as missing, both would issue CREATE TYPE, and the loser would + crash with a UniqueViolation on `pg_type_typname_nsp_index`. + + Simulate two racing processes with two independent engines hitting a brand new database at + the same time, and assert neither raises. + """ + base_url = CONFIG.instance()['POSTGRES_HOSTS']['osprey_db'] + fresh_url = base_url.rsplit('/', 1)[0] + '/osprey_test_concurrent_schema_create' + + try: + drop_database(fresh_url) + except ProgrammingError as e: + if not isinstance(e.orig, InvalidCatalogName): + raise + try: + create_database(fresh_url) + except ProgrammingError as e: + if not isinstance(e.orig, DuplicateDatabase): + raise + + try: + errors: list[Exception] = [] + barrier = threading.Barrier(2) + + def _create_schema() -> None: + engine = sqlalchemy.create_engine(fresh_url) + try: + barrier.wait(timeout=5) + postgres.create_schema(engine) + except Exception as e: + errors.append(e) + finally: + engine.dispose() + + threads = [threading.Thread(target=_create_schema) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + + assert not errors, errors + finally: + drop_database(fresh_url)