perf: send the publish NOTIFY in the INSERT's RETURNING (#235)

* perf: send the publish NOTIFY in the INSERT's RETURNING

* test: accept the uncast pg_notify bind that SQLAlchemy 2.0.0 renders

* ci: measure coverage with sys.monitoring

Author: lesnik512

docs/concepts/performance.md

@@ -45,16 +45,18 @@ why).
 | Need more parallel handler capacity | `max_workers` / `fetch_batch_size` | More concurrency, but mind the [connection budget](../usage/subscriber.md#connection-budget). |
 
 Single-publish-per-transaction workloads see no change from NOTIFY dedup (one row
-still emits one NOTIFY), and the batched-flush win is largest at low `max_workers`
+still emits one NOTIFY), but that NOTIFY rides in the INSERT's `RETURNING`, so it
+adds no round trip of its own. The batched-flush win is largest at low `max_workers`
 (where per-row deletes serialise) and narrows as worker parallelism rises.
 
 ## The three levers
 
 ### Producer: NOTIFY dedup (automatic)
 
 `broker.publish` / `publish_batch` emit one `pg_notify` per (transaction,
-queue), so N publishes to the same queue in one transaction cost one NOTIFY
-round-trip. Postgres already coalesces identical notifications per transaction
+queue), so N publishes to the same queue in one transaction cost at most one
+NOTIFY. A single `publish` puts it in the INSERT's `RETURNING`, so it costs no
+round trip; `publish_batch` sends it as one extra statement. Postgres already coalesces identical notifications per transaction
 at delivery, so extra NOTIFYs would be pure waste. It is default-on, has no knob, and is
 behavior-preserving: the subscriber still gets its wake, fired inline at the first
 publish. See

docs/introduction/how-it-works.md

@@ -47,11 +47,13 @@ async with session_factory() as session, session.begin():
 `publish_batch(*bodies, queue, session, ...)` does the same with a single
 round-trip for many rows.
 
-The producer also emits `SELECT pg_notify('outbox_<table>', queue)` on the
-caller's session right after the INSERT, except when the row is
-genuinely future-dated (a future `activate_in` / `activate_at`; a *past*
-`activate_at`, e.g. a recovered idempotency token, still notifies) or a
-`timer_id` conflict made the insert a no-op. NOTIFY is transactional, so listeners only see it
+A single `publish` also sends `pg_notify('outbox_<table>', queue)` in the
+INSERT's `RETURNING` clause, so the wakeup costs no extra round trip.
+`publish_batch` sends a separate `SELECT pg_notify(...)` after its INSERT.
+Neither notifies when the row is genuinely future-dated (a future
+`activate_in` / `activate_at`; a *past* `activate_at`, e.g. a recovered
+idempotency token, still notifies), and a `timer_id` conflict returns no
+row, so it sends no NOTIFY either. NOTIFY is transactional, so listeners only see it
 after the user's transaction commits, so atomicity with the row insert is
 automatic.
 

faststream_outbox/publisher/producer.py

@@ -13,7 +13,7 @@
 
 from faststream._internal.endpoint.utils import ParserComposition
 from faststream._internal.parser import DefaultCodec
-from sqlalchemy import Float, Table, bindparam, func, insert, text
+from sqlalchemy import Float, Table, bindparam, func, insert, literal, text
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 
 from faststream_outbox._scheduling import is_future_dated, resolve_next_attempt_client_side
@@ -33,8 +33,12 @@
     from sqlalchemy import Insert
 
 
-def _build_insert_stmts(table: Table) -> "dict[tuple[bool, bool], Insert]":
-    """Build the single-row INSERT once per ``(timer_id set, activate_in set)``; values arrive as parameters."""
+def _build_insert_stmts(table: Table, channel: str) -> "dict[tuple[bool, bool, bool], Insert]":
+    """Build the single-row INSERT once per ``(timer_id set, 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 (an ``ON CONFLICT DO NOTHING`` skip returns no row).
+    """
     # Server-side compute keeps timing immune to worker/DB clock skew (mirrors
     # client.mark_pending_with_lease).
     activate_in = func.now() + func.make_interval(0, 0, 0, 0, 0, 0, bindparam("activate_in_seconds", type_=Float))
@@ -43,12 +47,13 @@ def _build_insert_stmts(table: Table) -> "dict[tuple[bool, bool], Insert]":
         index_elements=[table.c.queue, table.c.timer_id],
         index_where=table.c.timer_id.is_not(None),
     )
-    return {
-        (False, False): plain.returning(table.c.id),
-        (False, True): plain.values(next_attempt_at=activate_in).returning(table.c.id),
-        (True, False): timer.returning(table.c.id),
-        (True, True): timer.values(next_attempt_at=activate_in).returning(table.c.id),
-    }
+    notify = func.pg_notify(literal(channel), table.c.queue)
+    stmts: dict[tuple[bool, bool, bool], Insert] = {}
+    for has_timer, base in ((False, plain), (True, timer)):
+        for has_activate_in, stmt in ((False, base), (True, base.values(next_attempt_at=activate_in))):
+            stmts[has_timer, has_activate_in, False] = stmt.returning(table.c.id)
+            stmts[has_timer, has_activate_in, True] = stmt.returning(table.c.id, notify)
+    return stmts
 
 
 class OutboxProducer:
@@ -80,7 +85,7 @@ def __init__(
         self._parser = ParserComposition(parser, default.parse_message)
         self._decoder = ParserComposition(decoder, default.decode_message)
         self._metrics_recorder = metrics_recorder
-        self._insert_stmts = _build_insert_stmts(table)
+        self._insert_stmts = _build_insert_stmts(table, self._channel)
 
     def _emit_metric(self, event: str, tags: "Mapping[str, typing.Any]") -> None:
         # Delegate to the shared swallow-and-DEBUG-log helper so the recorder-isolation
@@ -143,16 +148,20 @@ async def _do_publish(
             params["timer_id"] = cmd.timer_id
         # Skip NOTIFY only when the row is genuinely future-dated. A past activate_at
         # (e.g. a recovered idempotency token) is immediately eligible — fire NOTIFY.
-        now = utcnow()
-        is_future = is_future_dated(cmd.activate_in, cmd.activate_at, now)
-
-        stmt = self._insert_stmts[cmd.timer_id is not None, cmd.activate_in is not None]
+        # Future-dated rows skip NOTIFY: listeners can't act before the gate opens, and
+        # polling fires them at the next tick.
+        emitted = self._notified_queues(cmd.session)
+        notify = not is_future_dated(cmd.activate_in, cmd.activate_at, utcnow()) and (
+            emitted is None or cmd.queue not in emitted
+        )
+        stmt = self._insert_stmts[cmd.timer_id is not None, cmd.activate_in is not None, notify]
         result = await cmd.session.execute(stmt, params)
         row_id: int | None = result.scalar()
-        # Skip NOTIFY for future-dated rows (listeners can't act before the gate
-        # opens — polling fires them at next tick) and on conflict (no row landed).
-        if row_id is not None and not is_future:
-            await self._notify(cmd.session, cmd.queue)
+        if row_id is not None and notify:
+            # Re-read: under autobegin the transaction only exists once the INSERT has run.
+            emitted = self._notified_queues(cmd.session)
+            if emitted is not None:
+                emitted.add(cmd.queue)
         return row_id
 
     async def publish_batch(self, cmd: OutboxPublishCommand) -> None:
@@ -206,6 +215,14 @@ async def publish_batch(self, cmd: OutboxPublishCommand) -> None:
     async def request(self, cmd: OutboxPublishCommand) -> typing.NoReturn:
         raise NotImplementedError(_REQUEST_UNSUPPORTED_MSG)
 
+    def _notified_queues(self, session: typing.Any) -> set[str] | None:
+        """Return the queues already notified in *session*'s innermost transaction, or None outside one."""
+        sync_session = session.sync_session
+        txn = sync_session.get_nested_transaction() or sync_session.get_transaction()
+        if txn is None:
+            return None
+        return self._notified.setdefault(txn, set())
+
     async def _notify(self, session: typing.Any, queue: str) -> None:
         # ``pg_notify(:channel, :payload)`` -- parameterized so channel and payload
         # bind cleanly (raw NOTIFY accepts only literals -- injection-prone). Runs on
@@ -224,10 +241,8 @@ async def _notify(self, session: typing.Any, queue: str) -> None:
         # defeat the memo. The Session holds the sync transaction strongly for its lifetime,
         # so the key survives across publishes whether the caller uses `session.begin()` or
         # autobegin. No active transaction -> emit unconditionally.
-        sync_session = session.sync_session
-        txn = sync_session.get_nested_transaction() or sync_session.get_transaction()
-        if txn is not None:
-            emitted = self._notified.setdefault(txn, set())
+        emitted = self._notified_queues(session)
+        if emitted is not None:
             if queue in emitted:
                 return
             emitted.add(queue)

pyproject.toml

@@ -149,6 +149,9 @@ Changelog = "https://github.com/modern-python/faststream-outbox/releases"
 
 [tool.coverage.run]
 omit = ["benchmarks/*"]
+# The C tracer intermittently drops executed lines in async tests; 3.11 falls back to it.
+core = "sysmon"
+disable_warnings = ["no-sysmon"]
 
 [tool.coverage.report]
 fail_under = 100

tests/test_integration.py

@@ -9,6 +9,7 @@
 from typing import Any
 from unittest import mock
 
+import asyncpg
 import pytest
 from faststream import BaseMiddleware, Context, Depends
 from faststream.kafka import KafkaBroker, TestKafkaBroker
@@ -2529,6 +2530,16 @@ async def test_notify_dedup_spans_publish_batch_and_publish(pg_engine, outbox_ta
     assert count[0] == 1
 
 
+async def test_notify_dedup_spans_publish_and_publish_batch(pg_engine: AsyncEngine, outbox_table: Table) -> None:
+    broker = OutboxBroker(pg_engine, outbox_table=outbox_table)
+    session_factory = async_sessionmaker(pg_engine, expire_on_commit=False)
+    with _count_pg_notify(pg_engine) as count:
+        async with session_factory() as session, session.begin():
+            await broker.publish({"n": 1}, queue="orders", session=session)  # pg_notify in RETURNING
+            await broker.publish_batch({"n": 2}, {"n": 3}, queue="orders", session=session)  # deduped
+    assert count[0] == 1
+
+
 async def test_notify_still_skipped_for_future_dated(pg_engine, outbox_table) -> None:
     broker = OutboxBroker(pg_engine, outbox_table=outbox_table)
     session_factory = async_sessionmaker(pg_engine, expire_on_commit=False)
@@ -2538,6 +2549,57 @@ async def test_notify_still_skipped_for_future_dated(pg_engine, outbox_table) ->
     assert count[0] == 0  # future-dated -> no NOTIFY (unchanged), dedup did not break the skip
 
 
+async def test_publish_notifies_in_the_insert_round_trip(pg_engine: AsyncEngine, outbox_table: Table) -> None:
+    broker = OutboxBroker(pg_engine, outbox_table=outbox_table)
+    session_factory = async_sessionmaker(pg_engine, expire_on_commit=False)
+    statements: list[str] = []
+
+    def _listener(conn, cursor, statement, parameters, context, executemany) -> None:  # noqa: ARG001
+        statements.append(statement)
+
+    async with session_factory() as session, session.begin():
+        await session.connection()
+        event.listen(pg_engine.sync_engine, "before_cursor_execute", _listener)
+        try:
+            await broker.publish({"n": 1}, queue="orders", session=session)
+            await broker.publish({"n": 2}, queue="orders", session=session)
+        finally:
+            event.remove(pg_engine.sync_engine, "before_cursor_execute", _listener)
+
+    assert len(statements) == 2
+    assert "pg_notify" in statements[0].lower()
+    assert "pg_notify" not in statements[1].lower()
+
+
+async def test_notify_sent_after_a_timer_conflict_in_the_same_transaction(
+    pg_engine: AsyncEngine, outbox_table: Table
+) -> None:
+    broker = OutboxBroker(pg_engine, outbox_table=outbox_table)
+    session_factory = async_sessionmaker(pg_engine, expire_on_commit=False)
+    async with session_factory() as session, session.begin():
+        await broker.publish({"n": 0}, queue="orders", session=session, timer_id="t-1")
+
+    payloads: list[str] = []
+    received = asyncio.Event()
+    dsn = pg_engine.url.set(drivername="postgresql").render_as_string(hide_password=False)
+    listener = await asyncpg.connect(dsn)
+    try:
+
+        def _cb(_conn, _pid, _channel, payload) -> None:
+            payloads.append(payload)
+            received.set()
+
+        await listener.add_listener(f"outbox_{outbox_table.name}", _cb)
+        async with session_factory() as session, session.begin():
+            assert await broker.publish({"n": 1}, queue="orders", session=session, timer_id="t-1") is None
+            assert await broker.publish({"n": 2}, queue="orders", session=session) is not None
+        await asyncio.wait_for(received.wait(), timeout=2.0)
+    finally:
+        await listener.close()
+
+    assert payloads == ["orders"]
+
+
 async def test_notify_deduped_in_autobegin_transaction(pg_engine, outbox_table) -> None:
     # Autobegin (no explicit session.begin()): the dedup must still collapse N publishes
     # of one queue to one pg_notify. Guards against keying the memo on the GC-fragile async

tests/test_unit.py

@@ -612,21 +612,19 @@ async def test_broker_publish_rejects_non_async_session() -> None:
         await broker.publish(b"x", queue="orders", session=object())  # ty: ignore[invalid-argument-type]
 
 
-async def test_broker_publish_executes_insert_then_pg_notify_on_session() -> None:
+async def test_broker_publish_executes_one_insert_that_notifies() -> None:
     broker = _make_broker()
     session = _make_session_mock()
     await broker.publish({"order_id": 1}, queue="orders", session=session)
-    # Two execute calls: the INSERT, then SELECT pg_notify(...).
-    assert session.execute.await_count == 2
+    assert session.execute.await_count == 1
     insert_stmt, params = session.execute.await_args_list[0].args
-    assert "INSERT INTO" in str(insert_stmt)
+    compiled = insert_stmt.compile(dialect=postgresql.dialect())
+    assert "INSERT INTO" in str(compiled)
+    assert re.search(r"RETURNING outbox\.id, pg_notify\(%\(param_1\)s(::VARCHAR)?, outbox\.queue\)", str(compiled))
+    assert compiled.params["param_1"] == "outbox_outbox"
     assert params["queue"] == "orders"
     assert json.loads(params["payload"]) == {"order_id": 1}
     assert params["headers"]["content-type"] == "application/json"
-    notify_stmt, notify_params = session.execute.await_args_list[1].args
-    assert "pg_notify" in str(notify_stmt)
-    assert notify_params["channel"] == "outbox_outbox"
-    assert notify_params["payload"] == "orders"
 
 
 async def test_broker_publish_encodes_pydantic_model() -> None:
@@ -989,6 +987,7 @@ async def test_broker_publish_with_activate_in_skips_notify() -> None:
     insert_stmt = session.execute.await_args_list[0].args[0]
     assert "INSERT INTO" in str(insert_stmt)
     assert "next_attempt_at" in str(insert_stmt)
+    assert "pg_notify" not in str(insert_stmt)
 
 
 async def test_broker_publish_with_activate_at_skips_notify() -> None:
@@ -997,8 +996,9 @@ async def test_broker_publish_with_activate_at_skips_notify() -> None:
     fire = dt.datetime.now(tz=dt.UTC) + dt.timedelta(minutes=5)
     await broker.publish(b"x", queue="orders", session=session, activate_at=fire)
     assert session.execute.await_count == 1
-    _, params = session.execute.await_args_list[0].args
+    insert_stmt, params = session.execute.await_args_list[0].args
     assert params["next_attempt_at"] == fire
+    assert "pg_notify" not in str(insert_stmt)
 
 
 async def test_broker_publish_emits_notify_when_activate_at_is_past() -> None:
@@ -1008,19 +1008,18 @@ async def test_broker_publish_emits_notify_when_activate_at_is_past() -> None:
     session = _make_session_mock()
     past = dt.datetime.now(tz=dt.UTC) - dt.timedelta(seconds=5)
     await broker.publish(b"x", queue="orders", session=session, activate_at=past)
-    assert session.execute.await_count == 2
-    notify_stmt, notify_params = session.execute.await_args_list[1].args
-    assert "pg_notify" in str(notify_stmt)
-    assert notify_params["payload"] == "orders"
+    assert session.execute.await_count == 1
+    insert_stmt, _params = session.execute.await_args_list[0].args
+    assert "pg_notify" in str(insert_stmt)
 
 
 async def test_broker_publish_emits_notify_when_activate_in_is_zero() -> None:
     broker = _make_broker()
     session = _make_session_mock()
     await broker.publish(b"x", queue="orders", session=session, activate_in=dt.timedelta(0))
-    assert session.execute.await_count == 2
-    notify_stmt, _params = session.execute.await_args_list[1].args
-    assert "pg_notify" in str(notify_stmt)
+    assert session.execute.await_count == 1
+    insert_stmt, _params = session.execute.await_args_list[0].args
+    assert "pg_notify" in str(insert_stmt)
 
 
 async def test_broker_publish_rejects_naive_activate_at() -> None:
@@ -3669,7 +3668,7 @@ def _make_producer(recorder: typing.Any) -> OutboxProducer:
 async def test_producer_publish_emits_error_event_and_reraises_on_sql_failure() -> None:
     events, recorder = _events_recorder()
     producer = _make_producer(recorder)
-    session = AsyncMock(spec=AsyncSession)
+    session = _make_session_mock()
     session.execute.side_effect = RuntimeError("forced INSERT failure")
     cmd = OutboxPublishCommand({"x": 1}, queue="orders", session=session)