feat: add WebRTC p2p connection over a fake channel

packages/syft-rds/src/syft_rds/client.py

@@ -229,6 +229,17 @@ def pair_with_mock_drive_service_connection(
         )
         return cls._build_rds_pair_from_managers(ds_mgr, do_mgr)
 
+    @classmethod
+    def pair_with_mock_drive_and_mock_webrtc_connection(
+        cls, **kwargs: Any
+    ) -> tuple["SyftRDSClient", "SyftRDSClient"]:
+        """(ds, do) pair like ``pair_with_mock_drive_service_connection``, with
+        messages over WebRTC joined by a fake channel."""
+        ds_mgr, do_mgr = SyftboxManager.pair_with_mock_drive_and_mock_webrtc_connection(
+            collection_specs=DATASET_COLLECTION_SPECS, **kwargs
+        )
+        return cls._build_rds_pair_from_managers(ds_mgr, do_mgr)
+
     @classmethod
     def _pair_with_google_drive_testing_connection(
         cls, **kwargs: Any

syft/sync/connections/drive/drive_p2p.py

@@ -82,6 +82,8 @@ class DriveP2PConnection(P2PConnection):
     methods below.
     """
 
+    keeps_sent_messages = True
+
     def __init__(self, cached_client: DriveCachedClient):
         self._cached_client = cached_client
 

syft/sync/connections/protocols.py

@@ -55,6 +55,10 @@ class P2PConnection(RouterConnection, ABC):
       is sent back to the peer.
     """
 
+    # Whether sent messages stay where the peer reads them, as in Drive's
+    # outbox, so the owner compacts them. A live link drops them on receipt.
+    keeps_sent_messages: bool = False
+
     @abstractmethod
     def send(
         self, peer: str, msg_id: str, data: bytes, kind: MessageKind

syft/sync/connections/webrtc/channel.py

@@ -0,0 +1,13 @@
+"""The pipe under the WebRTC core: one per peer, carrying frames as bytes."""
+
+from abc import ABC, abstractmethod
+
+
+class Channel(ABC):
+    """One ordered, reliable pipe to one peer. An aiortc data channel in
+    production, a ``FakeChannel`` in tests. It hands each frame that arrives
+    to the core's ``on_frame``."""
+
+    @abstractmethod
+    def send(self, frame: bytes) -> None:
+        """Send ``frame`` to the peer."""

syft/sync/connections/webrtc/fake_channel.py

@@ -0,0 +1,26 @@
+"""An in-memory channel between two WebRTC cores, for tests."""
+
+from syft.sync.connections.webrtc.channel import Channel
+from syft.sync.connections.webrtc.webrtc_p2p import WebRTCP2PConnection
+
+
+class FakeChannel(Channel):
+    """One end of an in-memory pipe, in place of an aiortc data channel. A
+    frame reaches the other core inside ``send``: no threads, no network."""
+
+    def __init__(self, sender: str, receiver: WebRTCP2PConnection):
+        self.sender = sender
+        self.receiver = receiver
+
+    def send(self, frame: bytes) -> None:
+        self.receiver.on_frame(self.sender, frame)
+
+    @staticmethod
+    def connect(a: WebRTCP2PConnection, b: WebRTCP2PConnection) -> None:
+        """Join ``a`` and ``b`` and open the channel. Both ends are attached
+        before it opens, as with a real channel, so the first receipts have
+        a way back."""
+        a.attach(b.email, FakeChannel(sender=a.email, receiver=b))
+        b.attach(a.email, FakeChannel(sender=b.email, receiver=a))
+        a.on_open(b.email)
+        b.on_open(a.email)

syft/sync/connections/webrtc/message_queue.py

@@ -0,0 +1,91 @@
+"""A queue of messages kept on disk, one file per message, so it survives a
+restart and adding or removing one message never rewrites the others."""
+
+from collections.abc import Collection
+from pathlib import Path
+from typing import NamedTuple
+
+from pydantic import BaseModel, ConfigDict
+
+from syft.sync.connections.types import MessageKind
+from syft.sync.sync.caches.persisted_dict import locked_file
+
+# Bytes travel as base64 inside a JSON frame.
+BYTES_AS_BASE64 = ConfigDict(ser_json_bytes="base64", val_json_bytes="base64")
+
+
+class QueuedMessage(BaseModel):
+    """One message, as the sender's router encoded it."""
+
+    model_config = ConfigDict(frozen=True, **BYTES_AS_BASE64)
+
+    msg_id: str
+    kind: MessageKind
+    data: bytes
+
+
+class _Entry(NamedTuple):
+    """One message file, named ``<position>.<kind>.<msg_id>``."""
+
+    position: int
+    kind: MessageKind
+    msg_id: str
+    path: Path
+
+
+class MessageQueue:
+    """Messages in one folder, oldest first. Every change holds the folder lock."""
+
+    def __init__(self, folder: Path):
+        self.folder = folder
+
+    def messages(
+        self, kind: MessageKind | None = None, max_messages: int | None = None
+    ) -> list[QueuedMessage]:
+        """The messages, of ``kind`` when given, oldest first, at most
+        ``max_messages``. Only those files are read."""
+        with locked_file(self.folder, exclusive=False):
+            entries = [e for e in self._entries() if kind is None or e.kind == kind]
+            return [
+                QueuedMessage(msg_id=e.msg_id, kind=e.kind, data=e.path.read_bytes())
+                for e in entries[:max_messages]
+            ]
+
+    def add(self, message: QueuedMessage) -> None:
+        """Append ``message``, unless one with its ``msg_id`` is already here."""
+        _require_file_name(message.msg_id)
+        with locked_file(self.folder):
+            entries = self._entries()
+            if any(e.msg_id == message.msg_id for e in entries):
+                return
+            position = entries[-1].position + 1 if entries else 0
+            name = f"{position:012d}.{message.kind.value}.{message.msg_id}"
+            # Written under a hidden name first, so a reader never sees half.
+            tmp = self.folder / f".{name}.tmp"
+            tmp.write_bytes(message.data)
+            tmp.replace(self.folder / name)
+
+    def remove(self, msg_ids: Collection[str]) -> list[str]:
+        """Remove the messages with these ids; return the ids that were here."""
+        with locked_file(self.folder):
+            removed = [e for e in self._entries() if e.msg_id in msg_ids]
+            for entry in removed:
+                entry.path.unlink()
+        return [e.msg_id for e in removed]
+
+    def _entries(self) -> list[_Entry]:
+        self.folder.mkdir(parents=True, exist_ok=True)
+        entries = []
+        for path in self.folder.iterdir():
+            if path.name.startswith("."):
+                continue
+            position, kind, msg_id = path.name.split(".", 2)
+            entries.append(_Entry(int(position), MessageKind(kind), msg_id, path))
+        return sorted(entries, key=lambda e: e.position)
+
+
+def _require_file_name(msg_id: str) -> None:
+    """``msg_id`` comes from the peer and names a file, so it may not leave the
+    queue folder or hide from it."""
+    if msg_id != Path(msg_id).name or msg_id.startswith("."):
+        raise ValueError(f"msg_id {msg_id!r} is not a plain file name")

syft/sync/connections/webrtc/webrtc_p2p.py

@@ -0,0 +1,131 @@
+"""Messages between peers over WebRTC: the synchronous core.
+
+The core keeps the queues on disk, sends receipts and resends what the peer has
+not confirmed. The ``Channel`` under it only carries frames: an aiortc data
+channel in production, a ``FakeChannel`` in tests.
+"""
+
+from collections import defaultdict
+from collections.abc import Container
+from pathlib import Path
+
+from pydantic import BaseModel, ConfigDict
+
+from syft.sync.connections.protocols import P2PConnection
+from syft.sync.connections.types import MessageKind, RawMessage, SendStatus
+from syft.sync.connections.webrtc.channel import Channel
+from syft.sync.connections.webrtc.message_queue import (
+    BYTES_AS_BASE64,
+    MessageQueue,
+    QueuedMessage,
+)
+from syft.sync.sync.constants import CACHE_DIR
+
+WEBRTC_QUEUES_DIR = "webrtc"
+
+
+class Frame(BaseModel):
+    """What travels over a channel: a message, or the receipt for one."""
+
+    model_config = ConfigDict(frozen=True, **BYTES_AS_BASE64)
+
+    message: QueuedMessage | None = None
+    # The msg_id of a message the peer stored.
+    receipt: str | None = None
+
+
+class WebRTCP2PConnection(P2PConnection):
+    """Messages over one channel per peer, with a queue per peer and direction.
+
+    ``send`` stores the message in the peer's outgoing queue before sending it,
+    so a restart loses nothing, and returns ``QUEUED``. The peer stores it in
+    its incoming queue and sends a receipt back. The message then leaves the
+    outgoing queue and ``delivered()`` reports it.
+    """
+
+    def __init__(self, email: str, syftbox_folder: Path):
+        self.email = email
+        self._folder = Path(syftbox_folder) / CACHE_DIR / WEBRTC_QUEUES_DIR
+        # peer email -> its channel
+        self._channels: dict[str, Channel] = {}
+        # Peers whose channel is open and got every older message again.
+        self._open: set[str] = set()
+        self._delivered: list[str] = []
+
+    def attach(self, peer: str, channel: Channel) -> None:
+        """Use ``channel`` for ``peer`` from now on. Receipts go over it right
+        away; messages once it opens."""
+        self._channels[peer] = channel
+
+    def on_open(self, peer: str) -> None:
+        """``peer``'s channel opened: send it every message it has not
+        confirmed, oldest first. New messages follow them, so none overtakes
+        an older one."""
+        for message in self._outgoing(peer).messages():
+            self._send_frame(peer, Frame(message=message))
+        self._open.add(peer)
+
+    def send(
+        self, peer: str, msg_id: str, data: bytes, kind: MessageKind
+    ) -> SendStatus:
+        message = QueuedMessage(msg_id=msg_id, kind=kind, data=data)
+        # Stored first, and not under the folder lock while it travels: a
+        # FakeChannel brings the receipt back inside this call.
+        self._outgoing(peer).add(message)
+        if peer in self._open:
+            self._send_frame(peer, Frame(message=message))
+        return SendStatus.QUEUED
+
+    def on_frame(self, peer: str, data: bytes) -> None:
+        """Called by ``peer``'s channel with each frame that arrives."""
+        frame = Frame.model_validate_json(data)
+        if frame.message is not None:
+            self._incoming(peer).add(frame.message)
+            self._send_frame(peer, Frame(receipt=frame.message.msg_id))
+        elif frame.receipt is not None:
+            self._delivered += self._outgoing(peer).remove({frame.receipt})
+
+    def delivered(self) -> list[str]:
+        delivered, self._delivered = self._delivered, []
+        return delivered
+
+    def receive(
+        self,
+        peer: str,
+        kind: MessageKind,
+        after: str | None = None,
+        max_messages: int | None = None,
+        already_received: Container[str] = frozenset(),
+    ) -> list[RawMessage]:
+        """The messages of ``kind`` from ``peer`` in the incoming queue, in the
+        order they arrived. ``after`` and ``already_received`` are not needed:
+        the queue holds only messages not processed yet."""
+        return [
+            RawMessage(sender=peer, msg_id=m.msg_id, kind=m.kind, data=m.data)
+            for m in self._incoming(peer).messages(kind, max_messages)
+        ]
+
+    def mark_processed(self, message: RawMessage) -> None:
+        self.mark_all_processed([message])
+
+    def mark_all_processed(self, messages: list[RawMessage]) -> None:
+        """Remove ``messages`` from the incoming queues, one lock per sender."""
+        # sender email -> msg_ids processed
+        by_sender: dict[str, set[str]] = defaultdict(set)
+        for message in messages:
+            by_sender[message.sender].add(message.msg_id)
+        for sender, msg_ids in by_sender.items():
+            self._incoming(sender).remove(msg_ids)
+
+    def _send_frame(self, peer: str, frame: Frame) -> None:
+        """Send ``frame`` if ``peer`` has a channel. A message without one
+        waits in the outgoing queue until ``on_open``."""
+        channel = self._channels.get(peer)
+        if channel is not None:
+            channel.send(frame.model_dump_json().encode())
+
+    def _outgoing(self, peer: str) -> MessageQueue:
+        return MessageQueue(self._folder / "outgoing" / peer)
+
+    def _incoming(self, peer: str) -> MessageQueue:
+        return MessageQueue(self._folder / "incoming" / peer)