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..619379bdbfd 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, @@ -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_get_next_proposed_filechange_message( + def owner_receive_proposed_messages( self, sender_email: str - ) -> ProposedFileChangesMessage | None: - """The next message from ``sender_email``, or None when the inbox is empty. + ) -> 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. - A message whose format does not match the pair is archived unread, - with a warning. + 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 in the inbox. + SyftPeerNotReadyError: when the peer is not set up yet. Every + message stays unprocessed. """ - p2p = self.transport.p2p - while True: - result = p2p.next_inbox_message(sender_email) - if result is None: - return None - raw_data, file_id = result - try: - decrypted = self.peer_store.decrypt_and_verify_if_needed( - 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) - continue - msg = ProposedFileChangesMessage.from_compressed_data(decrypted) - msg.platform_id = file_id - return msg + 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( + raw.sender, raw.data + ) + 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 + + 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 @@ -321,15 +335,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..b85c64cd395 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/syft/sync/connections/drive/drive_connections.py @@ -10,9 +10,14 @@ 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, @@ -26,10 +31,40 @@ 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 +# Threads that download the messages of one receive at the same time. +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: + 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 +72,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 +97,61 @@ 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]: + """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(listed) + return [ + RawMessage( + sender=peer, + msg_id=m["file_name"], + kind=kind, + data=data, + handle=m["file_id"], + ) + for m, data in zip(listed, contents) + if data is not None + ] + + 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(m: dict) -> bytes | None: + if not hasattr(per_thread, "conn"): + per_thread.conn = self._conn.clone_for_thread() + return _download_or_skip(per_thread.conn, m) + + workers = min(DOWNLOAD_WORKERS, len(listed)) + with ThreadPoolExecutor(max_workers=workers) as pool: + return list(pool.map(download, 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. + _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``.""" diff --git a/syft/sync/connections/drive/gdrive_transport.py b/syft/sync/connections/drive/gdrive_transport.py index 25c90276081..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" @@ -395,12 +398,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 +421,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 +469,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 @@ -833,44 +847,44 @@ 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() + execute_with_retries(self._archive_inbox_file_request(file_id, sender_email)) - # 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_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. - 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") - ) - 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", - supportsAllDrives=True, - ) + 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: @@ -968,15 +982,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: @@ -1097,35 +1108,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, @@ -1317,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)) @@ -1558,17 +1551,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`) # @@ -1631,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 8e01312751b..bbf919d93ed 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, @@ -299,17 +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: - 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 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: @@ -760,18 +768,17 @@ 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 - ) -> 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 - ) - if message is not 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. + 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 + try: if same_email(message.sender_email, sender_email): self.handle_proposed_filechange_events_message(sender_email, message) else: @@ -779,16 +786,13 @@ def pull_and_process_next_proposed_filechange( f"Dropped a proposed change from {sender_email}: the message " 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 + return True + except Exception as e: + warnings.warn( + f"Archived proposed change {received.msg_id} from {sender_email}: " + f"applying it failed: {e!r}" ) - return message - elif raise_on_none: - raise ValueError("No proposed file change to process") - else: - return None + 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 8423e179fb0..f9b4744e0cf 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,24 +54,30 @@ 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.""" + 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/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_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 a91874f2b2d..31614b924f5 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 @@ -90,20 +94,193 @@ 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_an_archived_inbox_message_is_not_read_again(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 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, 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, batch and get_media. + A batch request also counts each update inside it.""" + 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), + ("batch", mock_drive_service.MockBatchHttpRequest), + ): + 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_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( 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 +321,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 +547,41 @@ 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 - assert router.owner_get_next_proposed_filechange_message(DS) is None - p2p.next_inbox_message.assert_called_once_with(DS) + p2p.receive.return_value = [] + assert router.owner_receive_proposed_messages(DS) == [] + p2p.receive.assert_called_once_with(DS, MessageKind.PROPOSED_CHANGES) -def test_the_router_archives_a_handled_message_through_p2p(): +def test_the_router_archives_unreadable_proposal_and_returns_the_rest(): 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) + 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_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_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_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) == [] 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()