From 64a7bfea35aca65d1de11318f5c8022ac218caf6 Mon Sep 17 00:00:00 2001 From: Daniel Bernstein Date: Fri, 10 Jul 2026 22:49:42 -0700 Subject: [PATCH 1/5] Fix UniqueViolation in equivalent_identifiers_refresh The refresh task rebuilds recursiveequivalentscache with delete-then-insert per parent identifier, assuming it is the only writer. It is not: the Identifier creation listener inserts (id, id) self-references from webapp transactions, and a redelivered task (acks_late + reject_on_worker_lost, plus self-replacement across thousands of batches) can run a second refresh chain concurrently on overlapping chains. A row committed by either between our delete and our flush collides with the constraint and aborts the whole task. Route all cache inserts through a shared _insert_cache_rows helper using INSERT ... ON CONFLICT DO NOTHING. Since every parent's rows are recomputed from the same source equivalencies, a row slipped in under us holds the same value we were about to write, so skipping the conflict converges on the correct chain. Also collapse the per-parent deletes into a single sorted IN (...) statement to reduce deadlock risk between racing refreshes. Co-Authored-By: Claude Opus 4.8 --- .../manager/sqlalchemy/refresh_equivalents.py | 91 ++++++++++++++----- .../sqlalchemy/test_refresh_equivalents.py | 65 +++++++++++++ 2 files changed, 133 insertions(+), 23 deletions(-) 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, From 7621a3b63dd57d57c69cc75de13b6db30dba5076 Mon Sep 17 00:00:00 2001 From: Jonathan Green Date: Wed, 15 Jul 2026 12:17:46 -0300 Subject: [PATCH 2/5] Insert cache rows in a stable order to avoid deadlocks Sort rows by (parent_identifier_id, identifier_id) in _insert_cache_rows so that two refreshes inserting overlapping keys acquire the unique-index entry locks in the same order rather than deadlocking. Drop the sorted() on the delete's parent IDs, which never controlled lock acquisition order (the single-statement delete already scans in index order), along with its misleading comment. --- .../manager/sqlalchemy/refresh_equivalents.py | 21 ++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/palace/manager/sqlalchemy/refresh_equivalents.py b/src/palace/manager/sqlalchemy/refresh_equivalents.py index 5281bd7408..09fa5441a3 100644 --- a/src/palace/manager/sqlalchemy/refresh_equivalents.py +++ b/src/palace/manager/sqlalchemy/refresh_equivalents.py @@ -41,8 +41,16 @@ def _insert_cache_rows(session: Session, rows: Iterable[dict[str, int]]) -> None 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. + + Rows are inserted in a stable ``(parent_identifier_id, identifier_id)`` + order so that two transactions inserting overlapping keys acquire the + unique-index entry locks in the same order, which keeps them from + deadlocking against each other. """ - for chunk in batched(rows, INSERT_CHUNK_SIZE): + ordered = sorted( + rows, key=lambda row: (row["parent_identifier_id"], row["identifier_id"]) + ) + for chunk in batched(ordered, INSERT_CHUNK_SIZE): session.execute( pg_insert(RecursiveEquivalencyCache) .values(list(chunk)) @@ -103,12 +111,11 @@ def process_identifier_ids(session: Session, identifier_ids: frozenset[int]) -> return # 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}) - ) + # fresh ones. The delete is a single statement, so its row locks are taken + # in the index's scan order — the same for any concurrent refresh — and the + # inserts are ordered by :func:`_insert_cache_rows`, so overlapping refreshes + # acquire locks consistently rather than deadlocking. + _delete_cache_rows(session, {parent_id for _, parent_id in chained_identifiers}) _insert_cache_rows( session, ( From f2ec80038efb1de995023293352e115552941c50 Mon Sep 17 00:00:00 2001 From: Jonathan Green Date: Wed, 15 Jul 2026 12:21:14 -0300 Subject: [PATCH 3/5] Test batching paths in refresh_equivalents Add coverage for the two paths that only trigger above a batch boundary: add_identity_equivalents spanning multiple yield_per partitions (inserting while the streaming cursor is open), and _insert_cache_rows spanning multiple INSERT_CHUNK_SIZE chunks. Both previously only ran with a single partition/chunk in tests. --- .../sqlalchemy/test_refresh_equivalents.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/tests/manager/sqlalchemy/test_refresh_equivalents.py b/tests/manager/sqlalchemy/test_refresh_equivalents.py index 8e8283625e..e6e4471245 100644 --- a/tests/manager/sqlalchemy/test_refresh_equivalents.py +++ b/tests/manager/sqlalchemy/test_refresh_equivalents.py @@ -152,6 +152,27 @@ def test_skips_existing(self, db: DatabaseTransactionFixture) -> None: # No duplicate rows should be added. assert after == before + def test_sweeps_across_multiple_partitions( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + ) -> None: + # With more missing identifiers than one yield_per partition, the sweep + # inserts self-references while the streaming cursor is still open and + # spans several partitions. A batch_size of 2 over 5 identifiers forces + # at least three partitions, exercising the interleaved-insert path. + identifiers = [db.identifier() for _ in range(5)] + db.session.flush() + recursive_equivalency_cache.drop() + + add_identity_equivalents(db.session, batch_size=2) + db.session.flush() + + for identifier in identifiers: + assert recursive_equivalency_cache.cache_for(identifier.id) == { + identifier.id + } + class TestInsertCacheRows: def test_skips_existing_rows( @@ -185,6 +206,35 @@ def test_empty_input(self, db: DatabaseTransactionFixture) -> None: _insert_cache_rows(db.session, []) assert db.session.query(RecursiveEquivalencyCache).count() == 0 + def test_batches_across_multiple_chunks( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + monkeypatch: pytest.MonkeyPatch, + ) -> None: + # More rows than a single INSERT chunk must all be inserted, across + # multiple statements. Shrink the chunk size so a handful of rows spans + # several chunks. + monkeypatch.setattr( + "palace.manager.sqlalchemy.refresh_equivalents.INSERT_CHUNK_SIZE", 2 + ) + parent = db.identifier() + others = [db.identifier() for _ in range(5)] + db.session.flush() + recursive_equivalency_cache.drop() + + rows = [ + {"parent_identifier_id": parent.id, "identifier_id": identifier.id} + for identifier in [parent, *others] + ] + _insert_cache_rows(db.session, rows) + db.session.flush() + + assert recursive_equivalency_cache.cache_for(parent.id) == { + parent.id, + *(identifier.id for identifier in others), + } + class TestRefreshEquivalentIdentifiers: def test_full_refresh( From f7943dd79f5272c56a3ce9b7fd453768b0510e69 Mon Sep 17 00:00:00 2001 From: Jonathan Green Date: Wed, 15 Jul 2026 12:26:09 -0300 Subject: [PATCH 4/5] Cover early-return branches in refresh_equivalents Add tests for the empty-guard returns in _delete_cache_rows and process_identifier_ids, plus the no-equivalencies path in refresh_equivalent_identifiers, bringing the module to 100% coverage. --- .../sqlalchemy/test_refresh_equivalents.py | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/manager/sqlalchemy/test_refresh_equivalents.py b/tests/manager/sqlalchemy/test_refresh_equivalents.py index e6e4471245..eae92d3ead 100644 --- a/tests/manager/sqlalchemy/test_refresh_equivalents.py +++ b/tests/manager/sqlalchemy/test_refresh_equivalents.py @@ -5,6 +5,7 @@ RecursiveEquivalencyCache, ) from palace.manager.sqlalchemy.refresh_equivalents import ( + _delete_cache_rows, _insert_cache_rows, add_identity_equivalents, process_identifier_ids, @@ -121,6 +122,23 @@ def test_tolerates_rows_the_delete_did_not_remove( assert recursive_equivalency_cache.cache_for(a.id) == {a.id, b.id} assert recursive_equivalency_cache.cache_for(b.id) == {a.id, b.id} + def test_no_chained_identifiers_is_noop( + self, + db: DatabaseTransactionFixture, + ) -> None: + # A non-empty set of identifier IDs that don't exist passes the empty + # check but yields no chains, so the function returns without touching + # the cache. + db.identifier() + db.session.flush() + before = db.session.query(RecursiveEquivalencyCache).count() + + process_identifier_ids(db.session, frozenset([-1, -2])) + db.session.flush() + + # Nothing was deleted or inserted. + assert db.session.query(RecursiveEquivalencyCache).count() == before + class TestAddIdentityEquivalents: def test_adds_self_references( @@ -236,6 +254,19 @@ def test_batches_across_multiple_chunks( } +class TestDeleteCacheRows: + def test_empty_input(self, db: DatabaseTransactionFixture) -> None: + # An empty parent set is a no-op: no DELETE is issued, so existing rows + # (e.g. the creation listener's self-references) are left in place. + db.identifier() + db.session.flush() + before = db.session.query(RecursiveEquivalencyCache).count() + + _delete_cache_rows(db.session, set()) + + assert db.session.query(RecursiveEquivalencyCache).count() == before + + class TestRefreshEquivalentIdentifiers: def test_full_refresh( self, @@ -256,3 +287,18 @@ def test_full_refresh( assert recursive_equivalency_cache.cache_for(b.id) == {a.id, b.id} # c should have a self-reference. assert recursive_equivalency_cache.cache_for(c.id) == {c.id} + + def test_no_equivalencies( + self, + db: DatabaseTransactionFixture, + recursive_equivalency_cache: RecursiveEquivalencyCacheFixture, + ) -> None: + # With no equivalencies at all, the chain computation is skipped and + # only the self-reference sweep runs. + a = db.identifier() + db.session.flush() + recursive_equivalency_cache.drop() + + refresh_equivalent_identifiers(db.session) + + assert recursive_equivalency_cache.cache_for(a.id) == {a.id} From cf62739562dd7b94ce8df3bb725724784678df38 Mon Sep 17 00:00:00 2001 From: Jonathan Green Date: Wed, 15 Jul 2026 12:41:39 -0300 Subject: [PATCH 5/5] Update comments [skip claude] --- .../manager/sqlalchemy/refresh_equivalents.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/src/palace/manager/sqlalchemy/refresh_equivalents.py b/src/palace/manager/sqlalchemy/refresh_equivalents.py index 09fa5441a3..ea82c99c86 100644 --- a/src/palace/manager/sqlalchemy/refresh_equivalents.py +++ b/src/palace/manager/sqlalchemy/refresh_equivalents.py @@ -34,18 +34,10 @@ 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. Rows are inserted in a stable ``(parent_identifier_id, identifier_id)`` - order so that two transactions inserting overlapping keys acquire the - unique-index entry locks in the same order, which keeps them from - deadlocking against each other. + order so that two transactions inserting overlapping keys don't + deadlock against each other. """ ordered = sorted( rows, key=lambda row: (row["parent_identifier_id"], row["identifier_id"]) @@ -111,10 +103,7 @@ def process_identifier_ids(session: Session, identifier_ids: frozenset[int]) -> return # Delete the old cache entries for every affected parent, then insert the - # fresh ones. The delete is a single statement, so its row locks are taken - # in the index's scan order — the same for any concurrent refresh — and the - # inserts are ordered by :func:`_insert_cache_rows`, so overlapping refreshes - # acquire locks consistently rather than deadlocking. + # fresh ones. _delete_cache_rows(session, {parent_id for _, parent_id in chained_identifiers}) _insert_cache_rows( session,