From 7d05f82f819c08a733aacd7a3c303b102ba56a54 Mon Sep 17 00:00:00 2001 From: Artur Shiriev Date: Mon, 26 Oct 2026 17:51:24 +0300 Subject: [PATCH] feat: queue a rearm when extending a timer whose handler is running --- AGENTS.md | 5 +- CONTEXT.md | 7 +- docs/operations/alembic.md | 43 ++++++ docs/usage/timers.md | 82 ++++++++---- faststream_outbox/broker.py | 24 ++-- faststream_outbox/client.py | 174 ++++++++++++++++++++----- faststream_outbox/schema.py | 11 ++ faststream_outbox/schema_validation.py | 21 ++- faststream_outbox/testing.py | 45 +++++-- tests/test_client_contract.py | 166 +++++++++++++++++++++-- tests/test_fake.py | 17 +++ tests/test_integration.py | 132 ++++++++++++++++++- tests/test_unit.py | 74 +++++++---- 13 files changed, 673 insertions(+), 128 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 93ec82e..7f7e074 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -34,7 +34,10 @@ will **not** tell you: runs in the database and one in the process. They cannot share an implementation, so `tests/test_client_contract.py` couples them by behaviour. A change to either adapter's insert, fetch, terminal, or DLQ semantics belongs in that suite. -- `schema.py`: the three partial indexes and the `_lease_ck` CHECK are load-bearing. +- `schema.py`: the three partial indexes and the `
_lease_ck` and `
_rearm_ck` CHECKs are + load-bearing. +- Every terminal write (delete, batched delete, DLQ, retry) must apply a queued rearm instead of + dropping it; `tests/test_client_contract.py` holds the rearm scenarios for both adapters. - `message.py`: the `DLQFailureReason` `Literal` is a **public contract**; operator queries key off those strings, so changing one is API-breaking. - `metrics/` (the recorder seam) and `prometheus/` + `opentelemetry/` (native middleware) are two diff --git a/CONTEXT.md b/CONTEXT.md index de50838..60730a0 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -69,9 +69,14 @@ _Avoid_: idempotency key, message key **Extend**: Re-publishing a `timer_id` with `timer_conflict="extend"`: an unleased timer's deadline moves to -the later of the stored and requested one, in place. A leased timer is left to its worker. +the later of the stored and requested one, in place. On a leased timer it queues a rearm. _Avoid_: refresh, reschedule, upsert +**Rearm**: +An extend queued on a leased timer, held in the `rearm_*` columns. The run's terminal write applies +it, turning the row back into a pending timer instead of deleting it. +_Avoid_: successor, follow-up timer + **Relay**: Forwarding a delivered row onward through a native FastStream publisher chain to a foreign broker. _Avoid_: bridge, forwarder diff --git a/docs/operations/alembic.md b/docs/operations/alembic.md index e814326..987fce6 100644 --- a/docs/operations/alembic.md +++ b/docs/operations/alembic.md @@ -200,6 +200,49 @@ table in a non-default schema, drop an `INVALID` index left by a failed concurrent build before retrying, and don't apply what `alembic revision --autogenerate` produces for this change. +## Upgrading to 0.19: add the rearm columns { #upgrading-to-019 } + +0.19 lets `publish(..., timer_conflict="extend")` reach a timer whose +handler is running: the extend is queued in three new nullable columns, +`rearm_at`, `rearm_payload` and `rearm_headers`, and applied when the run +ends. A CHECK, `outbox_rearm_ck`, keeps a queued rearm off unleased rows. +`validate_schema()` reports a table without them as drift. + +Run the migration **before** deploying 0.19, which reads and writes the new +columns on every completion. Adding a nullable column without a default +doesn't rewrite the table, and adding the CHECK as `NOT VALID` and +validating it separately avoids holding an exclusive lock while the table is +scanned. + +```python +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + + +def upgrade() -> None: + op.add_column("outbox", sa.Column("rearm_at", sa.DateTime(timezone=True), nullable=True)) + op.add_column("outbox", sa.Column("rearm_payload", sa.LargeBinary(), nullable=True)) + op.add_column("outbox", sa.Column("rearm_headers", postgresql.JSONB(), nullable=True)) + op.execute( + "ALTER TABLE outbox ADD CONSTRAINT outbox_rearm_ck " + "CHECK (rearm_at IS NULL OR acquired_token IS NOT NULL) NOT VALID" + ) + op.execute("ALTER TABLE outbox VALIDATE CONSTRAINT outbox_rearm_ck") + + +def downgrade() -> None: + op.drop_constraint("outbox_rearm_ck", "outbox", type_="check") + op.drop_column("outbox", "rearm_headers") + op.drop_column("outbox", "rearm_payload") + op.drop_column("outbox", "rearm_at") +``` + +Earlier versions ignore the new columns, and `validate_schema()` doesn't +report extra columns, so deploying after the migration is safe. Don't +downgrade while a rearm is queued: dropping the columns drops it, and the +timer is deleted when its run ends. + ## Adding the DLQ after the fact To opt into [Dead-letter queue](../usage/dlq.md) audit, add diff --git a/docs/usage/timers.md b/docs/usage/timers.md index 56ec180..4f2d3d5 100644 --- a/docs/usage/timers.md +++ b/docs/usage/timers.md @@ -129,23 +129,46 @@ of order cannot pull a timer in. When the requested deadline is later, the row also takes the new payload and headers, and its retry history (`attempts_count`, `first_attempt_at`, `last_attempt_at`) is cleared. When it is not later, the row keeps everything it had and the call still returns -its id. `deliveries_count` is never reset. - -An extend never touches a row a worker has leased, and returns `None` for -it. - -- An id means that once your transaction commits, the timer will not fire - before the deadline you asked for. Until the commit, workers skip the - locked row; after it, they see the new deadline. -- `None` means a worker claimed the timer first. Its handler is running, or - will run when an expired lease is reclaimed, for the deadline it claimed. - Publishing again does not help while that row exists. - -If a late heartbeat must win over an expiry that is already running, the -library cannot decide that for you. Have the handler re-check, in the same -transaction as its side effects, whether the timeout still applies (for -example against a last-activity column the heartbeat writes), and publish a -fresh timer from the heartbeat once the old row is gone. +its id. `deliveries_count` is not reset. + +Once your transaction commits, the timer will not fire before the deadline +you asked for. Until the commit, workers skip the locked row; after it, they +see the new deadline. + +### Extending a timer that is already running + +When a worker has claimed the timer, its handler runs for the deadline it +claimed, and the extend leaves that run alone. It queues a rearm on the row +instead and still returns the row id. When the run ends, the rearm turns the +row back into a pending timer with the queued deadline, payload and headers, +so every accepted heartbeat leads to a firing at or after its deadline: + +- If the handler succeeds, the timer starts over at the rearm deadline, with + its retry history and `deliveries_count` cleared. +- If the handler fails and would be retried, the rearm replaces the retry. + `deliveries_count` is kept, so `max_deliveries` still stops a handler that + fails on every run. +- If the handler fails terminally, the failed run goes to the DLQ (when one + is configured) and the timer starts over at the rearm deadline. +- If the worker dies, the expired lease is reclaimed and the run is + delivered again as usual; the rearm applies when that run ends. + +Queued rearms follow the same rule as pending timers: a rearm only moves +later, and an extend whose deadline is not later than the running deadline +queues nothing. + +A handler can rearm its own timer the same way, by publishing its own +`timer_id` with `timer_conflict="extend"`, directly or through an +`OutboxResponse`. + +The rearm doesn't undo what the running handler does. If a heartbeat must +win over an expiry that is already running, have the handler re-check, in +the same transaction as its side effects, whether the timeout still applies +(for example against a last-activity column the heartbeat writes). The +rearm then fires the handler again for the newer deadline. + +A rearm that becomes due when the run ends is picked up by polling, within +the subscriber's `max_fetch_interval`. No NOTIFY is sent for it. Every extend writes a new row version, even when the deadline does not move. `next_attempt_at` is indexed, so these updates are not HOT updates: @@ -155,25 +178,24 @@ Size autovacuum for your heartbeat rate (see ## Cancellation -`broker.cancel_timer(*, queue, timer_id, session)` issues a `DELETE` on -the caller's session, but only if the row is not yet leased: +`broker.cancel_timer(*, queue, timer_id, session)` cancels a timer on the +caller's session: ```python -deleted = await broker.cancel_timer( +cancelled = await broker.cancel_timer( queue="orders", timer_id="order-confirm-42", session=session, ) -# True if a row was deleted; False if it didn't exist or was already in flight. +# True if a pending timer or a queued rearm was dropped. ``` -The underlying SQL is `DELETE WHERE queue=? AND timer_id=? AND -acquired_token IS NULL`. The `acquired_token IS NULL` guard is -load-bearing: it preserves the lease-token invariant by refusing to -clobber a row whose handler is already running. If the timer fired in the -race window between your application logic deciding to cancel and the -`DELETE` landing, the delivery completes normally and `cancel_timer` -returns `False`. +It locks the `(queue, timer_id)` row first. An unleased row is deleted. A +leased row is never deleted, because its handler is already running: the +cancel only drops a queued rearm, so the timer does not come back after the +run, and returns `False` when there was none. If the timer fired in the race +window between your application logic deciding to cancel and the cancel +landing, the delivery completes normally. ## Latency floor @@ -198,3 +220,7 @@ on it, bind the `TestOutboxBroker` to a name and read `tb.fake_client.rows[0].next_attempt_at`. Pass `run_loops=True` if you need scheduled delivery to actually wait. See [Testing](./testing.md). + +A rearm applied when a run ends is not dispatched again in sync mode, so a +handler that rearms its own timer runs once per publish. The rearmed row +stays in `tb.fake_client.rows`. diff --git a/faststream_outbox/broker.py b/faststream_outbox/broker.py index ea0e14b..775d212 100644 --- a/faststream_outbox/broker.py +++ b/faststream_outbox/broker.py @@ -439,9 +439,10 @@ async def publish( # ty: ignore[invalid-method-override] returns ``None``. ``"extend"`` updates a row no worker has leased, in place: its deadline becomes the later of the stored and the requested one, and only when the requested one is later does the row take the new payload and - headers and drop its retry history. Either way it returns the row id. It - returns ``None`` only when a worker holds the lease: that worker's handler - runs for the deadline it claimed, and the caller has to handle that case. + headers and drop its retry history. On a leased row it queues the same + update as a rearm, which applies when the running handler finishes; see + the timers guide for how it interacts with retries and the DLQ. An extend + always returns the row id. Returns the row's ``id`` (BigInt PK), or ``None`` as described above. """ @@ -510,19 +511,18 @@ async def cancel_timer( timer_id: str, session: SessionOrConnection, ) -> bool: - """Delete a not-yet-leased timer row. Idempotent at the ``(queue, timer_id)`` level. + """Delete a not-yet-leased timer row, or drop the rearm queued on a leased one. Same transactional contract as :meth:`publish` — runs on the caller's session and - commits with their transaction. + commits with their transaction. Idempotent at the ``(queue, timer_id)`` level. - Returns ``True`` if *this call* deleted the row. ``False`` means the row is no - longer cancelable by you, which covers three cases: + Returns ``True`` if *this call* deleted the row or dropped its queued rearm. + ``False`` means there was nothing left for you to cancel, which covers three cases: - 1. A worker has claimed the row (``acquired_token IS NOT NULL``). The - ``acquired_token IS NULL`` guard prevents this call from clobbering the - lease. Normally the delivery then completes. If that worker died, the lease - stays set until ``lease_ttl_seconds`` passes and the row is delivered again; - ``cancel_timer`` can't stop that redelivery. + 1. A worker has claimed the row and no rearm is queued on it. The call never + deletes a leased row, so the delivery completes normally. If that worker + died, the lease stays set until ``lease_ttl_seconds`` passes and the row is + delivered again; ``cancel_timer`` can't stop that redelivery. 2. Another caller already canceled the same ``(queue, timer_id)`` in a concurrent transaction. 3. No row exists with that ``(queue, timer_id)`` — either the timer was never diff --git a/faststream_outbox/client.py b/faststream_outbox/client.py index 5b57df6..cf7a17c 100644 --- a/faststream_outbox/client.py +++ b/faststream_outbox/client.py @@ -27,6 +27,7 @@ Float, Integer, String, + and_, bindparam, case, cast, @@ -240,11 +241,33 @@ def branch(name: str, *lease_state: "ColumnElement[bool]") -> "CTE": ) +def _rearm_values(outbox: "Table") -> "dict[str, typing.Any]": + """``SET`` values that turn a finished run with a queued rearm into a fresh pending timer.""" + return { + "next_attempt_at": outbox.c.rearm_at, + "payload": outbox.c.rearm_payload, + "headers": outbox.c.rearm_headers, + "rearm_at": None, + "rearm_payload": None, + "rearm_headers": None, + "acquired_at": None, + "acquired_token": None, + "attempts_count": 0, + "deliveries_count": 0, + "first_attempt_at": None, + "last_attempt_at": None, + } + + def _build_dlq_cte(outbox: "Table", dlq: "Table") -> "Insert": """Build the DELETE-then-archive CTE once; the row id, lease token and failure context are bind parameters.""" deleted = ( delete(outbox) - .where(outbox.c.id == bindparam("message_id"), outbox.c.acquired_token == bindparam("acquired_token")) + .where( + outbox.c.id == bindparam("message_id"), + outbox.c.acquired_token == bindparam("acquired_token"), + outbox.c.rearm_at.is_(None), + ) .returning(*(outbox.c[col] for col, _ in _DLQ_PROJECTION)) .cte("deleted") ) @@ -255,33 +278,67 @@ def _build_dlq_cte(outbox: "Table", dlq: "Table") -> "Insert": return insert(dlq).from_select([*(col for _, col in _DLQ_PROJECTION), *_DLQ_INJECTED_COLUMNS], archived) +def _build_dlq_rearm_cte(outbox: "Table", dlq: "Table") -> "Insert": + """Like :func:`_build_dlq_cte` for a run with a queued rearm: archive the run, then apply the rearm.""" + run = ( + select(*(outbox.c[col] for col, _ in _DLQ_PROJECTION)) + .where( + outbox.c.id == bindparam("message_id"), + outbox.c.acquired_token == bindparam("acquired_token"), + outbox.c.rearm_at.is_not(None), + ) + .cte("run") + ) + rearmed = ( + update(outbox) + .where(outbox.c.id == run.c.id, outbox.c.acquired_token == bindparam("acquired_token")) + .values(**_rearm_values(outbox)) + .returning(outbox.c.id) + .cte("rearmed") + ) + archived = select( + *(run.c[col] for col, _ in _DLQ_PROJECTION), + *(bindparam(col, type_=dlq.c[col].type) for col in _DLQ_INJECTED_COLUMNS), + ).select_from(run.join(rearmed, rearmed.c.id == run.c.id)) + return insert(dlq).from_select([*(col for _, col in _DLQ_PROJECTION), *_DLQ_INJECTED_COLUMNS], archived) + + def _build_insert_stmts(table: "Table", channel: str) -> "dict[tuple[TimerConflict | None, bool, bool], Insert]": """Build the single-row INSERT once per ``(timer conflict, activate_in set, notify)``; values arrive as parameters. The notify variants call ``pg_notify`` in ``RETURNING``, so the NOTIFY costs no extra round trip and - fires only when a row actually landed or was extended (a skipped conflict returns no row). + fires only when a row actually landed or was extended (a skipped conflict returns no row). An extend + of a leased row writes only the ``rearm_*`` columns, which the run's terminal write applies. """ activate_in = func.now() + _seconds("activate_in_seconds") index_elements = [table.c.queue, table.c.timer_id] index_where = table.c.timer_id.is_not(None) excluded = pg_insert(table).excluded - later = excluded.next_attempt_at > table.c.next_attempt_at + pending_later = and_(table.c.acquired_token.is_(None), excluded.next_attempt_at > table.c.next_attempt_at) + rearm_later = and_( + table.c.acquired_token.is_not(None), + excluded.next_attempt_at > func.coalesce(table.c.rearm_at, table.c.next_attempt_at), + ) - def _if_later(name: str, value: "ColumnElement[typing.Any]") -> "ColumnElement[typing.Any]": - return case((later, value), else_=table.c[name]) + def _when( + condition: "ColumnElement[bool]", name: str, value: "ColumnElement[typing.Any]" + ) -> "ColumnElement[typing.Any]": + return case((condition, value), else_=table.c[name]) extend = pg_insert(table).on_conflict_do_update( index_elements=index_elements, index_where=index_where, set_={ - "next_attempt_at": func.greatest(table.c.next_attempt_at, excluded.next_attempt_at), - "payload": _if_later("payload", excluded.payload), - "headers": _if_later("headers", excluded.headers), - "attempts_count": _if_later("attempts_count", literal(0)), - "first_attempt_at": _if_later("first_attempt_at", null()), - "last_attempt_at": _if_later("last_attempt_at", null()), + "next_attempt_at": _when(pending_later, "next_attempt_at", excluded.next_attempt_at), + "payload": _when(pending_later, "payload", excluded.payload), + "headers": _when(pending_later, "headers", excluded.headers), + "attempts_count": _when(pending_later, "attempts_count", literal(0)), + "first_attempt_at": _when(pending_later, "first_attempt_at", null()), + "last_attempt_at": _when(pending_later, "last_attempt_at", null()), + "rearm_at": _when(rearm_later, "rearm_at", excluded.next_attempt_at), + "rearm_payload": _when(rearm_later, "rearm_payload", excluded.payload), + "rearm_headers": _when(rearm_later, "rearm_headers", excluded.headers), }, - where=table.c.acquired_token.is_(None), ) bases: dict[TimerConflict | None, Insert] = { None: insert(table), @@ -426,6 +483,7 @@ def __init__( self._fetch_stmt = _build_fetch_stmt(outbox_table) self._inserter = _OutboxInserter(outbox_table) self._dlq_cte = None if dlq_table is None else _build_dlq_cte(outbox_table, dlq_table) + self._dlq_rearm_cte = None if dlq_table is None else _build_dlq_rearm_cte(outbox_table, dlq_table) @property def table(self) -> "Table": @@ -554,7 +612,7 @@ async def delete_with_lease( msg = "OutboxClient.delete_with_lease requires a live AsyncConnection (got None)" raise TypeError(msg) if dlq_payload is not None: - if self._dlq_cte is None: + if self._dlq_cte is None or self._dlq_rearm_cte is None: # P10: a dlq_payload with no dlq_table would silently degrade to a plain # DELETE — losing the audit row. If the broker/client wiring ever desyncs, # fail loudly rather than drop forensics. @@ -567,10 +625,21 @@ async def delete_with_lease( "last_exception": dlq_payload["last_exception"], } result = await conn.execute(self._dlq_cte, params) + if (result.rowcount or 0) > 0: + return True + result = await conn.execute(self._dlq_rearm_cte, params) return (result.rowcount or 0) > 0 t = self._table - del_stmt = delete(t).where(t.c.id == message_id, t.c.acquired_token == acquired_token) + del_stmt = delete(t).where(t.c.id == message_id, t.c.acquired_token == acquired_token, t.c.rearm_at.is_(None)) result = await conn.execute(del_stmt) + if (result.rowcount or 0) > 0: + return True + rearm_stmt = ( + update(t) + .where(t.c.id == message_id, t.c.acquired_token == acquired_token, t.c.rearm_at.is_not(None)) + .values(**_rearm_values(t)) + ) + result = await conn.execute(rearm_stmt) return (result.rowcount or 0) > 0 async def delete_batch_with_lease( @@ -593,8 +662,32 @@ async def delete_batch_with_lease( if not pairs: return set() t = self._table + wanted = self._wanted_pairs(pairs) + stmt = ( + delete(t) + .where(t.c.id == wanted.c.id, t.c.acquired_token == wanted.c.token, t.c.rearm_at.is_(None)) + .returning(t.c.id, t.c.acquired_token) + ) + result = await conn.execute(stmt) + done = {(row_id, token) for row_id, token in result.all()} + missing = [pair for pair in pairs if pair not in done] + if not missing: + return done + wanted = self._wanted_pairs(missing) + rearm_stmt = ( + update(t) + .where(t.c.id == wanted.c.id, t.c.acquired_token == wanted.c.token, t.c.rearm_at.is_not(None)) + .values(**_rearm_values(t)) + .returning(t.c.id, wanted.c.token) + ) + result = await conn.execute(rearm_stmt) + return done | {(row_id, token) for row_id, token in result.all()} + + def _wanted_pairs(self, pairs: "Sequence[tuple[int, uuid.UUID]]") -> typing.Any: + """``unnest(:ids, :tokens) AS wanted(id, token)``, so the SQL text doesn't depend on ``len(pairs)``.""" + t = self._table ids_type, tokens_type = ARRAY(t.c.id.type), ARRAY(t.c.acquired_token.type) - wanted = ( + return ( func.unnest( cast(bindparam("ids", [pid for pid, _ in pairs], type_=ids_type), ids_type), cast(bindparam("tokens", [ptok for _, ptok in pairs], type_=tokens_type), tokens_type), @@ -602,13 +695,6 @@ async def delete_batch_with_lease( .table_valued("id", "token") .render_derived(name="wanted") ) - stmt = ( - delete(t) - .where(t.c.id == wanted.c.id, t.c.acquired_token == wanted.c.token) - .returning(t.c.id, t.c.acquired_token) - ) - result = await conn.execute(stmt) - return {(row_id, token) for row_id, token in result.all()} async def mark_pending_with_lease( self, @@ -636,15 +722,24 @@ async def mark_pending_with_lease( msg = "OutboxClient.mark_pending_with_lease requires a live AsyncConnection (got None)" raise TypeError(msg) t = self._table - next_attempt_at_expr = func.now() + _seconds("delay") + rearm = t.c.rearm_at.is_not(None) + + def _unless_rearm(name: str, value: "ColumnElement[typing.Any]") -> "ColumnElement[typing.Any]": + return case((rearm, _rearm_values(t)[name]), else_=value) + stmt = ( update(t) .where(t.c.id == message_id, t.c.acquired_token == acquired_token) .values( - next_attempt_at=next_attempt_at_expr, - attempts_count=attempts_count, - first_attempt_at=first_attempt_at, - last_attempt_at=last_attempt_at, + next_attempt_at=_unless_rearm("next_attempt_at", func.now() + _seconds("delay")), + payload=_unless_rearm("payload", t.c.payload), + headers=_unless_rearm("headers", t.c.headers), + attempts_count=case((rearm, 0), else_=attempts_count), + first_attempt_at=case((rearm, null()), else_=literal(first_attempt_at, t.c.first_attempt_at.type)), + last_attempt_at=case((rearm, null()), else_=literal(last_attempt_at, t.c.last_attempt_at.type)), + rearm_at=None, + rearm_payload=None, + rearm_headers=None, acquired_at=None, acquired_token=None, ) @@ -653,15 +748,28 @@ async def mark_pending_with_lease( return (result.rowcount or 0) > 0 async def cancel_timer(self, *, queue: str, timer_id: str, session: SessionOrConnection) -> bool: - """Delete a not-yet-leased ``(queue, timer_id)`` row on the caller's session.""" + """Delete an unleased ``(queue, timer_id)`` row, or drop a leased one's queued rearm, on *session*. + + The row is locked first, so a run that finishes concurrently either completes before the + cancel reads it or waits for the cancel to commit; its rearm can't slip past in between. + """ require_session_or_connection(session, "OutboxClient.cancel_timer") t = self._table - stmt = delete(t).where( - t.c.queue == queue, - t.c.timer_id == timer_id, - t.c.acquired_token.is_(None), + locked = ( + select(t.c.id, t.c.acquired_token).where(t.c.queue == queue, t.c.timer_id == timer_id).with_for_update() ) - result = await session.execute(stmt) + row = (await session.execute(locked)).first() + if row is None: + return False + if row.acquired_token is None: + await session.execute(delete(t).where(t.c.id == row.id)) + return True + clear = ( + update(t) + .where(t.c.id == row.id, t.c.rearm_at.is_not(None)) + .values(rearm_at=None, rearm_payload=None, rearm_headers=None) + ) + result = await session.execute(clear) return (result.rowcount or 0) > 0 # ty: ignore[unresolved-attribute] async def fetch_unprocessed( diff --git a/faststream_outbox/schema.py b/faststream_outbox/schema.py index 48bae23..86727e6 100644 --- a/faststream_outbox/schema.py +++ b/faststream_outbox/schema.py @@ -50,6 +50,7 @@ _TIMER_ID_UQ_SUFFIX = "_timer_id_uq" _LEASED_IDX_SUFFIX = "_leased_idx" _LEASE_CK_SUFFIX = "_lease_ck" +_REARM_CK_SUFFIX = "_rearm_ck" def validate_table_identifiers(table_name: str) -> None: @@ -70,6 +71,7 @@ def validate_table_identifiers(table_name: str) -> None: f"{table_name}{_TIMER_ID_UQ_SUFFIX}", f"{table_name}{_LEASED_IDX_SUFFIX}", f"{table_name}{_LEASE_CK_SUFFIX}", + f"{table_name}{_REARM_CK_SUFFIX}", ) ) longest = max(derived, key=len) @@ -112,6 +114,10 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table Column("acquired_at", DateTime(timezone=True), nullable=True), Column("acquired_token", Uuid, nullable=True), Column("timer_id", String(255), nullable=True), + # A timer extend that arrives while a worker holds the lease, applied when the run ends. + Column("rearm_at", DateTime(timezone=True), nullable=True), + Column("rearm_payload", LargeBinary, nullable=True), + Column("rearm_headers", JSONB, nullable=True), # P8: a lease is either fully set or fully unset. A half-set lease (e.g. a # manual UPDATE that sets acquired_at but not acquired_token) would be # permanently invisible to fetch, cancel_timer, and the metrics — this @@ -120,6 +126,11 @@ def make_outbox_table(metadata: "MetaData", table_name: str = "outbox") -> Table "(acquired_token IS NULL) = (acquired_at IS NULL)", name=f"{table_name}{_LEASE_CK_SUFFIX}", ), + # Only a running timer can carry a queued rearm; on an unleased row nothing would apply it. + CheckConstraint( + "rearm_at IS NULL OR acquired_token IS NOT NULL", + name=f"{table_name}{_REARM_CK_SUFFIX}", + ), ) # Partial index that backs the fetch query's unleased branch: an ordered scan per queue on # (next_attempt_at, id) that stops after `limit` unlocked rows. `id` is part of the key because diff --git a/faststream_outbox/schema_validation.py b/faststream_outbox/schema_validation.py index bd0db2b..5f6ac04 100644 --- a/faststream_outbox/schema_validation.py +++ b/faststream_outbox/schema_validation.py @@ -24,6 +24,7 @@ _LEASE_CK_SUFFIX, _LEASED_IDX_SUFFIX, _PENDING_IDX_SUFFIX, + _REARM_CK_SUFFIX, _TIMER_ID_UQ_SUFFIX, make_dlq_table, make_outbox_table, @@ -192,11 +193,12 @@ def _validate_index_predicates_sync(connection: "Connection", table: "Table", sc "AND con.contype = 'c'", ) -# Expected CHECK-constraint predicates, keyed by the constraint-name suffix make_outbox_table -# uses. Normalized form (lowercased, parens/whitespace collapsed, leading ``check`` stripped) -# of ``(acquired_token IS NULL) = (acquired_at IS NULL)``. +# Expected CHECK-constraint predicates and the invariant each enforces, keyed by the +# constraint-name suffix make_outbox_table uses. Predicates are in normalized form +# (lowercased, parens/whitespace collapsed, leading ``check`` stripped). _EXPECTED_CHECK_CONSTRAINTS = { - _LEASE_CK_SUFFIX: "acquired_token is null = acquired_at is null", + _LEASE_CK_SUFFIX: ("acquired_token is null = acquired_at is null", "the lease invariant"), + _REARM_CK_SUFFIX: ("rearm_at is null or acquired_token is not null", "the queued-rearm invariant"), } @@ -229,10 +231,10 @@ def _validate_check_constraints_sync(connection: "Connection", table: "Table", s # leading ``check`` keyword after normalizing away parens/case/whitespace. live_predicates = {_normalize_predicate(row["definition"]).removeprefix("check ").strip() for row in rows} errors: list[str] = [] - for suffix, want in _EXPECTED_CHECK_CONSTRAINTS.items(): + for suffix, (want, invariant) in _EXPECTED_CHECK_CONSTRAINTS.items(): if want not in live_predicates: errors.append( - f"missing CHECK constraint enforcing '{want}' (the lease invariant; name it e.g. {table.name}{suffix})", + f"missing CHECK constraint enforcing '{want}' ({invariant}; name it e.g. {table.name}{suffix})", ) return errors @@ -252,6 +254,8 @@ def _validate_check_constraints_sync(connection: "Connection", table: "Table", s _PENDING_IDX_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-015" _LEASED_IDX_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-017" +_REARM_UPGRADE_URL = "https://faststream-outbox.modern-python.org/operations/alembic/#upgrading-to-019" +_REARM_COLUMNS = frozenset({"rearm_at", "rearm_payload", "rearm_headers"}) def _compose_schema_mismatch_message(errors: list[str], *, has_blind_drift: bool) -> str: @@ -388,7 +392,10 @@ def _drift_entry_to_error(entry: "tuple[typing.Any, ...]", table_name: str) -> s return f"table '{table_name}' does not exist" if op == "add_column": col = entry[3] - return f"table '{table_name}' missing column '{col.name}'" + error = f"table '{table_name}' missing column '{col.name}'" + if col.name in _REARM_COLUMNS: + error += f"; a table created before 0.19 needs the rearm columns added: {_REARM_UPGRADE_URL}" + return error if op == "modify_type": col_name = entry[3] existing_type = entry[5] diff --git a/faststream_outbox/testing.py b/faststream_outbox/testing.py index b48ae05..af5b5bb 100644 --- a/faststream_outbox/testing.py +++ b/faststream_outbox/testing.py @@ -57,6 +57,21 @@ class _FakeRow: acquired_at: dt.datetime | None = None acquired_token: uuid.UUID | None = None timer_id: str | None = None + rearm_at: dt.datetime | None = None + rearm_payload: bytes | None = None + rearm_headers: dict[str, str] | None = None + + +def _apply_rearm(row: _FakeRow) -> bool: + """Turn a finished run with a queued rearm into a fresh pending timer; False when none is queued.""" + if row.rearm_at is None or row.rearm_payload is None: + return False + row.next_attempt_at, row.payload, row.headers = row.rearm_at, row.rearm_payload, row.rearm_headers + row.rearm_at = row.rearm_payload = row.rearm_headers = None + row.acquired_at = row.acquired_token = None + row.attempts_count = row.deliveries_count = 0 + row.first_attempt_at = row.last_attempt_at = None + return True def _claim_fake_row(row: _FakeRow, *, now: dt.datetime, token: uuid.UUID) -> None: @@ -174,9 +189,13 @@ async def insert( def _extend( row: _FakeRow, *, payload: bytes, headers: dict[str, str] | None, next_attempt_at: dt.datetime ) -> int | None: - """Mirror the real ``ON CONFLICT DO UPDATE ... WHERE acquired_token IS NULL`` extend.""" + """Mirror the real ``ON CONFLICT DO UPDATE`` extend, which queues a rearm on a leased row.""" if row.acquired_token is not None: - return None + if next_attempt_at > (row.rearm_at or row.next_attempt_at): + row.rearm_at = next_attempt_at + row.rearm_payload = payload + row.rearm_headers = dict(headers) if headers is not None else None + return row.id if next_attempt_at > row.next_attempt_at: row.next_attempt_at = next_attempt_at row.payload = payload @@ -321,7 +340,8 @@ async def delete_with_lease( for col in _DLQ_INJECTED_COLUMNS: dlq_row[col] = dlq_payload[col] self._dlq_rows.append(dlq_row) - del self._rows[i] + if not _apply_rearm(row): + del self._rows[i] return True return False @@ -340,7 +360,8 @@ async def delete_batch_with_lease( row = self._rows[i] if row.acquired_token is not None and (row.id, row.acquired_token) in wanted: deleted.add((row.id, row.acquired_token)) - del self._rows[i] + if not _apply_rearm(row): + del self._rows[i] return deleted async def mark_pending_with_lease( @@ -356,6 +377,10 @@ async def mark_pending_with_lease( ) -> bool: for row in self._rows: if row.id == message_id and _lease_token_matches(row.acquired_token, acquired_token): + deliveries = row.deliveries_count + if _apply_rearm(row): + row.deliveries_count = deliveries + return True row.next_attempt_at = utcnow() + dt.timedelta(seconds=max(0.0, delay_seconds)) row.attempts_count = attempts_count row.first_attempt_at = first_attempt_at @@ -366,12 +391,16 @@ async def mark_pending_with_lease( return False async def cancel_timer(self, *, queue: str, timer_id: str, session: SessionOrConnection) -> bool: - """Mirror :meth:`OutboxBroker.cancel_timer` — drop a not-yet-leased timer row.""" + """Mirror :meth:`OutboxBroker.cancel_timer` — drop an unleased timer row or a leased one's queued rearm.""" require_session_or_connection(session, "FakeOutboxClient.cancel_timer") for i, row in enumerate(self._rows): - if row.queue == queue and row.timer_id == timer_id and row.acquired_token is None: - del self._rows[i] - return True + if row.queue == queue and row.timer_id == timer_id: + if row.acquired_token is None: + del self._rows[i] + return True + queued = row.rearm_at is not None + row.rearm_at = row.rearm_payload = row.rearm_headers = None + return queued return False async def fetch_unprocessed( diff --git a/tests/test_client_contract.py b/tests/test_client_contract.py index 9528b72..adbb5da 100644 --- a/tests/test_client_contract.py +++ b/tests/test_client_contract.py @@ -41,7 +41,7 @@ from faststream_outbox._time import utcnow from faststream_outbox.client import OutboxClient from faststream_outbox.message import OutboxInnerMessage -from faststream_outbox.testing import FakeOutboxClient +from faststream_outbox.testing import FakeOutboxClient, fake_session pytestmark = pytest.mark.asyncio @@ -169,11 +169,17 @@ async def get_row(self, message_id: int) -> dict[str, object] | None: "deliveries_count": row.deliveries_count, "next_attempt_at": row.next_attempt_at, "timer_id": row.timer_id, + "rearm_at": row.rearm_at, + "rearm_payload": row.rearm_payload, + "rearm_headers": row.rearm_headers, } async def dlq_rows(self) -> list[dict[str, object]]: return [dict(r) for r in self._client.dlq_rows] + async def cancel(self, *, queue: str, timer_id: str) -> bool: + return await self._client.cancel_timer(queue=queue, timer_id=timer_id, session=fake_session()) + async def call(self, method: str, **kwargs: object) -> object: return await getattr(self._client, method)(**kwargs) @@ -307,8 +313,14 @@ async def get_row(self, message_id: int) -> dict[str, object] | None: return None row = dict(mapping) row["payload"] = bytes(row["payload"]) + if row["rearm_payload"] is not None: + row["rearm_payload"] = bytes(row["rearm_payload"]) return row + async def cancel(self, *, queue: str, timer_id: str) -> bool: + async with self._sessions() as session, session.begin(): + return await self._client.cancel_timer(queue=queue, timer_id=timer_id, session=session) + async def dlq_rows(self) -> list[dict[str, object]]: async with self._engine.connect() as conn: result = await conn.execute(select(self._dlq)) @@ -492,30 +504,156 @@ async def test_insert_extend_resets_the_retry_history(contract: _Harness) -> Non assert row["deliveries_count"] == 1 +async def _claimed_timer(contract: _Harness, *, payload: bytes = b"old") -> tuple[int, uuid.UUID]: + rid = await contract.insert(queue="orders", payload=payload, timer_id="t1") + assert rid is not None + [msg] = await contract.fetch(["orders"], limit=10, lease_ttl_seconds=60.0) + assert msg.acquired_token is not None + return rid, msg.acquired_token + + +async def _extend( + contract: _Harness, *, payload: bytes, minutes: float, headers: dict[str, str] | None = None +) -> int | None: + return await contract.insert( + queue="orders", + payload=payload, + headers=headers, + activate_at=utcnow() + dt.timedelta(minutes=minutes), + timer_id="t1", + timer_conflict="extend", + ) + + @pytest.mark.parametrize("acquired_age", [0.0, 3600.0]) -async def test_insert_extend_leaves_a_leased_timer_to_its_worker(contract: _Harness, acquired_age: float) -> None: - """INVARIANT: an extend never touches a leased timer and returns ``None`` for it. +async def test_insert_extend_queues_a_rearm_on_a_leased_timer(contract: _Harness, acquired_age: float) -> None: + """INVARIANT: an extend on a leased timer queues a rearm and leaves the running instance alone. - Once a worker holds the lease, its handler may already have acted on the old deadline. Rewriting - the row, or clearing the lease so the upsert lands, would report a refresh the handler never saw - and let a stale completion delete the refreshed timer. An expired lease is still a lease: the - fetch loop reclaims the row, so the caller must handle ``None`` the same way. + The worker's handler may already have acted on the deadline it claimed, so the extend must not + touch the lease, the claimed deadline or the payload the handler sees. Dropping the extend instead + would lose a heartbeat that arrived while the expiry ran; clearing the lease so the update lands + would let the stale completion delete the refreshed timer. An expired lease is still a lease. """ token = uuid.uuid4() rid = await contract.seed( queue="orders", payload=b"old", timer_id="t1", leased_token=token, acquired_age=acquired_age ) before = await contract.get_row(rid) + assert before is not None + at = utcnow() + dt.timedelta(minutes=10) extended = await contract.insert( - queue="orders", - payload=b"new", - activate_at=utcnow() + dt.timedelta(minutes=10), - timer_id="t1", - timer_conflict="extend", + queue="orders", payload=b"new", headers={"k": "v"}, activate_at=at, timer_id="t1", timer_conflict="extend" ) - assert extended is None - assert await contract.get_row(rid) == before + assert extended == rid + after = await contract.get_row(rid) + assert after is not None + assert (after["rearm_at"], after["rearm_payload"], after["rearm_headers"]) == (at, b"new", {"k": "v"}) + running = ("acquired_token", "acquired_at", "next_attempt_at", "payload", "headers", "deliveries_count") + assert {k: after[k] for k in running} == {k: before[k] for k in running} + + +async def test_insert_extend_never_moves_a_queued_rearm_earlier(contract: _Harness) -> None: + rid, _ = await _claimed_timer(contract) + assert await _extend(contract, payload=b"later", minutes=10) == rid + queued = await contract.get_row(rid) + assert queued is not None + assert await _extend(contract, payload=b"stale", minutes=1) == rid + row = await contract.get_row(rid) + assert row is not None + assert (row["rearm_at"], row["rearm_payload"]) == (queued["rearm_at"], b"later") + + +async def test_insert_extend_queues_nothing_behind_a_later_claimed_deadline(contract: _Harness) -> None: + token = uuid.uuid4() + rid = await contract.seed(queue="orders", timer_id="t1", next_attempt_offset=600.0, leased_token=token) + assert await _extend(contract, payload=b"stale", minutes=1) == rid + row = await contract.get_row(rid) + assert row is not None + assert row["rearm_at"] is None + + +async def test_delete_applies_a_queued_rearm(contract: _Harness) -> None: + rid, token = await _claimed_timer(contract) + assert await _extend(contract, payload=b"new", minutes=10, headers={"k": "v"}) == rid + queued = await contract.get_row(rid) + assert queued is not None assert await contract.delete(rid, token) + row = await contract.get_row(rid) + assert row is not None + assert row["acquired_token"] is None + assert row["acquired_at"] is None + assert (row["next_attempt_at"], row["payload"], row["headers"]) == (queued["rearm_at"], b"new", {"k": "v"}) + assert (row["rearm_at"], row["rearm_payload"], row["rearm_headers"]) == (None, None, None) + assert (row["attempts_count"], row["deliveries_count"]) == (0, 0) + + +async def test_delete_with_a_lost_lease_keeps_the_queued_rearm(contract: _Harness) -> None: + rid, _ = await _claimed_timer(contract) + assert await _extend(contract, payload=b"new", minutes=10) == rid + assert not await contract.delete(rid, uuid.uuid4()) + row = await contract.get_row(rid) + assert row is not None + assert row["rearm_payload"] == b"new" + + +async def test_delete_batch_applies_a_queued_rearm(contract: _Harness) -> None: + rid, token = await _claimed_timer(contract) + other = await contract.seed(queue="orders", leased_token=token) + assert await _extend(contract, payload=b"new", minutes=10) == rid + assert await contract.delete_batch([(rid, token), (other, token)]) == {(rid, token), (other, token)} + assert await contract.queue_ids("orders") == [rid] + row = await contract.get_row(rid) + assert row is not None + assert (row["acquired_token"], row["payload"], row["rearm_at"]) == (None, b"new", None) + + +async def test_delete_with_dlq_archives_the_failed_run_and_applies_a_queued_rearm(contract: _Harness) -> None: + rid, token = await _claimed_timer(contract, payload=b"failed") + assert await _extend(contract, payload=b"new", minutes=10) == rid + assert await contract.delete(rid, token, dlq_payload={"failure_reason": "rejected", "last_exception": "boom"}) + [archived] = await contract.dlq_rows() + assert (archived["original_id"], archived["payload"], archived["failure_reason"]) == (rid, b"failed", "rejected") + row = await contract.get_row(rid) + assert row is not None + assert (row["acquired_token"], row["payload"], row["deliveries_count"]) == (None, b"new", 0) + + +async def test_mark_pending_lets_a_queued_rearm_win(contract: _Harness) -> None: + """INVARIANT: a retry yields to a queued rearm, and keeps the run's delivery count. + + The rearm is newer work than the failed run, so the timer moves to the rearm deadline and + payload with a fresh retry history. ``deliveries_count`` is not reset, so ``max_deliveries`` + still stops a handler that fails on every run while heartbeats keep rearming it. + """ + rid, token = await _claimed_timer(contract) + assert await _extend(contract, payload=b"new", minutes=10) == rid + queued = await contract.get_row(rid) + assert queued is not None + now = utcnow() + assert await contract.mark_pending( + rid, token, delay_seconds=5.0, attempts_count=2, first_attempt_at=now, last_attempt_at=now + ) + row = await contract.get_row(rid) + assert row is not None + assert (row["next_attempt_at"], row["payload"], row["rearm_at"]) == (queued["rearm_at"], b"new", None) + assert (row["attempts_count"], row["first_attempt_at"], row["last_attempt_at"]) == (0, None, None) + assert (row["acquired_token"], row["deliveries_count"]) == (None, 1) + + +async def test_cancel_timer_drops_a_queued_rearm(contract: _Harness) -> None: + rid, token = await _claimed_timer(contract) + assert await _extend(contract, payload=b"new", minutes=10) == rid + assert await contract.cancel(queue="orders", timer_id="t1") + row = await contract.get_row(rid) + assert row is not None + assert (row["acquired_token"], row["rearm_at"]) == (token, None) + assert await contract.delete(rid, token) + assert await contract.queue_ids("orders") == [] + + +async def test_cancel_timer_on_a_leased_timer_without_a_rearm_returns_false(contract: _Harness) -> None: + await _claimed_timer(contract) + assert not await contract.cancel(queue="orders", timer_id="t1") async def test_insert_extend_same_timer_id_on_another_queue_lands(contract: _Harness) -> None: diff --git a/tests/test_fake.py b/tests/test_fake.py index a85d4b2..ecd0bca 100644 --- a/tests/test_fake.py +++ b/tests/test_fake.py @@ -323,6 +323,23 @@ async def handle(body: str) -> None: assert received == [] +async def test_handler_extending_its_own_timer_leaves_it_pending_without_redispatch() -> None: + broker = _make_broker() + received: list[str] = [] + + @broker.subscriber("timers") + async def handle(body: str) -> None: + received.append(body) + await broker.publish("again", queue="timers", timer_id="op-1", timer_conflict="extend", session=fake_session()) + + test_broker = TestOutboxBroker(broker) + async with test_broker: + await broker.publish("first", queue="timers", timer_id="op-1", session=fake_session()) + assert received == ["first"] + [row] = test_broker.fake_client.rows + assert (row.acquired_token, row.rearm_at, row.payload) == (None, None, encode_payload("again")[0]) + + @pytest.mark.parametrize( ("kwargs", "match"), [ diff --git a/tests/test_integration.py b/tests/test_integration.py index c4471f2..3bcde3f 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -959,7 +959,7 @@ async def test_extend_holds_off_a_claim_until_it_commits(pg_engine: AsyncEngine, assert row["acquired_token"] is None -async def test_extend_waits_for_an_uncommitted_claim_and_returns_none( +async def test_extend_waits_for_an_uncommitted_claim_and_queues_a_rearm( pg_engine: AsyncEngine, outbox_table: Table ) -> None: broker = OutboxBroker(pg_engine, outbox_table=outbox_table) @@ -987,10 +987,118 @@ async def _extend() -> int | None: extend = asyncio.create_task(_extend()) await asyncio.sleep(0.3) assert not extend.done() - assert await extend is None + assert await extend is not None row = await _timer_row(pg_engine, outbox_table, "op-1") assert row["acquired_token"] == token assert json.loads(row["payload"]) == {"v": 1} + assert json.loads(row["rearm_payload"]) == {"v": 2} + + +async def test_completion_waits_for_an_uncommitted_extend_and_applies_it( + pg_engine: AsyncEngine, outbox_table: Table +) -> None: + """INVARIANT: a completion never deletes a timer whose rearm is being queued. + + The extend locks the leased row, so the completion's ``DELETE ... WHERE rearm_at IS NULL`` waits + and re-checks the row once the extend commits, then falls through to applying the rearm. Reading + ``rearm_at`` before deleting, outside that single statement, would let a heartbeat that commits in + between vanish together with the finished run. + """ + broker = OutboxBroker(pg_engine, outbox_table=outbox_table) + client = OutboxClient(pg_engine, outbox_table) + session_factory = async_sessionmaker(pg_engine, expire_on_commit=False) + async with session_factory() as session, session.begin(): + await broker.publish({"v": 1}, queue="orders", session=session, timer_id="op-1") + async with pg_engine.connect() as conn: + [msg] = await client.fetch(conn, ["orders"], limit=10, lease_ttl_seconds=60.0) + token = msg.acquired_token + assert token is not None + + async def _complete() -> bool: + async with pg_engine.connect() as raw: + writer = await raw.execution_options(isolation_level="AUTOCOMMIT") + return await client.delete_with_lease(writer, msg.id, token) + + async with session_factory() as session, session.begin(): + await broker.publish( + {"v": 2}, + queue="orders", + session=session, + activate_in=dt.timedelta(minutes=10), + timer_id="op-1", + timer_conflict="extend", + ) + complete = asyncio.create_task(_complete()) + await asyncio.sleep(0.3) + assert not complete.done() + assert await complete + row = await _timer_row(pg_engine, outbox_table, "op-1") + assert row["acquired_token"] is None + assert json.loads(row["payload"]) == {"v": 2} + assert row["rearm_at"] is None + + +async def test_ping_during_a_running_expiry_fires_the_timer_again(pg_engine: AsyncEngine, outbox_table: Table) -> None: + received: list[dict] = [] + broker = OutboxBroker(pg_engine, outbox_table=outbox_table) + handler_started = asyncio.Event() + release_handler = asyncio.Event() + + @broker.subscriber("orders", min_fetch_interval=0.05, max_fetch_interval=0.1) + async def handle(body: dict) -> None: + received.append(body) + handler_started.set() + await release_handler.wait() + + session_factory = async_sessionmaker(pg_engine, expire_on_commit=False) + + async def _ping(n: int, activate_in: dt.timedelta) -> int | None: + async with session_factory() as session, session.begin(): + return await broker.publish( + {"ping": n}, + queue="orders", + session=session, + activate_in=activate_in, + timer_id="op-1", + timer_conflict="extend", + ) + + async with broker: + await _ping(1, dt.timedelta(0)) + await asyncio.wait_for(handler_started.wait(), timeout=3.0) + assert await _ping(2, dt.timedelta(seconds=0.3)) is not None + release_handler.set() + await _wait_until(lambda: len(received) == 2, timeout=3.0) + assert received == [{"ping": 1}, {"ping": 2}] + assert await _row_count(pg_engine, outbox_table) == 0 + + +async def test_handler_rearms_its_own_timer(pg_engine: AsyncEngine, outbox_table: Table) -> None: + received: list[int] = [] + broker = OutboxBroker(pg_engine, outbox_table=outbox_table) + session_factory = async_sessionmaker(pg_engine, expire_on_commit=False) + + @broker.subscriber("ticks", min_fetch_interval=0.05, max_fetch_interval=0.1) + async def handle(body: dict) -> None: + received.append(body["n"]) + if body["n"] < 3: + async with session_factory() as session, session.begin(): + await broker.publish( + {"n": body["n"] + 1}, + queue="ticks", + session=session, + activate_in=dt.timedelta(seconds=0.1), + timer_id="tick", + timer_conflict="extend", + ) + + async with broker: + async with session_factory() as session, session.begin(): + await broker.publish({"n": 1}, queue="ticks", session=session, timer_id="tick") + await _wait_until(lambda: len(received) == 3, timeout=5.0) + await asyncio.sleep(0.3) + assert received == [1, 2, 3] + assert await _row_count(pg_engine, outbox_table) == 0 async def test_concurrent_extends_keep_one_row_at_the_latest_deadline( @@ -2544,6 +2652,26 @@ async def handle(body: dict[str, Any]) -> None: assert await _row_count(pg_engine, outbox_table) == 1 +async def test_validate_schema_fails_when_rearm_check_constraint_missing( + pg_engine: AsyncEngine, + outbox_table: Table, +) -> None: + drop_sql = f'ALTER TABLE "{outbox_table.name}" DROP CONSTRAINT "{outbox_table.name}_rearm_ck"' + async with pg_engine.begin() as conn: + await conn.exec_driver_sql(drop_sql) + client = OutboxClient(pg_engine, outbox_table) + with pytest.raises(RuntimeError, match="missing CHECK constraint enforcing 'rearm_at is null or acquired_token"): + await client.validate_schema() + + +async def test_validate_schema_fails_when_rearm_columns_missing(pg_engine: AsyncEngine, outbox_table: Table) -> None: + async with pg_engine.begin() as conn: + await conn.exec_driver_sql(f'ALTER TABLE "{outbox_table.name}" DROP COLUMN rearm_headers') + client = OutboxClient(pg_engine, outbox_table) + with pytest.raises(RuntimeError, match=r"missing column 'rearm_headers'.*operations/alembic/#upgrading-to-019"): + await client.validate_schema() + + async def test_validate_schema_fails_when_lease_check_constraint_missing( pg_engine: AsyncEngine, outbox_table: Table, diff --git a/tests/test_unit.py b/tests/test_unit.py index 7f0fb25..46077d4 100644 --- a/tests/test_unit.py +++ b/tests/test_unit.py @@ -10,6 +10,7 @@ import warnings from collections.abc import AsyncGenerator, Sequence from contextlib import asynccontextmanager +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -125,6 +126,9 @@ def test_make_outbox_table_columns_present() -> None: "acquired_at", "acquired_token", "timer_id", + "rearm_at", + "rearm_payload", + "rearm_headers", } assert {c.name for c in t.columns} == expected assert t.name == "my_outbox" @@ -1152,27 +1156,41 @@ async def test_broker_cancel_timer_rejects_non_async_session() -> None: await broker.cancel_timer(queue="orders", timer_id="x", session=object()) # ty: ignore[invalid-argument-type] -async def test_broker_cancel_timer_emits_delete_with_lease_guard() -> None: - """INVARIANT: ``cancel_timer`` deletes only unleased rows -- the ``acquired_token IS NULL`` guard. +def _cancel_session(row: object) -> AsyncMock: + session = AsyncMock(spec=AsyncSession) + result = MagicMock() + result.first.return_value = row + result.rowcount = 1 + session.execute.return_value = result + return session + + +async def test_broker_cancel_timer_never_deletes_a_leased_row() -> None: + """INVARIANT: ``cancel_timer`` deletes only an unleased row; on a leased one it only drops the queued rearm. Cancelling a row whose handler is already in flight would delete it out from under the worker, and the worker's terminal write would then find nothing: the delivery has happened but no - record of it survives. The guard makes an in-flight cancel a no-op returning ``False``, letting - the delivery complete normally. Dropping it turns a race that is currently invisible into - silent double-processing. + record of it survives. A leased row keeps its lease and its run; what the cancel removes is the + rearm that would have turned it back into a pending timer. """ broker = _make_broker(engine=MagicMock()) - session = AsyncMock(spec=AsyncSession) - session.execute.return_value.rowcount = 1 - deleted = await broker.cancel_timer(queue="orders", timer_id="email-1", session=session) - assert deleted is True - delete_stmt = session.execute.await_args_list[0].args[0] - sql = str(delete_stmt) - assert "DELETE" in sql - assert "acquired_token IS NULL" in sql - params = delete_stmt.compile().params - assert params["queue_1"] == "orders" - assert params["timer_id_1"] == "email-1" + session = _cancel_session(SimpleNamespace(id=7, acquired_token=uuid.uuid4())) + assert await broker.cancel_timer(queue="orders", timer_id="email-1", session=session) is True + locked, cleared = (call.args[0] for call in session.execute.await_args_list) + assert "FOR UPDATE" in str(locked.compile(dialect=postgresql.dialect())) + assert locked.compile().params == {"queue_1": "orders", "timer_id_1": "email-1"} + sql = str(cleared) + assert sql.startswith("UPDATE") + assert "rearm_at IS NOT NULL" in sql + + +async def test_broker_cancel_timer_deletes_an_unleased_row() -> None: + broker = _make_broker(engine=MagicMock()) + session = _cancel_session(SimpleNamespace(id=7, acquired_token=None)) + assert await broker.cancel_timer(queue="orders", timer_id="email-1", session=session) is True + deleted = session.execute.await_args_list[1].args[0] + assert str(deleted).startswith("DELETE") + assert deleted.compile().params == {"id_1": 7} async def test_broker_fetch_unprocessed_rejects_non_async_session() -> None: @@ -1234,10 +1252,10 @@ async def test_broker_fetch_unprocessed_respects_explicit_limit() -> None: async def test_broker_cancel_timer_returns_false_when_nothing_deleted() -> None: broker = _make_broker(engine=MagicMock()) - session = AsyncMock(spec=AsyncSession) - session.execute.return_value.rowcount = 0 + session = _cancel_session(None) deleted = await broker.cancel_timer(queue="orders", timer_id="x", session=session) assert deleted is False + assert session.execute.await_count == 1 async def test_broker_publish_batch_executes_single_insert_for_many_rows() -> None: @@ -1416,6 +1434,11 @@ def test_compose_schema_mismatch_message_joins_multiple_errors() -> None: _CK_CONVENTION = {"ck": "ck_%(table_name)s_%(constraint_name)s"} _LEASE_CK_PREDICATE = "acquired_token is null = acquired_at is null" _LEASE_CK_DEFINITION = "CHECK (((acquired_token IS NULL) = (acquired_at IS NULL)))" +_REARM_CK_PREDICATE = "rearm_at is null or acquired_token is not null" +_REARM_CK_ROW = { + "name": "outbox_faststream_rearm_ck", + "definition": "CHECK (((rearm_at IS NULL) OR (acquired_token IS NOT NULL)))", +} def _mock_check_constraint_connection(rows: list[dict[str, str]]) -> MagicMock: @@ -1429,7 +1452,10 @@ def test_validate_check_constraints_passes_under_convention_resolved_name() -> N """A constraint carrying the convention-doubled name still passes — matched on predicate.""" table = make_outbox_table(MetaData(naming_convention=_CK_CONVENTION), table_name="outbox_faststream") connection = _mock_check_constraint_connection( - [{"name": "ck_outbox_faststream_outbox_faststream_lease_ck", "definition": _LEASE_CK_DEFINITION}], + [ + {"name": "ck_outbox_faststream_outbox_faststream_lease_ck", "definition": _LEASE_CK_DEFINITION}, + _REARM_CK_ROW, + ], ) assert _validate_check_constraints_sync(connection, table, None) == [] @@ -1438,7 +1464,7 @@ def test_validate_check_constraints_passes_under_literal_name() -> None: """The literal ``
_lease_ck`` a hand-written migration creates also passes — same predicate.""" table = make_outbox_table(MetaData(naming_convention=_CK_CONVENTION), table_name="outbox_faststream") connection = _mock_check_constraint_connection( - [{"name": "outbox_faststream_lease_ck", "definition": _LEASE_CK_DEFINITION}], + [{"name": "outbox_faststream_lease_ck", "definition": _LEASE_CK_DEFINITION}, _REARM_CK_ROW], ) assert _validate_check_constraints_sync(connection, table, None) == [] @@ -1446,13 +1472,17 @@ def test_validate_check_constraints_passes_under_literal_name() -> None: def test_validate_check_constraints_missing_describes_predicate() -> None: """When no live CHECK enforces the predicate, the error names the invariant, not a guessed name.""" table = make_outbox_table(MetaData(), table_name="outbox_faststream") - connection = _mock_check_constraint_connection([]) # constraint absent from the live DB + connection = _mock_check_constraint_connection([]) # constraints absent from the live DB errors = _validate_check_constraints_sync(connection, table, None) assert errors == [ ( f"missing CHECK constraint enforcing '{_LEASE_CK_PREDICATE}' " f"(the lease invariant; name it e.g. outbox_faststream_lease_ck)" ), + ( + f"missing CHECK constraint enforcing '{_REARM_CK_PREDICATE}' " + f"(the queued-rearm invariant; name it e.g. outbox_faststream_rearm_ck)" + ), ] @@ -1461,7 +1491,7 @@ def test_validate_check_constraints_drifted_predicate_reported_as_missing() -> N table = make_outbox_table(MetaData(), table_name="outbox_faststream") drifted = "CHECK ((acquired_token IS NOT NULL OR acquired_at IS NULL))" connection = _mock_check_constraint_connection( - [{"name": "outbox_faststream_lease_ck", "definition": drifted}], + [{"name": "outbox_faststream_lease_ck", "definition": drifted}, _REARM_CK_ROW], ) errors = _validate_check_constraints_sync(connection, table, None) assert errors == [