From 5eb09c1e36e97c82ad21a42691b038bed5709db4 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Fri, 16 Oct 2026 23:35:12 +0300 Subject: [PATCH] perf: send a message with one INSERT ... ON CONFLICT DO UPDATE The pre-check SELECT, the insert and the DuplicateKeyError recovery re-read become one statement. RETURNING xmax = 0 tells a new row from a replay, so last_message_id moves only when the message was created. --- AGENTS.md | 8 +- app/repositories/messages_repository.py | 19 +- app/use_cases/create_message.py | 37 +--- .../0005-upsert-via-duplicate-key-recovery.md | 13 +- docs/adr/0006-idempotency-scoped-per-chat.md | 9 +- ...009-coverage-exclusions-structural-only.md | 4 +- docs/adr/0012-message-send-is-one-upsert.md | 19 ++ tests/use_cases/test_create_message.py | 200 +++++++++--------- 8 files changed, 164 insertions(+), 145 deletions(-) create mode 100644 docs/adr/0012-message-send-is-one-upsert.md diff --git a/AGENTS.md b/AGENTS.md index 8cfe5c0..0f41521 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -135,10 +135,12 @@ back. `app`/`client` fixtures build the real app and run it through container_fixture="request_container")` (`tests/use_cases/conftest.py`) exposes every repository/use case provider as a same-named pytest fixture — the template predates this and hand-assembles dependencies instead. The -race-simulation pattern used to test the concurrent-retry paths without a second +race-simulation pattern used to test the concurrent-retry path without a second real connection is the `_Racing*Repository` classes in -`tests/use_cases/test_create_chat.py` and `tests/use_cases/test_create_message.py`; -the invariant each one pins is in the docstring on the test that uses it. +`tests/use_cases/test_create_chat.py`. `tests/use_cases/test_create_message.py` +uses a second real connection instead: it holds an uncommitted insert and polls +`pg_stat_activity` until the use case waits on its lock. The invariant each one +pins is in the docstring on the test that uses it. **Migrations**: `migrations/env.py` reads the shared `METADATA` and rewrites the DSN driver from `postgresql+asyncpg` → `postgresql` (Alembic uses sync diff --git a/app/repositories/messages_repository.py b/app/repositories/messages_repository.py index e5e82dd..1ff8caf 100644 --- a/app/repositories/messages_repository.py +++ b/app/repositories/messages_repository.py @@ -6,6 +6,7 @@ from advanced_alchemy.filters import LimitOffset from advanced_alchemy.repository import SQLAlchemyAsyncRepository from advanced_alchemy.service import SQLAlchemyAsyncRepositoryService +from sqlalchemy.dialects import postgresql from app.database import tables @@ -16,8 +17,22 @@ class BaseRepository(SQLAlchemyAsyncRepository[tables.MessagesTable]): repository_type = BaseRepository - async def fetch_by_idempotency_key(self, chat_id: int, idempotency_key: uuid.UUID) -> tables.MessagesTable | None: - return await self.get_one_or_none(chat_id=chat_id, idempotency_key=idempotency_key) + async def get_or_create( + self, *, chat_id: int, user_id: int, idempotency_key: uuid.UUID, text: str + ) -> tuple[tables.MessagesTable, bool]: + statement = postgresql.insert(tables.MessagesTable).values( + chat_id=chat_id, user_id=user_id, idempotency_key=idempotency_key, text=text + ) + statement = statement.on_conflict_do_update( + constraint="uk_messages_chat_id_idempotency_key", + set_={"idempotency_key": statement.excluded.idempotency_key}, + ) + # xmax is 0 only on a row version this statement inserted. + created: typing.Final = sa.literal_column("xmax = 0", sa.Boolean) + message, was_created = ( + await self.repository.session.execute(statement.returning(tables.MessagesTable, created)) + ).one() + return message, was_created async def fetch_latest_active(self, chat_id: int) -> tables.MessagesTable | None: """Return the newest non-deleted message in a chat, or None if none remains. diff --git a/app/use_cases/create_message.py b/app/use_cases/create_message.py index babb16f..30eeed3 100644 --- a/app/use_cases/create_message.py +++ b/app/use_cases/create_message.py @@ -1,7 +1,5 @@ import dataclasses -import typing -from advanced_alchemy.exceptions import DuplicateKeyError from db_retry import Transaction, postgres_retry from app.actor import Actor @@ -24,44 +22,19 @@ class CreateMessageUseCase: async def __call__( self, *, actor: Actor, chat_id: int, data: SendMessageRequest ) -> tuple[tables.MessagesTable, bool]: - message: tables.MessagesTable | None = None async with self.transaction: if not await self.chat_members_repository.is_member(chat_id, actor.id): msg = "Not a member of this chat" raise PermissionDeniedError(msg) - existing: typing.Final = await self.messages_repository.fetch_by_idempotency_key( - chat_id, data.idempotency_key + message, created = await self.messages_repository.get_or_create( + chat_id=chat_id, user_id=actor.id, idempotency_key=data.idempotency_key, text=data.text ) - if existing is not None: - await self.transaction.commit() - return existing, False - - try: - message = await self.messages_repository.create( - tables.MessagesTable( - chat_id=chat_id, - user_id=actor.id, - idempotency_key=data.idempotency_key, - text=data.text, - ) - ) - except DuplicateKeyError: - # Re-read outside the block: __aexit__ rolls back and detaches on an open transaction. - await self.transaction.rollback() - else: + if created: await self.chats_repository.update( tables.ChatsTable(id=chat_id, last_message_id=message.id), item_id=chat_id, attribute_names=["last_message_id"], ) - await self.transaction.commit() - - if message is None: - duplicate = await self.messages_repository.fetch_by_idempotency_key(chat_id, data.idempotency_key) - if duplicate is None: - msg = "Message send raced but the resulting row could not be found" - raise RuntimeError(msg) - return duplicate, False - - return message, True + await self.transaction.commit() + return message, created diff --git a/docs/adr/0005-upsert-via-duplicate-key-recovery.md b/docs/adr/0005-upsert-via-duplicate-key-recovery.md index bae8769..0daf4bb 100644 --- a/docs/adr/0005-upsert-via-duplicate-key-recovery.md +++ b/docs/adr/0005-upsert-via-duplicate-key-recovery.md @@ -1,7 +1,7 @@ # Upsert by recovering from DuplicateKeyError -`CreateChatUseCase` and `CreateMessageUseCase` both read first to see whether -the row already exists and then insert, but the read is only an optimisation: +`CreateChatUseCase` reads first to see whether the direct chat already exists +and then inserts, but the read is only an optimisation: the correctness guarantee is the `except DuplicateKeyError` branch, which rolls back and re-reads the row the winner committed. The pre-check alone does not make the insert safe. At READ COMMITTED two concurrent "open a DM with Bob" @@ -10,7 +10,8 @@ requests both miss the read, both insert, and the loser violates have received the existing chat. `@postgres_retry` does not rescue it either, because `db-retry` retries serialization and connection failures, not integrity violations, and `SELECT ... FOR UPDATE` has nothing to take, since the race is -between two inserts of a row that does not yet exist. Both recovery paths are -exercised by repository subclasses (`_RacingChatsRepository`, -`_RacingMessagesRepository`) whose `create` raises and whose lookup misses once, -simulating the database condition at a seam the tests already own. +between two inserts of a row that does not yet exist. The recovery path is +exercised by a repository subclass, `_RacingChatsRepository`, whose `create` +raises and whose lookup misses once, simulating the database condition at a seam +the tests already own. Message sends used this pattern too until +[0012](0012-message-send-is-one-upsert.md) replaced it with a single statement. diff --git a/docs/adr/0006-idempotency-scoped-per-chat.md b/docs/adr/0006-idempotency-scoped-per-chat.md index d61797b..1a02bec 100644 --- a/docs/adr/0006-idempotency-scoped-per-chat.md +++ b/docs/adr/0006-idempotency-scoped-per-chat.md @@ -1,8 +1,7 @@ # Idempotency is scoped per chat `messages` carries `UniqueConstraint("chat_id", "idempotency_key")`, and the -pre-check lookup, the constraint and the `DuplicateKeyError` recovery re-read -all filter on that same pair. Idempotency is a property of an operation, and the +send's `ON CONFLICT` names that constraint. Idempotency is a property of an operation, and the operation is "send this message to this chat", so two chats are two operations and a key reused across them is not a retry of anything. A global constraint on `idempotency_key` shipped first and made a reachable state look unreachable: @@ -10,6 +9,6 @@ cross-chat key reuse under concurrency raised `DuplicateKeyError` from the insert, the chat-scoped re-read missed, and control reached a guard whose only justification was "the unique constraint guarantees a match here". Worse, a client reusing a key across chats received the other chat's message while its -intended message was never written, a silent wrong-row return. The three -surfaces must agree, and aligning the constraint with the lookup was both the -cheaper direction and the one that matches the domain. +intended message was never written, a silent wrong-row return. The surfaces +had to agree, and aligning the constraint with the lookup was both the cheaper +direction and the one that matches the domain. diff --git a/docs/adr/0009-coverage-exclusions-structural-only.md b/docs/adr/0009-coverage-exclusions-structural-only.md index ce43236..e881870 100644 --- a/docs/adr/0009-coverage-exclusions-structural-only.md +++ b/docs/adr/0009-coverage-exclusions-structural-only.md @@ -7,8 +7,8 @@ explain itself, plus `exclude_also` for `if typing.TYPE_CHECKING:` blocks. No `# pragma: no cover` appears in `app/`, `tests/` or `migrations/`. "Awkward to reach" is not a warrant: where a branch looked untestable the answer was a repository subclass that raises the condition the database would raise, which is -how both `DuplicateKeyError` recovery paths and both defensive `is None` guards -became executed code. Pragmas on unreachable-in-production guards were argued +how the `DuplicateKeyError` recovery path in `CreateChatUseCase` and its +defensive `is None` guard became executed code. Pragmas on unreachable-in-production guards were argued for twice and refused, because an excluded branch is one nobody notices when it stops being unreachable, and the coverage number then asserts something untrue about what the tests exercise; tests written only to move the number were diff --git a/docs/adr/0012-message-send-is-one-upsert.md b/docs/adr/0012-message-send-is-one-upsert.md new file mode 100644 index 0000000..7a761b9 --- /dev/null +++ b/docs/adr/0012-message-send-is-one-upsert.md @@ -0,0 +1,19 @@ +# A message send is one INSERT ... ON CONFLICT DO UPDATE + +`MessagesRepository.get_or_create` writes a send as one `INSERT ... ON CONFLICT +ON CONSTRAINT uk_messages_chat_id_idempotency_key DO UPDATE SET idempotency_key += EXCLUDED.idempotency_key RETURNING ..., xmax = 0`. The no-op update makes a +replay return the existing row from the same statement, and a request racing an +uncommitted insert of the same key waits on that row's lock, then gets the +winner's row. `xmax = 0` is true only for a row version the statement inserted, +so the use case can answer 201 or 200 and skip the `last_message_id` update on a +replay. The read-then-insert with `DuplicateKeyError` recovery from +[0005](0005-upsert-via-duplicate-key-recovery.md) cost two round trips per new +message and re-read the winner outside the retried transaction. `ON CONFLICT DO +NOTHING` returns no row to the loser, and advanced-alchemy's `upsert` reads +first. The cost: every replay writes a row version and consumes a sequence +value, so a retried send leaves a gap in message ids that +[0001](0001-sequence-ids-not-snowflakes.md)'s client gap detection answers with +one empty catch-up fetch. Postgres does not document `xmax` in `RETURNING`; the +replay and concurrent tests in `tests/use_cases/test_create_message.py` would +catch a change. diff --git a/tests/use_cases/test_create_message.py b/tests/use_cases/test_create_message.py index d52bb85..822040c 100644 --- a/tests/use_cases/test_create_message.py +++ b/tests/use_cases/test_create_message.py @@ -1,59 +1,26 @@ +import asyncio +import re +import typing import uuid import pytest -from advanced_alchemy.exceptions import DuplicateKeyError +import sqlalchemy as sa +from db_retry import Transaction +from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession +from app import security from app.actor import Actor from app.database import tables +from app.database.resources import create_database_engine, create_session from app.exceptions import PermissionDeniedError +from app.repositories.chat_members_repository import ChatMembersRepository from app.repositories.chats_repository import ChatsRepository from app.repositories.messages_repository import MessagesRepository from app.schemas import api as schemas from app.use_cases.create_chat import CreateChatUseCase from app.use_cases.create_message import CreateMessageUseCase - - -class _RacingMessagesRepository(MessagesRepository): - """Simulate losing a concurrent-send-with-the-same-key race. - - The pre-check misses (as if the winner's row weren't committed/visible yet), the insert - then collides with the winner's now-committed row (DuplicateKeyError), and the recovery - re-read must find it. - """ - - _missed_precheck: bool = False - - async def fetch_by_idempotency_key(self, chat_id: int, idempotency_key: uuid.UUID) -> tables.MessagesTable | None: - if not self._missed_precheck: - self._missed_precheck = True - return None - return await super().fetch_by_idempotency_key(chat_id, idempotency_key) - - async def create(self, *_args: object, **_kwargs: object) -> tables.MessagesTable: - msg = "simulated race: another request already sent this message" - raise DuplicateKeyError(msg) - - -class _NeverFoundMessagesRepository(MessagesRepository): - """Simulates a race whose recovery re-read can never find the winner's row. - - Both `fetch_by_idempotency_key` calls (the pre-check and the post-rollback recovery - re-read) return `None`, and `create()` always raises `DuplicateKeyError` - a state the - unique constraint on `(chat_id, idempotency_key)` should make unreachable in production, - exercised here only to prove `CreateMessageUseCase` raises `RuntimeError` rather than - returning `None` silently. - """ - - async def fetch_by_idempotency_key( - self, - chat_id: int, # noqa: ARG002 - idempotency_key: uuid.UUID, # noqa: ARG002 - ) -> tables.MessagesTable | None: - return None - - async def create(self, *_args: object, **_kwargs: object) -> tables.MessagesTable: - msg = "simulated race: another request already sent this message" - raise DuplicateKeyError(msg) +from tests.api.helpers import count_statements +from tests.factories import UserFactory async def test_send_returns_created_true_on_first_call( @@ -106,57 +73,6 @@ async def test_non_member_cannot_send( ) -async def test_concurrent_duplicate_key_recovers_the_winners_message( - create_message_use_case: CreateMessageUseCase, direct_chat: tables.ChatsTable, alice: Actor -) -> None: - """INVARIANT: a message insert losing the (chat_id, idempotency_key) race returns the winner's row. - - Broken by dropping the explicit `await self.transaction.rollback()` before the - recovery read, or by widening the caught exception beyond DuplicateKeyError so - an unrelated integrity error is silently funnelled into recovery. - """ - key = uuid.uuid4() - winner, winner_created = await create_message_use_case( - actor=alice, chat_id=direct_chat.id, data=schemas.SendMessageRequest(idempotency_key=key, text="hi") - ) - assert winner_created is True - winner_id = winner.id - - racer = CreateMessageUseCase( - transaction=create_message_use_case.transaction, - chats_repository=create_message_use_case.chats_repository, - chat_members_repository=create_message_use_case.chat_members_repository, - messages_repository=_RacingMessagesRepository( - session=create_message_use_case.messages_repository.repository.session, auto_commit=False - ), - ) - loser, loser_created = await racer( - actor=alice, chat_id=direct_chat.id, data=schemas.SendMessageRequest(idempotency_key=key, text="hi again") - ) - assert loser_created is False - assert loser.id == winner_id - assert loser.text == "hi" - - -async def test_send_recovery_raises_if_the_winners_row_is_unreadable( - create_message_use_case: CreateMessageUseCase, direct_chat: tables.ChatsTable, alice: Actor -) -> None: - broken = CreateMessageUseCase( - transaction=create_message_use_case.transaction, - chats_repository=create_message_use_case.chats_repository, - chat_members_repository=create_message_use_case.chat_members_repository, - messages_repository=_NeverFoundMessagesRepository( - session=create_message_use_case.messages_repository.repository.session, auto_commit=False - ), - ) - with pytest.raises(RuntimeError, match="could not be found"): - await broken( - actor=alice, - chat_id=direct_chat.id, - data=schemas.SendMessageRequest(idempotency_key=uuid.uuid4(), text="hi"), - ) - - async def test_same_idempotency_key_in_two_different_chats_creates_two_messages( create_message_use_case: CreateMessageUseCase, create_chat_use_case: CreateChatUseCase, @@ -181,3 +97,97 @@ async def test_same_idempotency_key_in_two_different_chats_creates_two_messages( assert first.id != second.id assert first.chat_id == direct_chat.id assert second.chat_id == other_chat.id + + +async def test_send_writes_the_message_with_a_single_statement( + create_message_use_case: CreateMessageUseCase, + direct_chat: tables.ChatsTable, + alice: Actor, + db_session: AsyncSession, +) -> None: + with count_statements(db_session) as statements: + await create_message_use_case( + actor=alice, + chat_id=direct_chat.id, + data=schemas.SendMessageRequest(idempotency_key=uuid.uuid4(), text="hi"), + ) + + message_statements: typing.Final = [statement for statement in statements if re.search(r"\bmessages\b", statement)] + assert len(message_statements) == 1, message_statements + assert "ON CONFLICT" in message_statements[0] + + +async def _wait_until_blocked_on_a_lock(engine: AsyncEngine, pid: int) -> None: + wait_event_type = None + while wait_event_type != "Lock": + await asyncio.sleep(0.01) + async with engine.connect() as connection: + wait_event_type = await connection.scalar( + sa.text("SELECT wait_event_type FROM pg_stat_activity WHERE pid = :pid"), {"pid": pid} + ) + + +async def test_send_returns_the_message_a_concurrent_request_created_with_the_same_key() -> None: + """INVARIANT: a send racing an uncommitted insert of the same key waits for it and returns that row. + + It reports created=False and leaves last_message_id alone. Broken by switching to ON CONFLICT + DO NOTHING, whose RETURNING yields no row for the request that loses the race. + """ + engine: typing.Final = create_database_engine() + idempotency_key: typing.Final = uuid.uuid4() + setup_session: typing.Final = create_session(engine) + user: typing.Final = UserFactory.build( + username=f"racer-{idempotency_key}", + password_hash=security.hash_password("hunter2hunter2"), + display_name="Racer", + ) + setup_session.add(user) + await setup_session.flush() + chat: typing.Final = tables.ChatsTable(chat_type=tables.ChatType.GROUP, title="race", created_by_id=user.id) + setup_session.add(chat) + await setup_session.flush() + setup_session.add(tables.ChatMembersTable(chat_id=chat.id, user_id=user.id)) + await setup_session.commit() + await setup_session.close() + try: + async with engine.connect() as concurrent_request: + await concurrent_request.execute( + sa.insert(tables.MessagesTable).values( + chat_id=chat.id, user_id=user.id, idempotency_key=idempotency_key, text="first" + ) + ) + session: typing.Final = create_session(engine) + pid: typing.Final = await (await session.connection()).scalar(sa.text("SELECT pg_backend_pid()")) + use_case: typing.Final = CreateMessageUseCase( + transaction=Transaction(session=session), + chats_repository=ChatsRepository(session=session, auto_commit=False), + chat_members_repository=ChatMembersRepository(session=session, auto_commit=False), + messages_repository=MessagesRepository(session=session, auto_commit=False), + ) + pending: typing.Final = asyncio.create_task( + use_case( + actor=Actor(id=user.id), + chat_id=chat.id, + data=schemas.SendMessageRequest(idempotency_key=idempotency_key, text="second"), + ) + ) + await asyncio.wait_for(_wait_until_blocked_on_a_lock(engine, pid), timeout=5) + await concurrent_request.commit() + message, created = await pending + await session.close() + async with engine.connect() as connection: + last_message_id: typing.Final = await connection.scalar( + sa.select(tables.ChatsTable.last_message_id).where(tables.ChatsTable.id == chat.id) + ) + assert message.text == "first" + assert created is False + assert last_message_id is None + finally: + async with engine.begin() as connection: + await connection.execute(sa.delete(tables.MessagesTable).where(tables.MessagesTable.chat_id == chat.id)) + await connection.execute( + sa.delete(tables.ChatMembersTable).where(tables.ChatMembersTable.chat_id == chat.id) + ) + await connection.execute(sa.delete(tables.ChatsTable).where(tables.ChatsTable.id == chat.id)) + await connection.execute(sa.delete(tables.UsersTable).where(tables.UsersTable.id == user.id)) + await engine.dispose()