From aab63b7024280d3e180ba7df05eb775484625dd4 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Wed, 28 Oct 2026 17:15:45 +0100 Subject: [PATCH 1/3] refactor: read the owner inbox through p2p receive and mark_processed The DO reads proposed changes with P2PConnection.receive(peer, kind, max_messages) and finishes each with mark_processed (renamed from ack). The Drive-only inbox methods are removed. Drive does the same calls as before: one listing, one download and one archive per message. --- syft/sync/connections/base_connection.py | 7 +- syft/sync/connections/connection_router.py | 44 ++++++------ .../connections/drive/drive_connections.py | 67 +++++++++--------- .../connections/drive/gdrive_transport.py | 68 +++---------------- syft/sync/sync/datasite_owner_syncer.py | 14 ++-- syft/sync/transport/protocols.py | 34 ++++++---- syft/sync/transport/types.py | 2 +- .../p2p/test_message_protocol_downgrade.py | 9 +-- tests/unit/test_drive_connections.py | 57 ++++++++++------ tests/unit/test_encryption_policy.py | 7 +- tests/unit/test_sender_binding.py | 4 +- 11 files changed, 146 insertions(+), 167 deletions(-) diff --git a/syft/sync/connections/base_connection.py b/syft/sync/connections/base_connection.py index 54c5cc9eec1..e81c43fc18f 100644 --- a/syft/sync/connections/base_connection.py +++ b/syft/sync/connections/base_connection.py @@ -117,10 +117,9 @@ def watcher_send_raw_bytes_to_inbox( ) -> None: raise NotImplementedError() - def owner_download_next_raw_proposed_message_from_inbox( - self, sender_email: str - ) -> tuple[bytes, str] | None: - """Download next message from inbox as raw bytes. Returns (data, file_id) or None.""" + def owner_get_inbox_file_metadatas(self, sender_email: str) -> list[dict]: + """The messages in the inbox from ``sender_email``, oldest first, each + with ``file_id`` and ``file_name``.""" raise NotImplementedError() def owner_write_raw_bytes_to_outbox( diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index ebd18bb4ed2..5c0d24a171d 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -29,7 +29,7 @@ from syft.sync.peers.peer_store import PeerStore from syft.sync.transport.naming import PEERS_META_KEY from syft.sync.transport.protocols import Transport -from syft.sync.transport.types import MessageKind, SendStatus +from syft.sync.transport.types import MessageKind, RawMessage, SendStatus from syft.sync.utils.print_utils import ( print_peer_added_to_platform, print_peer_adding_to_platform, @@ -151,33 +151,38 @@ def watcher_send_proposed_file_changes_message( def owner_get_next_proposed_filechange_message( self, sender_email: str - ) -> ProposedFileChangesMessage | None: - """The next message from ``sender_email``, or None when the inbox is empty. + ) -> tuple[RawMessage, ProposedFileChangesMessage] | None: + """The next message from ``sender_email`` as received and as read, or + None when nothing is waiting. Pass the received one to + ``owner_mark_processed`` once it is handled. - A message whose format does not match the pair is archived unread, - with a warning. + A message whose format does not match the pair is marked processed + unread, with a warning. Raises: SyftPeerNotReadyError: when the message cannot be read yet. It - stays in the inbox. + stays unprocessed. """ p2p = self.transport.p2p while True: - result = p2p.next_inbox_message(sender_email) - if result is None: + received = p2p.receive( + sender_email, MessageKind.PROPOSED_CHANGES, max_messages=1 + ) + if not received: return None - raw_data, file_id = result + [raw] = received try: decrypted = self.peer_store.decrypt_and_verify_if_needed( - sender_email, raw_data + sender_email, raw.data ) except MessageFormatError as e: - warnings.warn(f"Archived message file {file_id} unread: {e}") - p2p.archive_inbox_message(sender_email, file_id) + warnings.warn(f"Archived message {raw.msg_id} unread: {e}") + p2p.mark_processed(raw) continue - msg = ProposedFileChangesMessage.from_compressed_data(decrypted) - msg.platform_id = file_id - return msg + return raw, ProposedFileChangesMessage.from_compressed_data(decrypted) + + def owner_mark_processed(self, message: RawMessage) -> None: + self.transport.p2p.mark_processed(message) def owner_write_event_messages_to_outbox( self, recipient_email: str, events_message: FileChangeEventsMessage @@ -321,15 +326,6 @@ def update_peer_state( peer_email, state, public_encryption_bundle ) - def owner_remove_proposed_filechange_from_inbox( - self, - proposed_filechange_message: ProposedFileChangesMessage, - sender_email: str, - ): - self.transport.p2p.archive_proposed_changes( - proposed_filechange_message, sender_email - ) - # ========================================================================= # ENCRYPTION BUNDLE EXCHANGE # ========================================================================= diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 83d700c80e2..3e0c4abd16b 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -26,10 +26,18 @@ if TYPE_CHECKING: from syft.sync.connections.base_connection import FileCollection from syft.sync.connections.drive.gdrive_transport import GDriveConnection - from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage from syft.sync.version.peer_info import PeerInfo +def _refuse_outbox_kind(kind: MessageKind) -> None: + """Only the inbox is read through ``receive`` so far.""" + if kind is not MessageKind.PROPOSED_CHANGES: + raise NotImplementedError( + f"Messages of kind {kind!r} are still read through the outbox " + "methods; they move to receive with send-and-forget." + ) + + class DriveP2PConnection(P2PConnection): """Messages through the shared inbox and outbox folders. @@ -37,11 +45,11 @@ class DriveP2PConnection(P2PConnection): data scientist's outbox. ``msg_id`` is the file name. A written file is ``DELIVERED``: it waits in the folder until the peer reads it. - Reading and archiving go through the Drive-only methods below, which keep - how Drive does it today: the inbox is read one message at a time and - archived, the outbox is read from a timestamp cursor and kept. ``receive`` - and ``ack`` replace them with send-and-forget, which changes that behaviour - (one rule for both folders, delete on ack, no read cursor). + The inbox is read through ``receive`` and a processed message is archived. + The outbox is still read through the Drive-only methods below, from a + timestamp cursor, and kept: the data scientist rebuilds a lost cache by + reading it again. It moves to ``receive`` with send-and-forget, once the + data scientist keeps what it received in its own storage. """ def __init__(self, connection: GDriveConnection): @@ -62,32 +70,27 @@ def delivered(self) -> list[str]: # Every Drive send is DELIVERED when it returns; none waits for a peer. return [] - def receive(self, max_messages: int | None = None) -> list[RawMessage]: - raise NotImplementedError( - "Messages are still read through the Drive methods; receiving " - "through P2PConnection lands with send-and-forget." - ) - - def ack(self, message: RawMessage) -> None: - raise NotImplementedError( - "Messages are still archived through the Drive methods; " - "acknowledging through P2PConnection lands with send-and-forget." - ) - - def next_inbox_message(self, sender: str) -> tuple[bytes, str] | None: - """The oldest message from ``sender`` in the inbox and its handle, or - None when the inbox is empty.""" - return self._conn.owner_download_next_raw_proposed_message_from_inbox(sender) - - def archive_inbox_message(self, sender: str, handle: str) -> None: - """Move the inbox message ``handle`` from ``sender`` out of the inbox.""" - self._conn.owner_archive_inbox_file(handle, sender) - - def archive_proposed_changes( - self, message: ProposedFileChangesMessage, sender: str - ) -> None: - """Move the handled ``message`` from ``sender`` out of the inbox.""" - self._conn.owner_remove_proposed_filechange_message_from_inbox(message, sender) + def receive( + self, peer: str, kind: MessageKind, max_messages: int | None = None + ) -> list[RawMessage]: + _refuse_outbox_kind(kind) + listed = self._conn.owner_get_inbox_file_metadatas(peer)[:max_messages] + return [ + RawMessage( + sender=peer, + msg_id=m["file_name"], + kind=kind, + data=self._conn.download_file(m["file_id"]), + handle=m["file_id"], + ) + for m in listed + ] + + def mark_processed(self, message: RawMessage) -> None: + # The archive folder is named after ``message.sender``, which receive + # set from the inbox folder, never from the message body. + _refuse_outbox_kind(message.kind) + self._conn.owner_archive_inbox_file(message.handle, message.sender) def outbox_messages(self, peer: str, since_timestamp: float | None) -> list[bytes]: """The messages in ``peer``'s outbox to me, newer than ``since_timestamp``.""" diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 25c90276081..8d48ec70ff0 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -833,31 +833,12 @@ def owner_write_event_messages_to_outbox( fname = events_message.message_filepath.as_string() self.owner_write_raw_bytes_to_outbox(recipient, fname, data) - def owner_remove_proposed_filechange_message_from_inbox( - self, - proposed_filechange_message: ProposedFileChangesMessage, - sender_email: str, - ): - """Archive a handled message. ``sender_email`` is the peer it came from. + def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: + """Move the inbox file ``file_id`` from ``sender_email`` to its archive folder. - The caller reads that email from the transport. The ``sender_email`` - field of the message is sender-supplied, and it names the archive folder, - so this method never reads it. + ``sender_email`` names the archive folder, so the caller reads it from + the transport, never from the sender-supplied message body. """ - fname = proposed_filechange_message.message_filename.as_string() - - # Use cached platform_id if available, otherwise fall back to name-based lookup - gdrive_id = proposed_filechange_message.platform_id - if gdrive_id is None: - gdrive_id = self.get_inbox_proposed_event_id_from_name(sender_email, fname) - if gdrive_id is None: - raise ValueError( - f"Event {fname} not found in inbox, event should already be created for this type of connection" - ) - self.owner_archive_inbox_file(gdrive_id, sender_email) - - def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: - """Move the inbox file ``file_id`` from ``sender_email`` to its archive folder.""" file_info = execute_with_retries( self.drive_service.files().get(fileId=file_id, fields="parents") ) @@ -1097,35 +1078,19 @@ def _get_valid_messages_from_file_metadatas( continue return res - def owner_download_next_raw_proposed_message_from_inbox( - self, sender_email: str - ) -> tuple[bytes, str] | None: + def owner_get_inbox_file_metadatas(self, sender_email: str) -> list[dict]: + """The messages in the inbox from ``sender_email``, oldest first, each + with ``file_id`` and ``file_name``, without downloading them.""" inbox_folder_id = self._get_own_datasite_inbox_id(sender_email) if inbox_folder_id is None: raise ValueError(f"Inbox folder for {sender_email} not found") file_metadatas = self.get_file_metadatas_from_folder(inbox_folder_id) valid_file_names = self._get_valid_messages_from_file_metadatas(file_metadatas) - if len(valid_file_names) == 0: - return None - first_file_name = sorted(valid_file_names, key=lambda x: x.submitted_timestamp)[ - 0 + name_to_id = {f["name"]: f["id"] for f in file_metadatas} + return [ + {"file_id": name_to_id[name.as_string()], "file_name": name.as_string()} + for name in sorted(valid_file_names, key=lambda x: x.submitted_timestamp) ] - first_file_id = [ - x for x in file_metadatas if x["name"] == first_file_name.as_string() - ][0]["id"] - raw_data = self.download_file(first_file_id) - return raw_data, first_file_id - - def owner_get_next_proposed_filechange_message( - self, sender_email: str - ) -> ProposedFileChangesMessage | None: - result = self.owner_download_next_raw_proposed_message_from_inbox(sender_email) - if result is None: - return None - raw_data, file_id = result - msg = ProposedFileChangesMessage.from_compressed_data(raw_data) - msg.platform_id = file_id - return msg def _find_p2p_folder_id( self, @@ -1558,17 +1523,6 @@ def get_syftbox_folder_id_from_drive(self) -> str | None: items = results.get("files", []) return items[0]["id"] if items else None - def get_inbox_proposed_event_id_from_name( - self, sender_email: str, name: str - ) -> str | None: - inbox_folder_id = self._get_own_datasite_inbox_id(sender_email) - query = f"name='{name}' and '{inbox_folder_id}' in parents and trashed=false" - results = execute_with_retries( - self.drive_service.files().list(q=query, fields="files(id, name)") - ) - items = results.get("files", []) - return items[0]["id"] if items else None - # ========================================================================= # GENERIC COLLECTION METHODS (parameterized on `prefix`) # diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 8e01312751b..63141593c64 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -764,10 +764,13 @@ def pull_and_process_next_proposed_filechange( self, sender_email: str, raise_on_none=True ) -> ProposedFileChangesMessage | None: # raise on none is useful for testing, shouldnt be used in production - message = self.connection_router.owner_get_next_proposed_filechange_message( - sender_email=sender_email + next_message = ( + self.connection_router.owner_get_next_proposed_filechange_message( + sender_email=sender_email + ) ) - if message is not None: + if next_message is not None: + received, message = next_message # The sender is the peer whose inbox held the message. The transport # read the message from that peer's folder and checked the signature # against that peer's key. The body is sender-supplied, so a claim @@ -780,10 +783,7 @@ def pull_and_process_next_proposed_filechange( f"claims sender {message.sender_email!r}." ) - # delete the message once we are done - self.connection_router.owner_remove_proposed_filechange_from_inbox( - message, sender_email - ) + self.connection_router.owner_mark_processed(received) return message elif raise_on_none: raise ValueError("No proposed file change to process") diff --git a/syft/sync/transport/protocols.py b/syft/sync/transport/protocols.py index 8423e179fb0..d319ab04455 100644 --- a/syft/sync/transport/protocols.py +++ b/syft/sync/transport/protocols.py @@ -27,13 +27,18 @@ class P2PConnection(ABC): - """Messages between peers, kept until the recipient acknowledges them. - - One rule for every message, whatever its kind: ``receive`` returns the - messages not acknowledged yet, and ``ack`` finishes one. On Drive a written - message is ``DELIVERED`` and ``ack`` removes it. On a live link such as - WebRTC a message is ``QUEUED`` until the peer acknowledges it, which - ``delivered()`` then reports. + """Messages between peers, kept until the recipient has processed them. + + Sending and receiving each have their own confirmation: + + - Sender: a message is delivered once it is stored where the peer will + find it. On Drive that is when the upload returns, so ``send`` returns + ``DELIVERED``. On a live link such as WebRTC ``send`` returns ``QUEUED``, + and ``delivered()`` reports the message once the peer has stored it and + sent a receipt back. + - Receiver: ``receive`` returns the messages not processed yet, and + ``mark_processed`` finishes one. This stays with the receiver: nothing + is sent back to the peer. """ @abstractmethod @@ -49,21 +54,22 @@ def send( @abstractmethod def delivered(self) -> list[str]: - """The ``msg_id`` of each queued message the peer acknowledged since - the last call. A connection that only returns ``DELIVERED`` returns - nothing.""" + """The ``msg_id`` of each queued message the peer stored since the last + call. A connection that only returns ``DELIVERED`` returns nothing.""" @abstractmethod - def receive(self, max_messages: int | None = None) -> list[RawMessage]: - """Messages from every peer that are not acknowledged yet, oldest first - per sender, at most ``max_messages`` of them. + def receive( + self, peer: str, kind: MessageKind, max_messages: int | None = None + ) -> list[RawMessage]: + """Messages of ``kind`` from ``peer`` that are not processed yet, oldest + first, at most ``max_messages`` of them. A peer whose channel does not exist yet has no messages. The connection sets each ``RawMessage.sender`` from where the message came from. """ @abstractmethod - def ack(self, message: RawMessage) -> None: + def mark_processed(self, message: RawMessage) -> None: """Record that ``message`` was processed; ``receive`` never returns it again.""" diff --git a/syft/sync/transport/types.py b/syft/sync/transport/types.py index 1f9cffa0d36..a3f17526867 100644 --- a/syft/sync/transport/types.py +++ b/syft/sync/transport/types.py @@ -28,7 +28,7 @@ class RawMessage(BaseModel): ``sender`` is set by the connection from where the message came from, never from its contents. ``handle`` is opaque: only the connection that - returned the message reads it, to acknowledge it. + returned the message reads it, to mark it processed. """ model_config = ConfigDict(frozen=True) diff --git a/tests/migrations/p2p/test_message_protocol_downgrade.py b/tests/migrations/p2p/test_message_protocol_downgrade.py index 864cedfcdfb..97febf9f831 100644 --- a/tests/migrations/p2p/test_message_protocol_downgrade.py +++ b/tests/migrations/p2p/test_message_protocol_downgrade.py @@ -24,6 +24,7 @@ from syft.sync.events.file_change_event import FileChangeEventsMessageV1 from syft.sync.messages.proposed_filechange import ProposedFileChangesMessageV1 from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.transport import MessageKind from syft.sync.utils.syftbox_utils import uncompress_data from syft_migration import MigrationError, ProtocolSchema @@ -119,10 +120,10 @@ def _wire_version(manager, peer_email: str, raw: bytes) -> str: def _raw_proposal_version(do_manager, ds_email: str) -> str: """The version field of the next proposal blob in the DO's inbox, unparsed.""" - raw, _ = do_manager._connection_router.connections[ - 0 - ].owner_download_next_raw_proposed_message_from_inbox(ds_email) - return _wire_version(do_manager, ds_email, raw) + [message] = do_manager._connection_router.transport.p2p.receive( + ds_email, MessageKind.PROPOSED_CHANGES, max_messages=1 + ) + return _wire_version(do_manager, ds_email, message.data) def _raw_outbox_versions(ds_manager, do_email: str) -> list[str]: diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index a91874f2b2d..1ec0981713b 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -90,20 +90,35 @@ def test_a_proposed_change_sent_through_p2p_reaches_the_owner_inbox(pair): DO, _proposal_name(1.0), b"job", MessageKind.PROPOSED_CHANGES ) assert status is SendStatus.DELIVERED - data, _ = do_conn.as_transport().p2p.next_inbox_message(DS) - assert data == b"job" + [message] = do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + assert (message.sender, message.kind, message.data) == ( + DS, + MessageKind.PROPOSED_CHANGES, + b"job", + ) + + +def test_inbox_messages_are_received_oldest_first_up_to_max(pair): + do_conn, ds_conn = pair + names = {submitted: _proposal_name(submitted) for submitted in (2.0, 1.0, 3.0)} + for name in names.values(): + ds_conn.as_transport().p2p.send(DO, name, b"job", MessageKind.PROPOSED_CHANGES) + p2p = do_conn.as_transport().p2p + + received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES, max_messages=2) + assert [m.msg_id for m in received] == [names[1.0], names[2.0]] -def test_an_archived_inbox_message_is_not_read_again(pair): +def test_processed_inbox_message_is_not_received_again(pair): do_conn, ds_conn = pair ds_conn.as_transport().p2p.send( DO, _proposal_name(1.0), b"job", MessageKind.PROPOSED_CHANGES ) p2p = do_conn.as_transport().p2p - _, handle = p2p.next_inbox_message(DS) - p2p.archive_inbox_message(DS, handle) - assert p2p.next_inbox_message(DS) is None + [message] = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + p2p.mark_processed(message) + assert p2p.receive(DS, MessageKind.PROPOSED_CHANGES) == [] def test_file_events_sent_through_p2p_reach_the_scientist_outbox(pair): @@ -144,16 +159,14 @@ def test_drive_never_has_sends_waiting_for_confirmation(pair): assert do_conn.as_transport().p2p.delivered() == [] -def test_receiving_through_p2p_is_not_built_yet(pair): - do_conn, _ = pair - p2p = do_conn.as_transport().p2p - message = RawMessage( - sender=DS, msg_id="m", kind=MessageKind.PROPOSED_CHANGES, data=b"" - ) +def test_receiving_file_events_through_p2p_is_not_built_yet(pair): + _, ds_conn = pair + p2p = ds_conn.as_transport().p2p + message = RawMessage(sender=DO, msg_id="m", kind=MessageKind.FILE_EVENTS, data=b"") with pytest.raises(NotImplementedError): - p2p.receive() + p2p.receive(DO, MessageKind.FILE_EVENTS) with pytest.raises(NotImplementedError): - p2p.ack(message) + p2p.mark_processed(message) # -- OwnerStore ------------------------------------------------------------------- @@ -372,16 +385,20 @@ def test_the_router_sends_proposed_changes_through_p2p(): def test_the_router_reads_the_inbox_through_p2p(): router, p2p = _router_with_mock_p2p() - p2p.next_inbox_message.return_value = None + p2p.receive.return_value = [] assert router.owner_get_next_proposed_filechange_message(DS) is None - p2p.next_inbox_message.assert_called_once_with(DS) + p2p.receive.assert_called_once_with( + DS, MessageKind.PROPOSED_CHANGES, max_messages=1 + ) -def test_the_router_archives_a_handled_message_through_p2p(): +def test_the_router_marks_a_handled_message_processed_through_p2p(): router, p2p = _router_with_mock_p2p() - message = ProposedFileChangesMessage(sender_email=DS, proposed_file_changes=[]) - router.owner_remove_proposed_filechange_from_inbox(message, DS) - p2p.archive_proposed_changes.assert_called_once_with(message, DS) + message = RawMessage( + sender=DS, msg_id="m", kind=MessageKind.PROPOSED_CHANGES, data=b"" + ) + router.owner_mark_processed(message) + p2p.mark_processed.assert_called_once_with(message) def test_the_router_reads_the_outbox_through_p2p(): diff --git a/tests/unit/test_encryption_policy.py b/tests/unit/test_encryption_policy.py index 741ab524664..978f8128bff 100644 --- a/tests/unit/test_encryption_policy.py +++ b/tests/unit/test_encryption_policy.py @@ -28,6 +28,7 @@ ) from syft.sync.syftbox_manager import SyftboxManager from syft.sync.sync.datasite_owner_syncer import OutgoingMessage +from syft.sync.transport import MessageKind from syft.sync.version.peer_info import PeerInfo from syft.sync.version.peer_manager import PeerManager, PeerManagerConfig from syft.sync.version.version_info import VersionInfo @@ -74,8 +75,10 @@ def _keyed_pair( def _raw_inbox_message(do_manager, sender_email: str) -> bytes: router = do_manager.datasite_owner_syncer.connection_router - raw, _ = router.transport.p2p.next_inbox_message(sender_email) - return raw + [message] = router.transport.p2p.receive( + sender_email, MessageKind.PROPOSED_CHANGES, max_messages=1 + ) + return message.data def _accepted_contents(do_manager) -> list[str]: diff --git a/tests/unit/test_sender_binding.py b/tests/unit/test_sender_binding.py index 883d7f4d5e3..c1a8a968b89 100644 --- a/tests/unit/test_sender_binding.py +++ b/tests/unit/test_sender_binding.py @@ -69,8 +69,8 @@ def test_sender_email_in_body_cannot_claim_the_datasite_owner(): def test_spoofed_sender_does_not_name_an_archive_folder(): """The DO archives a handled message under the peer that sent it. - `owner_remove_proposed_filechange_message_from_inbox` builds the archive - folder name from an email, and creates the folder. Read that email from the + `owner_archive_inbox_file` builds the archive folder name from an email, + and creates the folder. Read that email from the transport, so a claim in the body cannot name a folder in the DO's Drive. """ ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() From b73bc4c3183c3c7a39b55c3e86841af80a971bb0 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Wed, 28 Oct 2026 17:41:44 +0100 Subject: [PATCH 2/3] perf: list the owner inbox once and download its messages in parallel The DO received one inbox message per listing and looked up each file's parents before archiving it: 4N+1 Drive calls for N messages. It now lists once, downloads in parallel with one connection copy per thread, and archives straight out of the cached inbox folder: 2N+1. A newly created archive folder is now cached, and a mock thread copy no longer warms folder ids the real copy skips. --- syft/sync/connections/connection_router.py | 41 +++---- .../connections/drive/drive_connections.py | 35 +++++- .../connections/drive/gdrive_transport.py | 47 ++++---- syft/sync/sync/datasite_owner_syncer.py | 58 +++++---- tests/unit/test_drive_connections.py | 111 ++++++++++++++++-- 5 files changed, 205 insertions(+), 87 deletions(-) diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 5c0d24a171d..8ee0873b267 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -149,37 +149,32 @@ def watcher_send_proposed_file_changes_message( filename = proposed_file_changes_message.message_filename.as_string() self.transport.p2p.send(recipient, filename, data, MessageKind.PROPOSED_CHANGES) - def owner_get_next_proposed_filechange_message( - self, sender_email: str - ) -> tuple[RawMessage, ProposedFileChangesMessage] | None: - """The next message from ``sender_email`` as received and as read, or - None when nothing is waiting. Pass the received one to - ``owner_mark_processed`` once it is handled. + def owner_receive_proposed_messages(self, sender_email: str) -> list[RawMessage]: + """Every proposed-changes message from ``sender_email`` not processed + yet, oldest first, still encrypted.""" + return self.transport.p2p.receive(sender_email, MessageKind.PROPOSED_CHANGES) + + def owner_read_proposed_message( + self, message: RawMessage + ) -> ProposedFileChangesMessage | None: + """``message`` decrypted, verified and decoded. A message whose format does not match the pair is marked processed - unread, with a warning. + unread, with a warning, and gives None. Raises: SyftPeerNotReadyError: when the message cannot be read yet. It stays unprocessed. """ - p2p = self.transport.p2p - while True: - received = p2p.receive( - sender_email, MessageKind.PROPOSED_CHANGES, max_messages=1 + try: + decrypted = self.peer_store.decrypt_and_verify_if_needed( + message.sender, message.data ) - if not received: - return None - [raw] = received - try: - decrypted = self.peer_store.decrypt_and_verify_if_needed( - sender_email, raw.data - ) - except MessageFormatError as e: - warnings.warn(f"Archived message {raw.msg_id} unread: {e}") - p2p.mark_processed(raw) - continue - return raw, ProposedFileChangesMessage.from_compressed_data(decrypted) + except MessageFormatError as e: + warnings.warn(f"Archived message {message.msg_id} unread: {e}") + self.owner_mark_processed(message) + return None + return ProposedFileChangesMessage.from_compressed_data(decrypted) def owner_mark_processed(self, message: RawMessage) -> None: self.transport.p2p.mark_processed(message) diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 3e0c4abd16b..5d05f88acf8 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -10,6 +10,8 @@ from __future__ import annotations +import threading +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform @@ -29,7 +31,11 @@ from syft.sync.version.peer_info import PeerInfo -def _refuse_outbox_kind(kind: MessageKind) -> None: +# Threads that download the messages of one receive at the same time. +DOWNLOAD_WORKERS = 10 + + +def _require_inbox_kind(kind: MessageKind) -> None: """Only the inbox is read through ``receive`` so far.""" if kind is not MessageKind.PROPOSED_CHANGES: raise NotImplementedError( @@ -73,23 +79,42 @@ def delivered(self) -> list[str]: def receive( self, peer: str, kind: MessageKind, max_messages: int | None = None ) -> list[RawMessage]: - _refuse_outbox_kind(kind) + # One listing for the whole inbox, then the files download in parallel. + _require_inbox_kind(kind) listed = self._conn.owner_get_inbox_file_metadatas(peer)[:max_messages] + contents = self._download_all([m["file_id"] for m in listed]) return [ RawMessage( sender=peer, msg_id=m["file_name"], kind=kind, - data=self._conn.download_file(m["file_id"]), + data=data, handle=m["file_id"], ) - for m in listed + for m, data in zip(listed, contents) ] + def _download_all(self, file_ids: list[str]) -> list[bytes]: + """The bytes of each file, in order. Several files download over a + thread pool with one connection copy per thread, because the Drive + client is not thread-safe.""" + if len(file_ids) <= 1: + return [self._conn.download_file(f) for f in file_ids] + per_thread = threading.local() + + def download(file_id: str) -> bytes: + if not hasattr(per_thread, "conn"): + per_thread.conn = self._conn.clone_for_thread() + return per_thread.conn.download_file(file_id) + + workers = min(DOWNLOAD_WORKERS, len(file_ids)) + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(download, file_ids)) + def mark_processed(self, message: RawMessage) -> None: # The archive folder is named after ``message.sender``, which receive # set from the inbox folder, never from the message body. - _refuse_outbox_kind(message.kind) + _require_inbox_kind(message.kind) self._conn.owner_archive_inbox_file(message.handle, message.sender) def outbox_messages(self, peer: str, since_timestamp: float | None) -> list[bytes]: diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 8d48ec70ff0..2d97c68b731 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -395,12 +395,18 @@ def from_token_path( return res @classmethod - def from_service(cls, email: str, mock_service: Any) -> "GDriveConnection": + def from_service( + cls, + email: str, + mock_service: Any, + warm_syftbox_folder_id_cache: bool = True, + ) -> "GDriveConnection": """Create a GDriveConnection using a mock drive service for testing. Args: email: Email of the user mock_service: MockDriveService instance to use instead of real API + warm_syftbox_folder_id_cache: See ``setup``. Returns: GDriveConnection configured with the mock service @@ -412,7 +418,10 @@ def from_service(cls, email: str, mock_service: Any) -> "GDriveConnection": res = cls(email=email, token_path=None) if isinstance(mock_service, MockDriveService): mock_service = MockDriveService(mock_service._backing_store, email) - res.setup(drive_service=mock_service) + res.setup( + drive_service=mock_service, + warm_syftbox_folder_id_cache=warm_syftbox_folder_id_cache, + ) return res def setup( @@ -457,7 +466,9 @@ def copy(self) -> "GDriveConnection": ) if isinstance(self.drive_service, MockDriveService): - new_conn = GDriveConnection.from_service(self.email, self.drive_service) + new_conn = GDriveConnection.from_service( + self.email, self.drive_service, warm_syftbox_folder_id_cache=False + ) else: new_conn = GDriveConnection.from_token_path( self.email, self.token_path, warm_syftbox_folder_id_cache=False @@ -837,19 +848,16 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: """Move the inbox file ``file_id`` from ``sender_email`` to its archive folder. ``sender_email`` names the archive folder, so the caller reads it from - the transport, never from the sender-supplied message body. + the transport, never from the sender-supplied message body. The inbox + folder is the file's only parent and its id is cached, so the move is + one Drive call. """ - file_info = execute_with_retries( - self.drive_service.files().get(fileId=file_id, fields="parents") - ) - previous_parents = ",".join(file_info.get("parents", [])) - archive_folder_id = self.owner_get_archive_folder_id(sender_email) execute_with_retries( self.drive_service.files().update( fileId=file_id, - addParents=archive_folder_id, - removeParents=previous_parents, - fields="id, parents", + addParents=self.owner_get_archive_folder_id(sender_email), + removeParents=self._get_own_datasite_inbox_id(sender_email), + fields="id", supportsAllDrives=True, ) ) @@ -949,15 +957,12 @@ def get_archive_folder_id_from_drive(self, sender_email: str) -> str | None: return items[0]["id"] if items else None def owner_get_archive_folder_id(self, sender_email: str) -> str: - if sender_email in self.archive_folder_id_cache: - return self.archive_folder_id_cache[sender_email] - else: - archive_folder_id = self.get_archive_folder_id_from_drive(sender_email) - if archive_folder_id: - self.archive_folder_id_cache[sender_email] = archive_folder_id - return archive_folder_id - else: - return self.create_archive_folder(sender_email) + if sender_email not in self.archive_folder_id_cache: + archive_folder_id = self.get_archive_folder_id_from_drive( + sender_email + ) or self.create_archive_folder(sender_email) + self.archive_folder_id_cache[sender_email] = archive_folder_id + return self.archive_folder_id_cache[sender_email] @staticmethod def _extract_timestamp_from_filename(filename: str) -> float | None: diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 63141593c64..777258fc31b 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -35,6 +35,7 @@ ) from syft.sync.messages.proposed_filechange import ProposedFileChangesMessage from syft.sync.peers.exceptions import SyftPeerNotReadyError +from syft.sync.transport.types import RawMessage from syft.sync.sync.caches.datasite_owner_cache import ( DataSiteOwnerEventCache, DataSiteOwnerEventCacheConfig, @@ -306,10 +307,12 @@ def _pull_proposals(self, peer_email: str) -> None: """ proposed_count = 0 try: - while msg := self.pull_and_process_next_proposed_filechange( - peer_email, raise_on_none=False - ): - proposed_count += len(msg.proposed_file_changes) + received = self.connection_router.owner_receive_proposed_messages( + peer_email + ) + for message in received: + if processed := self.process_proposed_message(message): + proposed_count += len(processed.proposed_file_changes) except SyftPeerNotReadyError as e: warnings.warn(f"Messages from {peer_email} stay in the inbox: {e}") except Exception as e: @@ -760,35 +763,28 @@ def _plan_events(self, plan: DeliveryPlan) -> list[FileChangeEvent]: events = [self._create_current_filestate_event(p) for p in plan.send] return events + [self._deletion_event(p) for p in plan.delete] - def pull_and_process_next_proposed_filechange( - self, sender_email: str, raise_on_none=True + def process_proposed_message( + self, received: RawMessage ) -> ProposedFileChangesMessage | None: - # raise on none is useful for testing, shouldnt be used in production - next_message = ( - self.connection_router.owner_get_next_proposed_filechange_message( - sender_email=sender_email - ) - ) - if next_message is not None: - received, message = next_message - # The sender is the peer whose inbox held the message. The transport - # read the message from that peer's folder and checked the signature - # against that peer's key. The body is sender-supplied, so a claim - # there that disagrees is a spoof attempt. - if same_email(message.sender_email, sender_email): - self.handle_proposed_filechange_events_message(sender_email, message) - else: - logger.warning( - f"Dropped a proposed change from {sender_email}: the message " - f"claims sender {message.sender_email!r}." - ) - - self.connection_router.owner_mark_processed(received) - return message - elif raise_on_none: - raise ValueError("No proposed file change to process") - else: + """Apply one received proposal and mark it processed. None when it was + unreadable and archived unread.""" + message = self.connection_router.owner_read_proposed_message(received) + if message is None: return None + # The sender is the peer whose inbox held the message. The transport + # read the message from that peer's folder and checked the signature + # against that peer's key. The body is sender-supplied, so a claim + # there that disagrees is a spoof attempt. + sender_email = received.sender + if same_email(message.sender_email, sender_email): + self.handle_proposed_filechange_events_message(sender_email, message) + else: + logger.warning( + f"Dropped a proposed change from {sender_email}: the message " + f"claims sender {message.sender_email!r}." + ) + self.connection_router.owner_mark_processed(received) + return message def check_write_permission(self, sender_email: str, path: str) -> bool: """Check if sender has write access to the given path.""" diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 1ec0981713b..875a9fad7be 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -5,7 +5,10 @@ """ import inspect +import threading +import time import uuid +from collections import Counter from unittest.mock import Mock import pytest @@ -18,6 +21,7 @@ CollectionFolder, GDriveConnection, ) +from syft.sync.connections.drive import mock_drive_service from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.events.file_change_event import FileChangeEventsMessage @@ -101,12 +105,107 @@ def test_a_proposed_change_sent_through_p2p_reaches_the_owner_inbox(pair): def test_inbox_messages_are_received_oldest_first_up_to_max(pair): do_conn, ds_conn = pair names = {submitted: _proposal_name(submitted) for submitted in (2.0, 1.0, 3.0)} - for name in names.values(): - ds_conn.as_transport().p2p.send(DO, name, b"job", MessageKind.PROPOSED_CHANGES) + for submitted, name in names.items(): + ds_conn.as_transport().p2p.send( + DO, name, f"job {submitted}".encode(), MessageKind.PROPOSED_CHANGES + ) p2p = do_conn.as_transport().p2p received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES, max_messages=2) - assert [m.msg_id for m in received] == [names[1.0], names[2.0]] + assert [(m.msg_id, m.data) for m in received] == [ + (names[1.0], b"job 1.0"), + (names[2.0], b"job 2.0"), + ] + + +def _count_drive_calls(monkeypatch) -> Counter: + """Count mock Drive calls by kind: list, get, update and get_media.""" + calls = Counter() + lock = threading.Lock() + + def count(kind: str) -> None: + with lock: + calls[kind] += 1 + + for kind, request_class in ( + ("list", mock_drive_service.MockListRequest), + ("get", mock_drive_service.MockGetRequest), + ("update", mock_drive_service.MockUpdateRequest), + ): + original = request_class.execute + + def counted(self, *args, _kind=kind, _original=original, **kwargs): + count(_kind) + return _original(self, *args, **kwargs) + + monkeypatch.setattr(request_class, "execute", counted) + original_get_media = mock_drive_service.MockFilesResource.get_media + + def counted_get_media(self, *args, **kwargs): + count("get_media") + return original_get_media(self, *args, **kwargs) + + monkeypatch.setattr( + mock_drive_service.MockFilesResource, "get_media", counted_get_media + ) + return calls + + +def _send_proposals(ds_conn, count: int) -> list[str]: + names = [_proposal_name(float(i)) for i in range(count)] + for name in names: + ds_conn.as_transport().p2p.send( + DO, name, name.encode(), MessageKind.PROPOSED_CHANGES + ) + return names + + +def test_receiving_inbox_lists_once_and_downloads_each_message(pair, monkeypatch): + do_conn, ds_conn = pair + p2p = do_conn.as_transport().p2p + p2p.receive(DS, MessageKind.PROPOSED_CHANGES) # finds the inbox folder once + names = _send_proposals(ds_conn, 3) + + calls = _count_drive_calls(monkeypatch) + received = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + + assert [(m.msg_id, m.data) for m in received] == [(n, n.encode()) for n in names] + assert calls == Counter(list=1, get_media=3) + + +def test_inbox_messages_download_at_the_same_time(pair, monkeypatch): + do_conn, ds_conn = pair + _send_proposals(ds_conn, 3) + running, most_running = [0], [0] + lock = threading.Lock() + original_download = GDriveConnection.download_file + + def slow_download(self, file_id: str) -> bytes: + with lock: + running[0] += 1 + most_running[0] = max(most_running[0], running[0]) + time.sleep(0.05) + with lock: + running[0] -= 1 + return original_download(self, file_id) + + monkeypatch.setattr(GDriveConnection, "download_file", slow_download) + do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + assert most_running[0] > 1 + + +def test_marking_inbox_message_processed_is_one_drive_call(pair, monkeypatch): + do_conn, ds_conn = pair + _send_proposals(ds_conn, 2) + p2p = do_conn.as_transport().p2p + first, second = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + p2p.mark_processed(first) # finds or creates the archive folder once + + calls = _count_drive_calls(monkeypatch) + p2p.mark_processed(second) + + assert calls == Counter(update=1) + assert p2p.receive(DS, MessageKind.PROPOSED_CHANGES) == [] def test_processed_inbox_message_is_not_received_again(pair): @@ -386,10 +485,8 @@ def test_the_router_sends_proposed_changes_through_p2p(): def test_the_router_reads_the_inbox_through_p2p(): router, p2p = _router_with_mock_p2p() p2p.receive.return_value = [] - assert router.owner_get_next_proposed_filechange_message(DS) is None - p2p.receive.assert_called_once_with( - DS, MessageKind.PROPOSED_CHANGES, max_messages=1 - ) + assert router.owner_receive_proposed_messages(DS) == [] + p2p.receive.assert_called_once_with(DS, MessageKind.PROPOSED_CHANGES) def test_the_router_marks_a_handled_message_processed_through_p2p(): From 8eb0e55d8140f6249bac9c27b6fb7b22416a7e24 Mon Sep 17 00:00:00 2001 From: Koen van der Veen Date: Thu, 29 Oct 2026 12:07:42 +0100 Subject: [PATCH 3/3] fix: archive bad proposals with a warning instead of blocking the inbox A proposal that could not be read or applied stayed in the inbox and held back every later proposal from that peer. Now it is archived with a warning. Proposals only stay while the peer is not set up yet or Drive cannot be reached. A file that fails to download is skipped with a warning instead of failing the whole receive. The router reads everything in one call, and processed proposals are archived in one Drive batch request per 100. --- syft/sync/connections/connection_router.py | 54 +++++++---- .../connections/drive/drive_connections.py | 60 +++++++++--- .../connections/drive/gdrive_transport.py | 58 ++++++++---- syft/sync/sync/datasite_owner_syncer.py | 54 ++++++----- syft/sync/transport/protocols.py | 5 + tests/unit/test_delivery_tracker.py | 8 +- tests/unit/test_drive_connections.py | 94 ++++++++++++++++++- tests/unit/test_encryption.py | 14 ++- tests/unit/test_owner_inbox.py | 65 +++++++++++++ 9 files changed, 323 insertions(+), 89 deletions(-) create mode 100644 tests/unit/test_owner_inbox.py diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 8ee0873b267..619379bdbfd 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -149,35 +149,49 @@ def watcher_send_proposed_file_changes_message( filename = proposed_file_changes_message.message_filename.as_string() self.transport.p2p.send(recipient, filename, data, MessageKind.PROPOSED_CHANGES) - def owner_receive_proposed_messages(self, sender_email: str) -> list[RawMessage]: - """Every proposed-changes message from ``sender_email`` not processed - yet, oldest first, still encrypted.""" - return self.transport.p2p.receive(sender_email, MessageKind.PROPOSED_CHANGES) + def owner_receive_proposed_messages( + self, sender_email: str + ) -> list[tuple[RawMessage, ProposedFileChangesMessage]]: + """Every readable proposal waiting from ``sender_email``, oldest first, + as received and as read. Pass the received ones to + ``owner_mark_all_processed`` once they are handled. - def owner_read_proposed_message( - self, message: RawMessage - ) -> ProposedFileChangesMessage | None: - """``message`` decrypted, verified and decoded. - - A message whose format does not match the pair is marked processed - unread, with a warning, and gives None. + A message that cannot be read (wrong format for the pair, bad + signature, cannot be decrypted or decoded) is marked processed unread, + with a warning, so it does not hold back the others. Raises: - SyftPeerNotReadyError: when the message cannot be read yet. It - stays unprocessed. + SyftPeerNotReadyError: when the peer is not set up yet. Every + message stays unprocessed. """ + received = self.transport.p2p.receive( + sender_email, MessageKind.PROPOSED_CHANGES + ) + readable = [] + for raw in received: + if (message := self._read_proposed_message(raw)) is not None: + readable.append((raw, message)) + return readable + + def _read_proposed_message( + self, raw: RawMessage + ) -> ProposedFileChangesMessage | None: + """``raw`` verified, decrypted and decoded, or None when it cannot be + read; that message is then marked processed with a warning.""" try: decrypted = self.peer_store.decrypt_and_verify_if_needed( - message.sender, message.data + raw.sender, raw.data ) - except MessageFormatError as e: - warnings.warn(f"Archived message {message.msg_id} unread: {e}") - self.owner_mark_processed(message) + return ProposedFileChangesMessage.from_compressed_data(decrypted) + except SyftPeerNotReadyError: + raise + except Exception as e: + warnings.warn(f"Archived message {raw.msg_id} unread: {e!r}") + self.transport.p2p.mark_processed(raw) return None - return ProposedFileChangesMessage.from_compressed_data(decrypted) - def owner_mark_processed(self, message: RawMessage) -> None: - self.transport.p2p.mark_processed(message) + def owner_mark_all_processed(self, messages: list[RawMessage]) -> None: + self.transport.p2p.mark_all_processed(messages) def owner_write_event_messages_to_outbox( self, recipient_email: str, events_message: FileChangeEventsMessage diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 5d05f88acf8..b85c64cd395 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -11,10 +11,13 @@ from __future__ import annotations import threading +import warnings from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING +from syft.sync.connections.drive.gdrive_retry import raise_unavailable from syft.sync.platforms.gdrive_files_platform import GdriveFilesPlatform +from syft.sync.transport.exceptions import TransportUnavailableError from syft.sync.transport.protocols import ( Collections, OwnerStore, @@ -35,6 +38,24 @@ DOWNLOAD_WORKERS = 10 +def _download_or_skip(connection: GDriveConnection, listed: dict) -> bytes | None: + """The bytes of one listed inbox file, or None with a warning when this + file cannot be downloaded, for example because it was deleted. + + Raises: + TransportUnavailableError: when Drive cannot be reached. + """ + name = listed["file_name"] + try: + with raise_unavailable(f"Drive is unreachable; inbox message {name} waits"): + return connection.download_file(listed["file_id"]) + except TransportUnavailableError: + raise + except Exception as e: + warnings.warn(f"Skipped inbox message {name}: download failed: {e!r}") + return None + + def _require_inbox_kind(kind: MessageKind) -> None: """Only the inbox is read through ``receive`` so far.""" if kind is not MessageKind.PROPOSED_CHANGES: @@ -79,10 +100,18 @@ def delivered(self) -> list[str]: def receive( self, peer: str, kind: MessageKind, max_messages: int | None = None ) -> list[RawMessage]: - # One listing for the whole inbox, then the files download in parallel. + """One listing for the whole inbox, then the files download in parallel. + + A file that cannot be downloaded is left out with a warning, so it does + not hold back the others. + + Raises: + TransportUnavailableError: when Drive cannot be reached. Every + message stays in the inbox. + """ _require_inbox_kind(kind) listed = self._conn.owner_get_inbox_file_metadatas(peer)[:max_messages] - contents = self._download_all([m["file_id"] for m in listed]) + contents = self._download_all(listed) return [ RawMessage( sender=peer, @@ -92,24 +121,26 @@ def receive( handle=m["file_id"], ) for m, data in zip(listed, contents) + if data is not None ] - def _download_all(self, file_ids: list[str]) -> list[bytes]: - """The bytes of each file, in order. Several files download over a - thread pool with one connection copy per thread, because the Drive - client is not thread-safe.""" - if len(file_ids) <= 1: - return [self._conn.download_file(f) for f in file_ids] + def _download_all(self, listed: list[dict]) -> list[bytes | None]: + """The bytes of each listed file, in order; None for a file that could + not be downloaded. Several files download over a thread pool with one + connection copy per thread, because the Drive client is not + thread-safe.""" + if len(listed) <= 1: + return [_download_or_skip(self._conn, m) for m in listed] per_thread = threading.local() - def download(file_id: str) -> bytes: + def download(m: dict) -> bytes | None: if not hasattr(per_thread, "conn"): per_thread.conn = self._conn.clone_for_thread() - return per_thread.conn.download_file(file_id) + return _download_or_skip(per_thread.conn, m) - workers = min(DOWNLOAD_WORKERS, len(file_ids)) + workers = min(DOWNLOAD_WORKERS, len(listed)) with ThreadPoolExecutor(max_workers=workers) as pool: - return list(pool.map(download, file_ids)) + return list(pool.map(download, listed)) def mark_processed(self, message: RawMessage) -> None: # The archive folder is named after ``message.sender``, which receive @@ -117,6 +148,11 @@ def mark_processed(self, message: RawMessage) -> None: _require_inbox_kind(message.kind) self._conn.owner_archive_inbox_file(message.handle, message.sender) + def mark_all_processed(self, messages: list[RawMessage]) -> None: + for message in messages: + _require_inbox_kind(message.kind) + self._conn.owner_archive_inbox_files([(m.handle, m.sender) for m in messages]) + def outbox_messages(self, peer: str, since_timestamp: float | None) -> list[bytes]: """The messages in ``peer``'s outbox to me, newer than ``since_timestamp``.""" return self._conn.watcher_download_raw_events_from_outbox(peer, since_timestamp) diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 2d97c68b731..653099b1dd1 100644 --- a/syft/sync/connections/drive/gdrive_transport.py +++ b/syft/sync/connections/drive/gdrive_transport.py @@ -5,6 +5,7 @@ import logging import pickle import re +import warnings from functools import lru_cache from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional @@ -67,6 +68,8 @@ # Timeout for Google API requests (in seconds) GOOGLE_API_TIMEOUT = 120 # 2 minutes +# Google Drive batch API has a limit of 100 requests per batch +DRIVE_BATCH_SIZE = 100 SYFTBOX_FOLDER = "SyftBox" GOOGLE_FOLDER_MIME_TYPE = "application/vnd.google-apps.folder" @@ -848,18 +851,40 @@ def owner_archive_inbox_file(self, file_id: str, sender_email: str) -> None: """Move the inbox file ``file_id`` from ``sender_email`` to its archive folder. ``sender_email`` names the archive folder, so the caller reads it from - the transport, never from the sender-supplied message body. The inbox - folder is the file's only parent and its id is cached, so the move is - one Drive call. + the transport, never from the sender-supplied message body. """ - execute_with_retries( - self.drive_service.files().update( - fileId=file_id, - addParents=self.owner_get_archive_folder_id(sender_email), - removeParents=self._get_own_datasite_inbox_id(sender_email), - fields="id", - supportsAllDrives=True, - ) + execute_with_retries(self._archive_inbox_file_request(file_id, sender_email)) + + def owner_archive_inbox_files(self, files: list[tuple[str, str]]) -> None: + """Move each inbox file, given as ``(file_id, sender_email)``, to its + sender's archive folder, up to ``DRIVE_BATCH_SIZE`` moves per request. + + A move that fails is warned about and its file stays in the inbox. + """ + + def callback(request_id, response, exception): + if exception is not None: + warnings.warn( + f"Could not archive inbox file {request_id}: {exception!r}" + ) + + for i in range(0, len(files), DRIVE_BATCH_SIZE): + batch = self.drive_service.new_batch_http_request(callback=callback) + for file_id, sender_email in files[i : i + DRIVE_BATCH_SIZE]: + request = self._archive_inbox_file_request(file_id, sender_email) + batch.add(request, request_id=file_id) + batch_execute_with_retries(batch) + + def _archive_inbox_file_request(self, file_id: str, sender_email: str): + """The request that moves an inbox file to its archive folder. The inbox + folder is the file's only parent and both folder ids are cached, so the + move is one Drive call.""" + return self.drive_service.files().update( + fileId=file_id, + addParents=self.owner_get_archive_folder_id(sender_email), + removeParents=self._get_own_datasite_inbox_id(sender_email), + fields="id", + supportsAllDrives=True, ) def _has_permission(self, file_id: str, email: str) -> bool: @@ -1287,10 +1312,8 @@ def callback(request_id, response, exception): f"Failed to delete {request_id}: error status {response.get('status')}" ) - # Google Drive batch API has a limit of 100 requests per batch - BATCH_SIZE = 100 - for i in range(0, len(file_ids), BATCH_SIZE): - chunk = file_ids[i : i + BATCH_SIZE] + for i in range(0, len(file_ids), DRIVE_BATCH_SIZE): + chunk = file_ids[i : i + DRIVE_BATCH_SIZE] batch = self.drive_service.new_batch_http_request(callback=callback) for file_id in chunk: batch.add(self.drive_service.files().delete(fileId=file_id)) @@ -1590,9 +1613,8 @@ def callback(request_id, response, exception): if exception and "alreadyShared" not in str(exception): raise exception - BATCH_SIZE = 100 - for i in range(0, len(users), BATCH_SIZE): - chunk = users[i : i + BATCH_SIZE] + for i in range(0, len(users), DRIVE_BATCH_SIZE): + chunk = users[i : i + DRIVE_BATCH_SIZE] batch = self.drive_service.new_batch_http_request(callback=callback) for user_email in chunk: permission = { diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 777258fc31b..bbf919d93ed 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -300,19 +300,24 @@ def sync(self, peer_emails: list[str], recompute_hashes: bool = True): self._save_rolling_state() def _pull_proposals(self, peer_email: str) -> None: - """Apply every proposed change waiting from ``peer_email``. - - A proposal that fails is logged as an error, so that it does not hold - back the other peers or the owner's changes. + """Apply every proposed change waiting from ``peer_email``, then archive + them all in one go. + + A proposal that cannot be read or applied is archived with a warning, + so it never holds back the ones after it. Proposals stay in the inbox + only while the peer is not set up yet or Drive cannot be reached; that + is logged, so it does not hold back the other peers or the owner's + changes. """ proposed_count = 0 try: received = self.connection_router.owner_receive_proposed_messages( peer_email ) - for message in received: - if processed := self.process_proposed_message(message): - proposed_count += len(processed.proposed_file_changes) + for raw, message in received: + if self._apply_proposed_message(raw, message): + proposed_count += len(message.proposed_file_changes) + self.connection_router.owner_mark_all_processed([r for r, _ in received]) except SyftPeerNotReadyError as e: warnings.warn(f"Messages from {peer_email} stay in the inbox: {e}") except Exception as e: @@ -763,28 +768,31 @@ def _plan_events(self, plan: DeliveryPlan) -> list[FileChangeEvent]: events = [self._create_current_filestate_event(p) for p in plan.send] return events + [self._deletion_event(p) for p in plan.delete] - def process_proposed_message( - self, received: RawMessage - ) -> ProposedFileChangesMessage | None: - """Apply one received proposal and mark it processed. None when it was - unreadable and archived unread.""" - message = self.connection_router.owner_read_proposed_message(received) - if message is None: - return None + def _apply_proposed_message( + self, received: RawMessage, message: ProposedFileChangesMessage + ) -> bool: + """Apply one proposal; False, with a warning, when applying it failed. + Either way the caller archives it, so it does not hold back the next.""" # The sender is the peer whose inbox held the message. The transport # read the message from that peer's folder and checked the signature # against that peer's key. The body is sender-supplied, so a claim # there that disagrees is a spoof attempt. sender_email = received.sender - if same_email(message.sender_email, sender_email): - self.handle_proposed_filechange_events_message(sender_email, message) - else: - logger.warning( - f"Dropped a proposed change from {sender_email}: the message " - f"claims sender {message.sender_email!r}." + try: + if same_email(message.sender_email, sender_email): + self.handle_proposed_filechange_events_message(sender_email, message) + else: + logger.warning( + f"Dropped a proposed change from {sender_email}: the message " + f"claims sender {message.sender_email!r}." + ) + return True + except Exception as e: + warnings.warn( + f"Archived proposed change {received.msg_id} from {sender_email}: " + f"applying it failed: {e!r}" ) - self.connection_router.owner_mark_processed(received) - return message + return False def check_write_permission(self, sender_email: str, path: str) -> bool: """Check if sender has write access to the given path.""" diff --git a/syft/sync/transport/protocols.py b/syft/sync/transport/protocols.py index d319ab04455..f9b4744e0cf 100644 --- a/syft/sync/transport/protocols.py +++ b/syft/sync/transport/protocols.py @@ -73,6 +73,11 @@ def mark_processed(self, message: RawMessage) -> None: """Record that ``message`` was processed; ``receive`` never returns it again.""" + def mark_all_processed(self, messages: list[RawMessage]) -> None: + """``mark_processed`` for each of ``messages``.""" + for message in messages: + self.mark_processed(message) + class OwnerStore(ABC): """The data owner's own storage: event log, checkpoints, rolling state. diff --git a/tests/unit/test_delivery_tracker.py b/tests/unit/test_delivery_tracker.py index 273c08ccc4b..9623e9ce288 100644 --- a/tests/unit/test_delivery_tracker.py +++ b/tests/unit/test_delivery_tracker.py @@ -428,7 +428,7 @@ def test_large_catch_up_split_into_messages(monkeypatch): assert sorted(len(m.events) for m in messages) == [1, 2] -def test_outdated_proposal_logged_without_blocking_owner_changes(caplog): +def test_outdated_proposal_archived_without_blocking_owner_changes(): ds_manager, do_manager = _pair() shared = do_manager.datasite_owner_syncer.perm_context.open("shared/") shared.grant_write_access(ds_manager.email) @@ -437,12 +437,12 @@ def test_outdated_proposal_logged_without_blocking_owner_changes(caplog): ds_manager.sync() # The DS edits a.txt from v1 while the DO edits it too, so the - # proposal is outdated and its handling fails on every sync. + # proposal is outdated and applying it fails; it is archived with a warning. ds_manager._send_file_change(f"{do_manager.email}/shared/a.txt", "ds edit") _write(do_manager, "shared/a.txt", "do edit") _write(do_manager, "shared/b.txt", "new") - do_manager.sync(auto_checkpoint=False) - assert ProposedEventFileOutdatedException.__name__ in caplog.text + with pytest.warns(UserWarning, match=ProposedEventFileOutdatedException.__name__): + do_manager.sync(auto_checkpoint=False) assert ( do_manager.syftbox_folder / do_manager.email / "shared/a.txt" ).read_text() == "do edit" diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 875a9fad7be..31614b924f5 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -119,7 +119,8 @@ def test_inbox_messages_are_received_oldest_first_up_to_max(pair): def _count_drive_calls(monkeypatch) -> Counter: - """Count mock Drive calls by kind: list, get, update and get_media.""" + """Count mock Drive calls by kind: list, get, update, batch and get_media. + A batch request also counts each update inside it.""" calls = Counter() lock = threading.Lock() @@ -131,6 +132,7 @@ def count(kind: str) -> None: ("list", mock_drive_service.MockListRequest), ("get", mock_drive_service.MockGetRequest), ("update", mock_drive_service.MockUpdateRequest), + ("batch", mock_drive_service.MockBatchHttpRequest), ): original = request_class.execute @@ -208,6 +210,67 @@ def test_marking_inbox_message_processed_is_one_drive_call(pair, monkeypatch): assert p2p.receive(DS, MessageKind.PROPOSED_CHANGES) == [] +def test_marking_inbox_messages_processed_is_one_batch_request(pair, monkeypatch): + do_conn, ds_conn = pair + _send_proposals(ds_conn, 4) + p2p = do_conn.as_transport().p2p + first, *rest = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + p2p.mark_processed(first) # finds or creates the archive folder once + + calls = _count_drive_calls(monkeypatch) + p2p.mark_all_processed(rest) + + assert calls == Counter(batch=1, update=3) # one request holding three moves + assert p2p.receive(DS, MessageKind.PROPOSED_CHANGES) == [] + + +def test_move_that_fails_in_batch_warns_and_rest_are_archived(pair): + do_conn, ds_conn = pair + _send_proposals(ds_conn, 2) + p2p = do_conn.as_transport().p2p + first, second = p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + gone = first.model_copy(update={"handle": "no-such-file"}) + + with pytest.warns(UserWarning, match="Could not archive inbox file no-such-file"): + p2p.mark_all_processed([gone, second]) + + assert [m.msg_id for m in p2p.receive(DS, MessageKind.PROPOSED_CHANGES)] == [ + first.msg_id + ] + + +def test_inbox_message_that_fails_to_download_is_skipped_with_warning( + pair, monkeypatch +): + do_conn, ds_conn = pair + names = _send_proposals(ds_conn, 3) + [broken_id] = [ + m["file_id"] + for m in do_conn.owner_get_inbox_file_metadatas(DS) + if m["file_name"] == names[1] + ] + original_download = GDriveConnection.download_file + + def download_all_but_broken(self, file_id: str) -> bytes: + if file_id == broken_id: + raise ValueError("file is gone") + return original_download(self, file_id) + + monkeypatch.setattr(GDriveConnection, "download_file", download_all_but_broken) + with pytest.warns(UserWarning, match=f"Skipped inbox message {names[1]}"): + received = do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + + assert [m.msg_id for m in received] == [names[0], names[2]] + + +def test_receiving_inbox_while_drive_is_down_raises(pair, monkeypatch): + do_conn, ds_conn = pair + _send_proposals(ds_conn, 2) + monkeypatch.setattr(GDriveConnection, "download_file", _drive_down) + with pytest.raises(TransportUnavailableError): + do_conn.as_transport().p2p.receive(DS, MessageKind.PROPOSED_CHANGES) + + def test_processed_inbox_message_is_not_received_again(pair): do_conn, ds_conn = pair ds_conn.as_transport().p2p.send( @@ -489,13 +552,36 @@ def test_the_router_reads_the_inbox_through_p2p(): p2p.receive.assert_called_once_with(DS, MessageKind.PROPOSED_CHANGES) -def test_the_router_marks_a_handled_message_processed_through_p2p(): +def test_the_router_archives_unreadable_proposal_and_returns_the_rest(): + router, p2p = _router_with_mock_p2p() + proposal = ProposedFileChangesMessage(sender_email=DS, proposed_file_changes=[]) + unreadable = RawMessage( + sender=DS, msg_id="bad", kind=MessageKind.PROPOSED_CHANGES, data=b"garbage" + ) + readable = RawMessage( + sender=DS, + msg_id="good", + kind=MessageKind.PROPOSED_CHANGES, + data=proposal.as_compressed_data(), + ) + p2p.receive.return_value = [unreadable, readable] + + with pytest.warns(UserWarning, match="Archived message bad unread"): + pairs = router.owner_receive_proposed_messages(DS) + + assert [(raw.msg_id, message.id) for raw, message in pairs] == [ + ("good", proposal.id) + ] + p2p.mark_processed.assert_called_once_with(unreadable) + + +def test_the_router_marks_handled_messages_processed_through_p2p(): router, p2p = _router_with_mock_p2p() message = RawMessage( sender=DS, msg_id="m", kind=MessageKind.PROPOSED_CHANGES, data=b"" ) - router.owner_mark_processed(message) - p2p.mark_processed.assert_called_once_with(message) + router.owner_mark_all_processed([message]) + p2p.mark_all_processed.assert_called_once_with([message]) def test_the_router_reads_the_outbox_through_p2p(): diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index 51567c8ef2a..365450ac0df 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -610,8 +610,9 @@ def test_verified_encrypted_message_flow(): assert len(matching) > 0 -def test_tampered_inbox_message_logs_error(caplog): - """Tampered message in mock inbox fails signature verification on sync.""" +def test_tampered_inbox_message_is_archived_with_warning(): + """Tampered message in mock inbox cannot be read on sync, so it is archived + with a warning instead of blocking the inbox.""" ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( encryption=True, ) @@ -633,9 +634,6 @@ def test_tampered_inbox_message_logs_error(caplog): tampered[-1] ^= 0xFF mock_file.content = bytes(tampered) - # DO sync logs the failed signature verification - do_manager.sync() - assert any( - r.levelname == "ERROR" and ds_manager.email in r.getMessage() - for r in caplog.records - ) + with pytest.warns(UserWarning, match="unread"): + do_manager.sync() + assert do_conn.get_file_metadatas_from_folder(inbox_id) == [] diff --git a/tests/unit/test_owner_inbox.py b/tests/unit/test_owner_inbox.py new file mode 100644 index 00000000000..d6c88bfd346 --- /dev/null +++ b/tests/unit/test_owner_inbox.py @@ -0,0 +1,65 @@ +"""The DO reads every proposal waiting from a peer in one sync. A proposal that +cannot be read or applied is archived with a warning, so it never holds back +the proposals after it.""" + +import time +import uuid + +import pytest + +from syft.sync.syftbox_manager import SyftboxManager +from syft.sync.sync.datasite_owner_syncer import DatasiteOwnerSyncer +from syft.sync.transport import MessageKind +from tests.unit.test_sender_binding import send_spoofed_message, sync_do, written_paths + + +def _pair_with_write_access() -> tuple[SyftboxManager, SyftboxManager]: + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection() + perms = do_manager.datasite_owner_syncer.perm_context.open(".") + perms.grant_write_access(ds_manager.email) + return ds_manager, do_manager + + +def _send_proposal(ds_manager: SyftboxManager, do_email: str, path: str) -> None: + send_spoofed_message(ds_manager, do_email, ds_manager.email, path) + + +def _inbox(do_manager: SyftboxManager, ds_email: str) -> list: + p2p = do_manager.datasite_owner_syncer.connection_router.transport.p2p + return p2p.receive(ds_email, MessageKind.PROPOSED_CHANGES) + + +def test_proposal_that_fails_to_apply_is_archived_and_next_is_applied(monkeypatch): + ds_manager, do_manager = _pair_with_write_access() + _send_proposal(ds_manager, do_manager.email, "first.txt") + _send_proposal(ds_manager, do_manager.email, "second.txt") + original_handle = DatasiteOwnerSyncer.handle_proposed_filechange_events_message + + def fail_on_first(self, sender_email, message): + if str(message.proposed_file_changes[0].path_in_datasite) == "first.txt": + raise RuntimeError("cannot apply") + return original_handle(self, sender_email, message) + + monkeypatch.setattr( + DatasiteOwnerSyncer, "handle_proposed_filechange_events_message", fail_on_first + ) + with pytest.warns(UserWarning, match="applying it failed"): + sync_do(do_manager) + + assert "second.txt" in written_paths(do_manager) + assert "first.txt" not in written_paths(do_manager) + assert _inbox(do_manager, ds_manager.email) == [] + + +def test_unreadable_proposal_is_archived_and_next_is_applied(): + ds_manager, do_manager = _pair_with_write_access() + older = f"msgv2_{time.time() - 60}_{uuid.uuid4().hex}.tar.gz" + ds_p2p = ds_manager.datasite_watcher_syncer.connection_router.transport.p2p + ds_p2p.send(do_manager.email, older, b"garbage", MessageKind.PROPOSED_CHANGES) + _send_proposal(ds_manager, do_manager.email, "after_garbage.txt") + + with pytest.warns(UserWarning, match=f"Archived message {older} unread"): + sync_do(do_manager) + + assert "after_garbage.txt" in written_paths(do_manager) + assert _inbox(do_manager, ds_manager.email) == []