refactor: read the owner inbox through p2p receive and mark_processed

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(

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
     # =========================================================================

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,22 +31,52 @@
 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.
 
     A proposed change goes to the data owner's inbox and file events go to the
     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``."""