From af01910b2a410df39e4f39e35e7f1ed1ff0d09d2 Mon Sep 17 00:00:00 2001 From: Olivier 'reivilibre Date: Tue, 15 Sep 2026 11:40:59 +0100 Subject: [PATCH] Refactor the federation transmission code to delineate transaction preparation and completion. (#20166) A key refactoring for, and split out of, https://github.com/element-hq/synapse/pull/20165 Would be easier to land first to isolate the diff. Should be a standalone change with no behavioural change. Motivation is that #20165 will round-robin between 'main queue' transactions and 'sticky event' transactions. To keep the data flow clear, I wanted to insert a typed struct (well, `attrs` dataclass) as an interface between the 'preparation' of a transaction and its 'completion'. Doing this whilst keeping the asynchronous context manager style did not lead to a readable result in my opinion. (I would also say the async context manager is a touch 'magic' / obscures control flow, but I suspect this is largely down to opinion.) Replace _TransactionQueueManager with prepare/complete transaction methods --------- Signed-off-by: Olivier 'reivilibre --- changelog.d/20166.misc | 1 + .../sender/per_destination_queue.py | 282 +++++++++++------- 2 files changed, 178 insertions(+), 105 deletions(-) create mode 100644 changelog.d/20166.misc diff --git a/changelog.d/20166.misc b/changelog.d/20166.misc new file mode 100644 index 0000000000..716e910f8f --- /dev/null +++ b/changelog.d/20166.misc @@ -0,0 +1 @@ +Refactor the federation transmission code to delineate transaction preparation and completion. \ No newline at end of file diff --git a/synapse/federation/sender/per_destination_queue.py b/synapse/federation/sender/per_destination_queue.py index 32f8630c9d..4c0918e0ac 100644 --- a/synapse/federation/sender/per_destination_queue.py +++ b/synapse/federation/sender/per_destination_queue.py @@ -22,7 +22,6 @@ import datetime import logging from collections import OrderedDict -from types import TracebackType from typing import TYPE_CHECKING, Hashable, Iterable import attr @@ -79,6 +78,52 @@ CATCHUP_RETRY_INTERVAL = 60 * 60 * 1000 MAX_PRESENCE_STATES_PER_EDU = 50 +@attr.s(slots=True, auto_attribs=True, frozen=True) +class _PreparedTransaction: + """ + A transaction that has been prepared for sending: what to send, along with the + information that is useful for marking the transaction as complete once it has + been successfully sent. + + Produced by `PerDestinationQueue._prepare_transaction` and consumed by + `PerDestinationQueue._complete_transaction`. + """ + + pdus: list[EventBase] + """ + The PDUs to send in this transaction. + """ + + edus: list[Edu] + """ + The EDUs to send in this transaction. + """ + + to_device_message_stream_id: int | None + """ + This is the stream ID of the latest to-device message (`device_federation_outbox`) to be + sent (or None if none sent). + + When the transaction completes, to-device messages up to this point will be deleted from + the outbox. + """ + + device_list_stream_id: int | None + """ + This is the stream ID of the latest device list to be sent (or None if none sent). + + When the transaction completes, we will mark device lists up to this point as having been + sent. + """ + + last_stream_ordering: int | None + """ + This is the stream ordering of the last PDU that was sent (or None if none sent). + + When the transaction completes, this should be stored as our position in the events stream. + """ + + class PerDestinationQueue: """ Manages the per-destination transmission queues. @@ -343,7 +388,7 @@ class PerDestinationQueue: ) async def _transaction_transmission_loop(self) -> None: - pending_pdus: list[EventBase] = [] + transaction: _PreparedTransaction | None = None try: self.transmission_loop_running = True # This will throw if we wouldn't retry. We do this here so we fail @@ -366,46 +411,48 @@ class PerDestinationQueue: while self._transmission_loop_enabled: self._new_data_to_send = False + # Clear `transaction` (as it's used in the exception handler + # for error reporting) + transaction = None + transaction = await self._prepare_transaction() - async with _TransactionQueueManager(self) as ( - pending_pdus, # noqa: F811 - pending_edus, - ): - if not pending_pdus and not pending_edus: - logger.debug("TX [%s] Nothing to send", self._destination) + if transaction is None: + logger.debug("TX [%s] Nothing to send", self._destination) - # If we've gotten told about new things to send during - # checking for things to send, we try looking again. - # Otherwise new PDUs or EDUs might arrive in the meantime, - # but not get sent because we currently have an - # `_active_transmission_loop` running. - if self._new_data_to_send: - continue - else: - return + # If we've gotten told about new things to send during + # checking for things to send, we try looking again. + # Otherwise new PDUs or EDUs might arrive in the meantime, + # but not get sent because we currently have an + # `_active_transmission_loop` running. + if self._new_data_to_send: + continue + else: + return - if pending_pdus: - logger.debug( - "TX [%s] len(pending_pdus_by_dest[dest]) = %d", - self._destination, - len(pending_pdus), - ) - - await self._transaction_manager.send_new_transaction( - self._destination, pending_pdus, pending_edus + if transaction.pdus: + logger.debug( + "TX [%s] len(pending_pdus_by_dest[dest]) = %d", + self._destination, + len(transaction.pdus), ) - sent_transactions_counter.labels( - **{SERVER_NAME_LABEL: self.server_name} + await self._transaction_manager.send_new_transaction( + self._destination, transaction.pdus, transaction.edus + ) + + sent_transactions_counter.labels( + **{SERVER_NAME_LABEL: self.server_name} + ).inc() + sent_edus_counter.labels(**{SERVER_NAME_LABEL: self.server_name}).inc( + len(transaction.edus) + ) + for edu in transaction.edus: + sent_edus_by_type.labels( + type=edu.edu_type, + **{SERVER_NAME_LABEL: self.server_name}, ).inc() - sent_edus_counter.labels( - **{SERVER_NAME_LABEL: self.server_name} - ).inc(len(pending_edus)) - for edu in pending_edus: - sent_edus_by_type.labels( - type=edu.edu_type, - **{SERVER_NAME_LABEL: self.server_name}, - ).inc() + + await self._complete_transaction(transaction) except NotRetryingDestination as e: logger.debug( @@ -455,16 +502,19 @@ class PerDestinationQueue: "TX [%s] Failed to send transaction: %s", self._destination, e ) - for p in pending_pdus: - logger.info( - "Failed to send event %s to %s", p.event_id, self._destination - ) + if transaction is not None: + for p in transaction.pdus: + logger.info( + "Failed to send event %s to %s", p.event_id, self._destination + ) except Exception: logger.exception("TX [%s] Failed to send transaction", self._destination) - for p in pending_pdus: - logger.info( - "Failed to send event %s to %s", p.event_id, self._destination - ) + + if transaction is not None: + for p in transaction.pdus: + logger.info( + "Failed to send event %s to %s", p.event_id, self._destination + ) finally: # We want to be *very* sure we clear this after we stop processing self.active_transmission_loop = None @@ -736,21 +786,31 @@ class PerDestinationQueue: self._catching_up = True self._pending_pdus = [] + async def _prepare_transaction(self) -> _PreparedTransaction | None: + """ + Work out what should go in the next transaction to this destination, by + calculating what we want to send and preparing the information that is + useful once we have completed the transaction. -@attr.s(slots=True, auto_attribs=True) -class _TransactionQueueManager: - """A helper async context manager for pulling stuff off the queues and - tracking what was last successfully sent, etc. - """ + Side effects: + - Dequeues pending EDUs + - `_pending_presence` + - `_pending_receipt_edus` + - `_pending_edus` (currently unused in practice) + - `_pending_edus_keyed` + - Advances our devices stream positions (only if there is nothing to send + in that range, so we don't lose any updates) - queue: PerDestinationQueue + PDUs are not dequeued until acknowledged by `_complete_transaction`. - _device_stream_id: int | None = None - _device_list_id: int | None = None - _last_stream_ordering: int | None = None - _pdus: list[EventBase] = attr.Factory(list) + Returns: + - the prepared transaction; or + - None if there is nothing to send and no progress to record + + Once the prepared transaction has been sent successfully, + `_complete_transaction` must be called with it. + """ - async def __aenter__(self) -> tuple[list[EventBase], list[Edu]]: # First we calculate the EDUs we want to send, if any. # There's a maximum number of EDUs that can be sent with a transaction, @@ -767,30 +827,30 @@ class _TransactionQueueManager: pending_edus = [] # Add presence EDU. - if self.queue._pending_presence: + if self._pending_presence: # Only send max 50 presence entries in the EDU, to bound the amount # of data we're sending. presence_to_add: list[JsonDict] = [] while ( - self.queue._pending_presence + self._pending_presence and len(presence_to_add) < MAX_PRESENCE_STATES_PER_EDU ): - _, presence = self.queue._pending_presence.popitem(last=False) + _, presence = self._pending_presence.popitem(last=False) presence_to_add.append( - format_user_presence_state(presence, self.queue._clock.time_msec()) + format_user_presence_state(presence, self._clock.time_msec()) ) pending_edus.append( Edu( - origin=self.queue.server_name, - destination=self.queue._destination, + origin=self.server_name, + destination=self._destination, edu_type=EduTypes.PRESENCE, content={"push": presence_to_add}, ) ) # Add read receipt EDUs. - pending_edus.extend(self.queue._get_receipt_edus(limit=5)) + pending_edus.extend(self._get_receipt_edus(limit=5)) edu_limit = MAX_EDUS_PER_TRANSACTION - len(pending_edus) # Next, prioritize to-device messages so that existing encryption channels @@ -799,91 +859,103 @@ class _TransactionQueueManager: ( to_device_edus, device_stream_id, - ) = await self.queue._get_to_device_message_edus( + ) = await self._get_to_device_message_edus( edu_limit - NUMBER_OF_RESERVED_EDUS_PER_TRANSACTION ) + device_stream_id_upon_completion: int | None = None if to_device_edus: - self._device_stream_id = device_stream_id + # We can advance our position in the device stream after the transaction completes. + device_stream_id_upon_completion = device_stream_id else: - self.queue._last_device_stream_id = device_stream_id + # We can advance our position in the device stream immediately, as there's nothing to send. + self._last_device_stream_id = device_stream_id pending_edus.extend(to_device_edus) edu_limit -= len(to_device_edus) # Add device list update EDUs. - device_update_edus, dev_list_id = await self.queue._get_device_update_edus( - edu_limit - ) + device_update_edus, dev_list_id = await self._get_device_update_edus(edu_limit) + device_list_id_upon_completion: int | None = None if device_update_edus: - self._device_list_id = dev_list_id + # We can advance our position in the device list stream after the transaction completes. + device_list_id_upon_completion = dev_list_id else: - self.queue._last_device_list_stream_id = dev_list_id + # We can advance our position in the device list stream immediately, as there's nothing to send. + self._last_device_list_stream_id = dev_list_id pending_edus.extend(device_update_edus) edu_limit -= len(device_update_edus) # Finally add any other types of EDUs if there is room. - other_edus = self.queue._pop_pending_edus(edu_limit) + other_edus = self._pop_pending_edus(edu_limit) pending_edus.extend(other_edus) edu_limit -= len(other_edus) - while edu_limit > 0 and self.queue._pending_edus_keyed: - _, val = self.queue._pending_edus_keyed.popitem() + while edu_limit > 0 and self._pending_edus_keyed: + _, val = self._pending_edus_keyed.popitem() pending_edus.append(val) edu_limit -= 1 # Now we look for any PDUs to send, by getting up to 50 PDUs from the # queue - self._pdus = self.queue._pending_pdus[:50] + pdus = self._pending_pdus[:50] - if not self._pdus and not pending_edus: - return [], [] + if not pdus and not pending_edus: + # There is nothing to send. There's also nothing to record upon + # completion: the only progress we could have made without sending + # anything is advancing our positions in the device streams, and that + # has already been done above. + return None - if self._pdus: - self._last_stream_ordering = self._pdus[ - -1 - ].internal_metadata.stream_ordering - assert self._last_stream_ordering + last_stream_ordering: int | None = None + if pdus: + last_stream_ordering = pdus[-1].internal_metadata.stream_ordering + assert last_stream_ordering - return self._pdus, pending_edus + return _PreparedTransaction( + pdus=pdus, + edus=pending_edus, + to_device_message_stream_id=device_stream_id_upon_completion, + device_list_stream_id=device_list_id_upon_completion, + last_stream_ordering=last_stream_ordering, + ) - async def __aexit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - if exc_type is not None: - # Failed to send transaction, so we bail out. - return + async def _complete_transaction(self, transaction: _PreparedTransaction) -> None: + """ + Handle the fact that a transaction has been successfully completed. + Must not be called if sending the transaction failed, as it records how far + through the various streams we have now got. + """ # Successfully sent transactions, so we remove pending PDUs from the queue - if self._pdus: - self.queue._pending_pdus = self.queue._pending_pdus[len(self._pdus) :] + if transaction.pdus: + self._pending_pdus = self._pending_pdus[len(transaction.pdus) :] # Succeeded to send the transaction so we record where we have sent up # to in the various streams - if self._device_stream_id: - await self.queue._store.delete_device_msgs_for_remote( - self.queue._destination, self._device_stream_id + if transaction.to_device_message_stream_id: + await self._store.delete_device_msgs_for_remote( + self._destination, transaction.to_device_message_stream_id ) - self.queue._last_device_stream_id = self._device_stream_id + self._last_device_stream_id = transaction.to_device_message_stream_id # also mark the device updates as sent - if self._device_list_id: + if transaction.device_list_stream_id: logger.info( - "Marking as sent %r %r", self.queue._destination, self._device_list_id + "Marking as sent %r %r", + self._destination, + transaction.device_list_stream_id, ) - await self.queue._store.mark_as_sent_devices_by_remote( - self.queue._destination, self._device_list_id + await self._store.mark_as_sent_devices_by_remote( + self._destination, transaction.device_list_stream_id ) - self.queue._last_device_list_stream_id = self._device_list_id + self._last_device_list_stream_id = transaction.device_list_stream_id - if self._last_stream_ordering: + if transaction.last_stream_ordering: # we sent some PDUs and it was successful, so update our # last_successful_stream_ordering in the destinations table. - await self.queue._store.set_destination_last_successful_stream_ordering( - self.queue._destination, self._last_stream_ordering + await self._store.set_destination_last_successful_stream_ordering( + self._destination, transaction.last_stream_ordering )