diff --git a/src/palace/manager/sqlalchemy/refresh_equivalents.py b/src/palace/manager/sqlalchemy/refresh_equivalents.py index 1042335b45..5281bd7408 100644 --- a/src/palace/manager/sqlalchemy/refresh_equivalents.py +++ b/src/palace/manager/sqlalchemy/refresh_equivalents.py @@ -12,7 +12,11 @@ from __future__ import annotations +from collections.abc import Collection, Iterable +from itertools import batched + from sqlalchemy import and_, delete, select, union +from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.orm import Session from palace.manager.sqlalchemy.model.identifier import ( @@ -21,6 +25,43 @@ RecursiveEquivalencyCache, ) +# Number of rows sent per INSERT statement. +INSERT_CHUNK_SIZE = 10_000 + + +def _insert_cache_rows(session: Session, rows: Iterable[dict[str, int]]) -> None: + """ + Insert (parent_identifier_id, identifier_id) rows into the cache. + + Rows that already exist are skipped rather than raising a unique violation. + We are not the only writer of this table: the ``Identifier`` creation + listener inserts self-reference rows from webapp transactions, and a + redelivered Celery task can run a second refresh chain concurrently with the + first. Because every parent's rows are deleted and recomputed from the same + source data, a row a concurrent transaction commits under us holds the same + value we were about to write, so ignoring the conflict converges on the + correct chain instead of aborting the whole batch. + """ + for chunk in batched(rows, INSERT_CHUNK_SIZE): + session.execute( + pg_insert(RecursiveEquivalencyCache) + .values(list(chunk)) + .on_conflict_do_nothing( + index_elements=["parent_identifier_id", "identifier_id"] + ) + ) + + +def _delete_cache_rows(session: Session, parent_ids: Collection[int]) -> None: + """Delete every cache row belonging to any of the given parent identifiers.""" + if not parent_ids: + return + session.execute( + delete(RecursiveEquivalencyCache).where( + RecursiveEquivalencyCache.parent_identifier_id.in_(parent_ids) + ) + ) + def process_identifier_ids(session: Session, identifier_ids: frozenset[int]) -> None: """ @@ -58,25 +99,23 @@ def process_identifier_ids(session: Session, identifier_ids: frozenset[int]) -> .add_columns(Identifier.id) ) chained_identifiers = session.execute(qu).fetchall() + if not chained_identifiers: + return - # Delete old cache entries for every affected parent, then insert fresh ones. - completed: set[int] = set() - new_rows: list[RecursiveEquivalencyCache] = [] - for link_id, parent_id in chained_identifiers: - if parent_id not in completed: - session.execute( - delete(RecursiveEquivalencyCache).where( - RecursiveEquivalencyCache.parent_identifier_id == parent_id - ) - ) - new_rows.append( - RecursiveEquivalencyCache( - parent_identifier_id=parent_id, identifier_id=link_id - ) - ) - completed.add(parent_id) - - session.add_all(new_rows) + # Delete the old cache entries for every affected parent, then insert the + # fresh ones. The parents are deleted in a single statement, in a stable + # (sorted) order, so that two refreshes racing on overlapping chains are + # less likely to deadlock against each other. + _delete_cache_rows( + session, sorted({parent_id for _, parent_id in chained_identifiers}) + ) + _insert_cache_rows( + session, + ( + {"parent_identifier_id": parent_id, "identifier_id": link_id} + for link_id, parent_id in chained_identifiers + ), + ) def add_identity_equivalents(session: Session, batch_size: int = 200) -> None: @@ -86,6 +125,10 @@ def add_identity_equivalents(session: Session, batch_size: int = 200) -> None: This ensures that queries against the cache always return at least the identifier itself, even when it has no equivalencies. + + An Identifier created by another transaction after this scan begins gets its + self-reference from the ``Identifier`` creation listener, which may commit + before we do; the insert tolerates that (see :func:`_insert_cache_rows`). """ missing_q = ( select(Identifier.id) @@ -100,11 +143,13 @@ def add_identity_equivalents(session: Session, batch_size: int = 200) -> None: .execution_options(yield_per=batch_size) ) - for (identifier_id,) in session.execute(missing_q): - session.add( - RecursiveEquivalencyCache( - parent_identifier_id=identifier_id, identifier_id=identifier_id - ) + for partition in session.execute(missing_q).partitions(): + _insert_cache_rows( + session, + ( + {"parent_identifier_id": identifier_id, "identifier_id": identifier_id} + for (identifier_id,) in partition + ), ) diff --git a/tests/manager/sqlalchemy/test_refresh_equivalents.py b/tests/manager/sqlalchemy/test_refresh_equivalents.py index 938a824f6f..8e8283625e 100644 --- a/tests/manager/sqlalchemy/test_refresh_equivalents.py +++ b/tests/manager/sqlalchemy/test_refresh_equivalents.py @@ -1,8 +1,11 @@ +import pytest + from palace.manager.sqlalchemy.model.identifier import ( Equivalency, RecursiveEquivalencyCache, ) from palace.manager.sqlalchemy.refresh_equivalents import ( + _insert_cache_rows, add_identity_equivalents, process_identifier_ids, refresh_equivalent_identifiers, @@ -89,6 +92,35 @@ def test_replaces_stale_cache( assert b.id not in recursive_equivalency_cache.cache_for(a.id) + def test_tolerates_rows_the_delete_did_not_remove( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # A concurrent writer can commit a cache row for a parent after we have + # already deleted that parent's rows, so the rows we are about to insert + # may collide. Simulate that by disabling the delete entirely: the insert + # must skip the existing rows rather than raise a unique violation. + a = db.identifier() + b = db.identifier() + db.session.add(Equivalency(input_id=a.id, output_id=b.id, strength=1.0)) + db.session.flush() + + process_identifier_ids(db.session, frozenset([a.id, b.id])) + db.session.flush() + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + + monkeypatch.setattr( + "palace.manager.sqlalchemy.refresh_equivalents._delete_cache_rows", + lambda session, parent_ids: None, + ) + process_identifier_ids(db.session, frozenset([a.id, b.id])) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + assert recursive_equivalency_cache.cache_for(b.id) == {a.id, b.id} + class TestAddIdentityEquivalents: def test_adds_self_references( @@ -121,6 +153,39 @@ def test_skips_existing(self, db: DatabaseTransactionFixture) -> None: assert after == before +class TestInsertCacheRows: + def test_skips_existing_rows( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + ) -> None: + # A row committed by a concurrent writer — the Identifier creation + # listener, or a second refresh chain — must not abort the insert. + a = db.identifier() + b = db.identifier() + db.session.flush() + # a and b already have self-references from the creation listener. + + rows = [ + {"parent_identifier_id": a.id, "identifier_id": a.id}, + {"parent_identifier_id": a.id, "identifier_id": b.id}, + ] + _insert_cache_rows(db.session, rows) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + + # Re-inserting the same rows is a no-op. + _insert_cache_rows(db.session, rows) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} + + def test_empty_input(self, db: DatabaseTransactionFixture) -> None: + _insert_cache_rows(db.session, []) + assert db.session.query(RecursiveEquivalencyCache).count() == 0 + + class TestRefreshEquivalentIdentifiers: def test_full_refresh( self,