refactor: read the DS outbox through receive

syft/sync/connections/connection_router.py

@@ -1,5 +1,6 @@
 import logging
 import warnings
+from collections.abc import Container
 from typing import TYPE_CHECKING, List, Optional
 
 from pydantic import BaseModel
@@ -13,9 +14,6 @@
 from syft.sync.checkpoints.checkpoint import Checkpoint, IncrementalCheckpoint
 from syft.sync.checkpoints.rolling_state import RollingState
 from syft.sync.connections.base_connection import FileCollection
-from syft.sync.encryption_policy import (
-    MessageFormatError,
-)
 from syft.sync.events.file_change_event import (
     FileChangeEventsMessage,
 )
@@ -109,7 +107,7 @@ def owner_receive_proposed_messages(
     ) -> 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.
+        ``mark_all_processed`` once they are handled.
 
         A message that cannot be read (wrong format for the pair, bad
         signature, cannot be decrypted or decoded) is marked processed unread,
@@ -145,7 +143,8 @@ def _read_proposed_message(
             self.connections.p2p_con.mark_processed(raw)
             return None
 
-    def owner_mark_all_processed(self, messages: list[RawMessage]) -> None:
+    def mark_all_processed(self, messages: list[RawMessage]) -> None:
+        """Finish ``messages``, as received by either role."""
         self.connections.p2p_con.mark_all_processed(messages)
 
     def owner_write_event_messages_to_outbox(
@@ -159,43 +158,56 @@ def owner_write_event_messages_to_outbox(
             recipient_email, fname, data, MessageKind.FILE_EVENTS
         )
 
-    def watcher_get_events_messages(
-        self, peer_email: str, since_timestamp: float | None
-    ) -> List[FileChangeEventsMessage]:
-        """The outbox messages of ``peer_email``, without the skipped ones."""
-        raw_list = self.connections.p2p_con.outbox_messages(peer_email, since_timestamp)
-        messages = [
-            self.watcher_read_events_message(peer_email, raw) for raw in raw_list
-        ]
-        return [m for m in messages if m is not None]
+    def watcher_receive_events_messages(
+        self,
+        sender_email: str,
+        after: str | None = None,
+        already_received: Container[str] = frozenset(),
+    ) -> list[tuple[RawMessage, FileChangeEventsMessage | None]]:
+        """The events messages from ``sender_email`` after the message
+        ``after``, and any that landed after it, without those in
+        ``already_received``, oldest first, as received and as read. Pass the
+        received ones to ``mark_all_processed`` once they are applied.
+
+        A message that can never be read (wrong format for the pair, bad
+        signature, cannot be decrypted or decoded) comes with None and a
+        warning, so the caller moves past it. The list ends, with a warning,
+        at a message that cannot be read yet because the peer is not set up;
+        it comes again at the next receive.
+        """
+        received = self.connections.p2p_con.receive(
+            sender_email,
+            MessageKind.FILE_EVENTS,
+            after=after,
+            already_received=already_received,
+        )
+        read = []
+        for raw in received:
+            try:
+                read.append((raw, self._read_events_message(raw)))
+            except SyftPeerNotReadyError as e:
+                warnings.warn(f"Outbox messages of {sender_email} wait: {e}")
+                break
+        return read
 
-    def watcher_read_events_message(
-        self, peer_email: str, raw: bytes
-    ) -> FileChangeEventsMessage | None:
-        """Verify, decrypt and parse one outbox message of ``peer_email``.
+    def _read_events_message(self, raw: RawMessage) -> FileChangeEventsMessage | None:
+        """``raw`` verified, decrypted and decoded, or None with a warning when
+        it can never be read.
 
-        Returns None when the message is skipped, with a warning: its format
-        does not match the pair, or it cannot be read now.
+        Raises:
+            SyftPeerNotReadyError: when the peer is not set up yet.
         """
         try:
-            data = self.peer_store.decrypt_and_verify_if_needed(peer_email, raw)
-        except (
-            MessageFormatError,
-            SyftPeerNotReadyError,
-        ) as e:
-            warnings.warn(f"Skipped an outbox message of {peer_email}: {e}")
+            data = self.peer_store.decrypt_and_verify_if_needed(raw.sender, raw.data)
+            return FileChangeEventsMessage.from_compressed_data(data)
+        except SyftPeerNotReadyError:
+            raise
+        except Exception as e:
+            warnings.warn(f"Skipped unreadable outbox message {raw.msg_id}: {e!r}")
             return None
-        return FileChangeEventsMessage.from_compressed_data(data)
-
-    def watcher_get_outbox_file_metadatas(
-        self, peer_email: str, since_timestamp: float | None
-    ) -> List[dict]:
-        return self.connections.p2p_con.outbox_message_metadatas(
-            peer_email, since_timestamp
-        )
 
     def download_outbox_message(self, file_id: str) -> bytes:
-        """Download one outbox message by its handle (thread-safe)."""
+        """Download one message of my own outbox by its handle (thread-safe)."""
         return self.connections.p2p_con.clone_for_thread().fetch_message(file_id)
 
     def owner_get_outbox_file_metadatas(self, recipient_email: str) -> List[dict]:

syft/sync/connections/drive/drive_p2p.py

@@ -2,6 +2,7 @@
 
 import threading
 import warnings
+from collections.abc import Container
 from concurrent.futures import ThreadPoolExecutor
 
 from syft.sync.connections.drive.gdrive_client import DRIVE_BATCH_SIZE, GDriveClient
@@ -14,6 +15,7 @@
     execute_with_retries,
     raise_unavailable,
 )
+from syft.sync.events.file_change_event import FileChangeEventsMessageFileName
 from syft.sync.messages.proposed_filechange import FileNameParseError, MessageFileName
 from syft.sync.connections.exceptions import ConnectionUnavailableError
 from syft.sync.connections.protocols import P2PConnection
@@ -22,32 +24,36 @@
 # Threads that download the messages of one receive at the same time.
 DOWNLOAD_WORKERS = 10
 
+# How far before ``after`` the outbox is listed again, for messages that land
+# after newer ones (a retried send, or notebook and syft-bg both writing).
+# The ones in that window the DS already has are in ``already_received``.
+OUTBOX_LOOKBACK_SECONDS = 60
+
 
 def _download_or_skip(client: GDriveClient, listed: dict) -> bytes | None:
-    """The bytes of one listed inbox file, or None with a warning when this
+    """The bytes of one listed message file, or None with a warning when this
     file cannot be downloaded, for example because it was deleted.
 
     Raises:
         ConnectionUnavailableError: when Drive cannot be reached.
     """
     name = listed["file_name"]
     try:
-        with raise_unavailable(f"Drive is unreachable; inbox message {name} waits"):
+        with raise_unavailable(f"Drive is unreachable; message {name} waits"):
             return client.download_file(listed["file_id"])
     except ConnectionUnavailableError:
         raise
     except Exception as e:
-        warnings.warn(f"Skipped inbox message {name}: download failed: {e!r}")
+        warnings.warn(f"Skipped 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."
-        )
+def _before_first_failure(
+    listed: list[dict], contents: list[bytes | None]
+) -> tuple[list[dict], list[bytes | None]]:
+    """``listed`` and ``contents`` up to the first file that failed to download."""
+    end = contents.index(None) if None in contents else len(contents)
+    return listed[:end], contents[:end]
 
 
 def _valid_message_names(file_metadatas: list[dict]) -> list[MessageFileName]:
@@ -67,11 +73,13 @@ class DriveP2PConnection(P2PConnection):
     outbox. ``msg_id`` is the file name. A written file is ``DELIVERED``: it
     waits in the folder until the peer reads it.
 
-    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 DS rebuilds a lost cache by reading it
-    again. It moves to ``receive`` with send-and-forget, once the DS keeps what
-    it received in its own storage.
+    Both are read through ``receive``. A processed inbox message is archived.
+    The outbox keeps every message, so the DS rebuilds a lost cache by reading
+    it again, and ``receive`` lists only the messages from shortly before
+    ``after`` on, leaving out the ones already received. That
+    ends with send-and-forget, once the DS keeps what it received in its own
+    storage. The DO still compacts its own outbox through the Drive-only
+    methods below.
     """
 
     def __init__(self, cached_client: DriveCachedClient):
@@ -100,20 +108,38 @@ def delivered(self) -> list[str]:
         return []
 
     def receive(
-        self, peer: str, kind: MessageKind, max_messages: int | None = None
+        self,
+        peer: str,
+        kind: MessageKind,
+        after: str | None = None,
+        max_messages: int | None = None,
+        already_received: Container[str] = frozenset(),
     ) -> list[RawMessage]:
-        """One listing for the whole inbox, then the files download in parallel.
+        """One listing of the folder, 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.
+        Proposed changes come from my inbox, which archives processed messages,
+        so ``after`` is not needed there. File events come from ``peer``'s
+        outbox, which keeps them. Those from ``OUTBOX_LOOKBACK_SECONDS`` before
+        ``after`` on are listed, and the ones in ``already_received`` are left
+        out before downloading.
+
+        A file that cannot be downloaded is left out with a warning. In the
+        inbox it does not hold back the others. In the outbox it ends the list,
+        because the reader passes the last message it got as ``after`` and
+        would never get the missed one.
 
         Raises:
             ConnectionUnavailableError: when Drive cannot be reached. Every
-                message stays in the inbox.
+                message stays in its folder.
         """
-        _require_inbox_kind(kind)
-        listed = self._list_inbox(peer)[:max_messages]
+        if kind is MessageKind.PROPOSED_CHANGES:
+            listed = self._list_inbox(peer)
+        else:
+            listed = self._list_outbox(peer, after, already_received)
+        listed = listed[:max_messages]
         contents = self._download_all(listed)
+        if kind is MessageKind.FILE_EVENTS:
+            listed, contents = _before_first_failure(listed, contents)
         return [
             RawMessage(
                 sender=peer,
@@ -144,18 +170,20 @@ def download(m: dict) -> bytes | None:
             return list(pool.map(download, listed))
 
     def mark_processed(self, message: RawMessage) -> None:
+        # The outbox keeps file events; the DS passes ``after`` instead.
+        if message.kind is MessageKind.FILE_EVENTS:
+            return
         # 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)
         execute_with_retries(self._archive_request(message.handle, message.sender))
 
     def mark_all_processed(self, messages: list[RawMessage]) -> None:
-        """Archive every message, up to ``DRIVE_BATCH_SIZE`` moves per request.
+        """Archive every inbox message, up to ``DRIVE_BATCH_SIZE`` moves per
+        request. File events stay in the outbox.
 
         A move that fails is warned about and its file stays in the inbox.
         """
-        for message in messages:
-            _require_inbox_kind(message.kind)
+        messages = [m for m in messages if m.kind is MessageKind.PROPOSED_CHANGES]
 
         def callback(request_id, response, exception):
             if exception is not None:
@@ -172,42 +200,6 @@ def callback(request_id, response, exception):
                 )
             batch_execute_with_retries(batch)
 
-    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._cached_client.client.download_file(m["file_id"])
-            for m in self.outbox_message_metadatas(peer, since_timestamp)
-        ]
-
-    def outbox_message_metadatas(
-        self, peer: str, since_timestamp: float | None
-    ) -> list[dict]:
-        """Like ``outbox_messages``, without downloading: each entry's
-        ``file_id`` is the handle to ``fetch_message`` it with."""
-        folder_id = self._cached_client.get_peer_datasite_outbox_id(peer)
-        if folder_id is None:
-            return []
-
-        file_metadatas = self._cached_client.get_file_metadatas_from_folder(
-            folder_id, since_timestamp=since_timestamp
-        )
-        valid_fname_objs = valid_events_from_file_metadatas(file_metadatas)
-        name_to_id = {f["name"]: f["id"] for f in file_metadatas}
-
-        result = []
-        for fname_obj in sorted(valid_fname_objs, key=lambda x: x.timestamp):
-            if since_timestamp is None or fname_obj.timestamp > since_timestamp:
-                file_name = fname_obj.as_string()
-                if file_name in name_to_id:
-                    result.append(
-                        {
-                            "file_id": name_to_id[file_name],
-                            "file_name": file_name,
-                            "timestamp": fname_obj.timestamp,
-                        }
-                    )
-        return result
-
     def fetch_message(self, handle: str) -> bytes:
         return self._cached_client.client.download_file(handle)
 
@@ -244,6 +236,37 @@ def _list_inbox(self, peer: str) -> list[dict]:
             )
         ]
 
+    def _list_outbox(
+        self, peer: str, after: str | None, already_received: Container[str]
+    ) -> list[dict]:
+        """The messages in ``peer``'s outbox to me from ``OUTBOX_LOOKBACK_SECONDS``
+        before the message ``after`` on, without those in ``already_received``,
+        oldest first, each with ``file_id`` and ``file_name``, without
+        downloading them. An outbox whose folder is not found holds none.
+
+        The look-back lists ``after`` and the messages just before it again,
+        which the DS received in earlier syncs; ``already_received`` skips them."""
+        folder_id = self._cached_client.get_peer_datasite_outbox_id(peer)
+        if folder_id is None:
+            return []
+        since = None
+        if after is not None:
+            sent_at = FileChangeEventsMessageFileName.from_string(after).timestamp
+            since = sent_at - OUTBOX_LOOKBACK_SECONDS
+        file_metadatas = self._cached_client.get_file_metadatas_from_folder(
+            folder_id, since_timestamp=since
+        )
+        name_to_id = {f["name"]: f["id"] for f in file_metadatas}
+        names = sorted(
+            valid_events_from_file_metadatas(file_metadatas), key=lambda x: x.timestamp
+        )
+        return [
+            {"file_id": name_to_id[name.as_string()], "file_name": name.as_string()}
+            for name in names
+            if name.as_string() in name_to_id
+            and name.as_string() not in already_received
+        ]
+
     def _archive_request(self, file_id: str, sender: 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

syft/sync/connections/protocols.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Container
 from typing import TYPE_CHECKING, Self
 
 from pydantic import BaseModel, ConfigDict
@@ -72,19 +73,32 @@ def delivered(self) -> list[str]:
 
     @abstractmethod
     def receive(
-        self, peer: str, kind: MessageKind, max_messages: int | None = None
+        self,
+        peer: str,
+        kind: MessageKind,
+        after: str | None = None,
+        max_messages: int | None = None,
+        already_received: Container[str] = frozenset(),
     ) -> list[RawMessage]:
         """Messages of ``kind`` from ``peer`` that are not processed yet, oldest
         first, at most ``max_messages`` of them.
 
+        ``after`` is the ``msg_id`` of the newest message the receiver
+        processed. A connection that keeps processed messages, as Drive keeps
+        the outbox, returns the later ones with no gap: a message it cannot
+        return now ends the list. It may also list back a bit before ``after``
+        for late messages, skipping the ones in ``already_received``. A
+        connection that drops processed messages ignores both.
+
         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 mark_processed(self, message: RawMessage) -> None:
         """Record that ``message`` was processed; ``receive`` never returns it
-        again."""
+        again, or, on a connection that keeps processed messages, not once
+        ``after`` is at or past it."""
 
     def mark_all_processed(self, messages: list[RawMessage]) -> None:
         """``mark_processed`` for each of ``messages``."""