feat: keep the DS's watcher state on its own Drive and restore it at startup

syft/sync/checkpoints/checkpoint.py

@@ -12,10 +12,22 @@
 """
 
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List
-
-from pydantic import BaseModel, Field
+from typing import TYPE_CHECKING, Any, Dict, List
+
+from pydantic import (
+    BaseModel,
+    Field,
+    computed_field,
+    field_serializer,
+    model_validator,
+)
 
+from syft.sync.utils.file_content import (
+    ContentType,
+    content_to_json,
+    content_type_of,
+    decode_content,
+)
 from syft.sync.utils.syftbox_utils import (
     compress_data,
     create_event_timestamp,
@@ -48,13 +60,32 @@ def _check_version(version: int, kind: str) -> None:
 # Default compacting threshold: merge after this many incremental checkpoints
 DEFAULT_COMPACTING_THRESHOLD = 4
 
+# Default threshold for a checkpoint: this many events in the DO's rolling
+# state, or this many paths in the rolling state that a DS keeps for a DO
+DEFAULT_CHECKPOINT_EVENT_THRESHOLD = 50
+
 
 class CheckpointFile(BaseModel):
     """Represents a single file in the checkpoint."""
 
     path: str  # Relative path in datasite (e.g., "public/syft_datasets/data.csv")
     hash: str  # Hash of the content (SHA256 hex string)
-    content: str  # File content
+    content: str | bytes  # File content: bytes for a binary file
+
+    @computed_field
+    @property
+    def content_type(self) -> ContentType | None:
+        """Saved with the file, so that a read gives bytes back as bytes."""
+        return content_type_of(self.content)
+
+    @field_serializer("content", when_used="json")
+    def _content_as_text(self, content: str | bytes) -> str | None:
+        return content_to_json(content)
+
+    @model_validator(mode="before")
+    @classmethod
+    def _content_from_text(cls, data: Any) -> Any:
+        return decode_content(data)
 
 
 class Checkpoint(BaseModel):
@@ -133,6 +164,35 @@ def get_file_contents(self) -> Dict[str, str]:
         """Get file contents as a dict mapping path to content."""
         return {f.path: f.content for f in self.files}
 
+    def with_events(self, events: List["FileChangeEvent"]) -> "Checkpoint":
+        """A new checkpoint: these files with ``events`` applied in order. A
+        change sets its file, and a delete removes it.
+
+        Raises:
+            ValueError: if a change has no content or no hash.
+        """
+        files = {file.path: file for file in self.files}
+        last_event_timestamp = self.last_event_timestamp
+        for event in events:
+            path = str(event.path_in_datasite)
+            if event.is_deleted:
+                files.pop(path, None)
+            elif event.content is None or event.new_hash is None:
+                raise ValueError(
+                    f"The change of {path} in {self.email} has no content or no hash."
+                )
+            else:
+                files[path] = CheckpointFile(
+                    path=path, hash=event.new_hash, content=event.content
+                )
+            if last_event_timestamp is None or event.timestamp > last_event_timestamp:
+                last_event_timestamp = event.timestamp
+        return Checkpoint(
+            email=self.email,
+            files=list(files.values()),
+            last_event_timestamp=last_event_timestamp,
+        )
+
     def as_compressed_data(self) -> bytes:
         """Compress checkpoint for storage."""
         return compress_data(self.model_dump_json().encode("utf-8"))
@@ -204,59 +264,6 @@ def from_compressed_data(cls, data: bytes) -> "IncrementalCheckpoint":
         return checkpoint
 
 
-def compact_incremental_checkpoints(
-    checkpoints: List[IncrementalCheckpoint],
-    email: str,
-) -> Checkpoint:
-    """
-    Merge multiple incremental checkpoints into a single full Checkpoint.
-
-    Events are merged in order (by sequence number), with later events
-    overwriting earlier ones for the same file path. Deleted files are excluded.
-
-    Args:
-        checkpoints: List of incremental checkpoints to merge.
-        email: Email for the resulting checkpoint.
-
-    Returns:
-        A full Checkpoint with the merged state.
-    """
-    # Sort by sequence number to process in order
-    sorted_checkpoints = sorted(checkpoints, key=lambda c: c.sequence_number)
-
-    # Merge events: later events overwrite earlier ones for same path
-    merged_events: Dict[str, "FileChangeEvent"] = {}
-    last_event_timestamp = None
-
-    for cp in sorted_checkpoints:
-        for event in cp.events:
-            merged_events[str(event.path_in_datasite)] = event
-            if event.timestamp is not None:
-                if (
-                    last_event_timestamp is None
-                    or event.timestamp > last_event_timestamp
-                ):
-                    last_event_timestamp = event.timestamp
-
-    # Convert to full Checkpoint (only non-deleted files)
-    files = []
-    for event in merged_events.values():
-        if not event.is_deleted and event.content is not None:
-            files.append(
-                CheckpointFile(
-                    path=str(event.path_in_datasite),
-                    hash=str(event.new_hash),
-                    content=event.content,
-                )
-            )
-
-    return Checkpoint(
-        email=email,
-        files=files,
-        last_event_timestamp=last_event_timestamp,
-    )
-
-
 # Import here to avoid circular imports
 from syft.sync.events.file_change_event import FileChangeEvent  # noqa: E402
 

syft/sync/checkpoints/rolling_state.py

@@ -61,6 +61,9 @@ class RollingState(BaseModel):
     base_checkpoint_timestamp: float  # Timestamp of the checkpoint this builds on
     last_event_timestamp: float | None = None  # Timestamp of last event included
     events: List[FileChangeEvent] = Field(default_factory=list)
+    # The ids of the events messages that a DS received from this owner. The
+    # owner's own rolling state keeps none.
+    received_msg_ids: List[str] = Field(default_factory=list)
 
     @property
     def filename(self) -> str:

syft/sync/connections/connection_router.py

@@ -575,3 +575,51 @@ def get_rolling_state(self) -> RollingState | None:
     def delete_rolling_state(self) -> None:
         """Delete rolling state from the storage backend."""
         self.connections.owner_store_con.delete_rolling_state()
+
+    # =========================================================================
+    # WATCHER STATE (not encrypted until there is key recovery: a new key, as
+    # an enclave gets at each boot, could not read it back)
+    # =========================================================================
+
+    def watcher_save_checkpoint(self, peer_email: str, checkpoint: Checkpoint) -> None:
+        """Save the files received from ``peer_email``; it replaces the previous
+        checkpoint."""
+        self.connections.watcher_store_con.put_checkpoint(
+            peer_email, checkpoint.as_compressed_data()
+        )
+
+    def watcher_load_checkpoint(self, peer_email: str) -> Checkpoint | None:
+        """The saved checkpoint of ``peer_email``, or None if there is none.
+
+        Raises:
+            ValueError: if a later client saved it.
+        """
+        data = self.connections.watcher_store_con.get_checkpoint(peer_email)
+        if data is None:
+            return None
+        return Checkpoint.from_compressed_data(data)
+
+    def watcher_save_rolling_state(
+        self, peer_email: str, rolling_state: RollingState
+    ) -> None:
+        """Save the changes received from ``peer_email`` since its checkpoint,
+        and the ids of the messages received; it replaces the previous rolling
+        state."""
+        self.connections.watcher_store_con.put_rolling_state(
+            peer_email, rolling_state.as_compressed_data()
+        )
+
+    def watcher_load_rolling_state(self, peer_email: str) -> RollingState | None:
+        """The saved rolling state of ``peer_email``, or None if there is none.
+
+        Raises:
+            ValueError: if a later client saved it.
+        """
+        data = self.connections.watcher_store_con.get_rolling_state(peer_email)
+        if data is None:
+            return None
+        return RollingState.from_compressed_data(data)
+
+    def watcher_peers_with_saved_state(self) -> list[str]:
+        """The DOs whose checkpoint or rolling state the DS saved."""
+        return self.connections.watcher_store_con.peers_with_saved_state()

syft/sync/connections/drive/drive_connections.py

@@ -1,33 +1,13 @@
 """The router's connections on Google Drive.
 
 Each connection has its own module: ``drive_p2p``, ``drive_owner_store``,
-``drive_collections`` and ``drive_peering``. They share one ``DriveCachedClient``,
+``drive_watcher_store``, ``drive_collections`` and ``drive_peering``. They share one ``DriveCachedClient``,
 which holds the Drive client. ``SyftboxManager._drive_connections`` builds them.
 """
 
 from syft.sync.connections.drive.drive_owner_store import DriveOwnerStoreConnection
 from syft.sync.connections.drive.drive_cached_client import DriveCachedClient
-from syft.sync.connections.protocols import Connections, WatcherStoreConnection
-
-
-class DriveWatcherStoreConnection(WatcherStoreConnection):
-    """The DS's storage on Drive. Not built yet: today the DS keeps what it
-    received in a local cache."""
-
-    def put_checkpoint(self, peer: str, name: str, data: bytes) -> str:
-        raise NotImplementedError("Watcher storage on Drive is not built yet")
-
-    def latest_checkpoint(self, peer: str) -> bytes | None:
-        raise NotImplementedError("Watcher storage on Drive is not built yet")
-
-    def put_rolling_state(self, peer: str, name: str, data: bytes) -> str:
-        raise NotImplementedError("Watcher storage on Drive is not built yet")
-
-    def get_rolling_state(self, peer: str) -> bytes | None:
-        raise NotImplementedError("Watcher storage on Drive is not built yet")
-
-    def reset(self, peer: str) -> None:
-        raise NotImplementedError("Watcher storage on Drive is not built yet")
+from syft.sync.connections.protocols import Connections
 
 
 def drive_cached_client_of(connections: Connections) -> DriveCachedClient:

syft/sync/connections/drive/drive_watcher_store.py

@@ -0,0 +1,130 @@
+"""The DS's own storage on Google Drive.
+
+The folder /SyftBox/{email}-{version}-watcher-state is shared with nobody. It
+holds one folder per DO, and each DO folder holds one checkpoint file and one
+rolling-state file. A save replaces the content of the file in place, so a
+read gets the old content or the new content, never a mix.
+"""
+
+from syft.sync.connections.drive.drive_cached_client import DriveCachedClient
+from syft.sync.connections.drive.gdrive_retry import (
+    execute_with_retries,
+    raise_unavailable,
+)
+from syft.sync.connections.protocols import WatcherStoreConnection
+from syft.version import SYFT_VERSION
+
+WATCHER_STATE_FOLDER_SUFFIX = "-watcher-state"
+CHECKPOINT_FILE_NAME = "checkpoint.tar.gz"
+ROLLING_STATE_FILE_NAME = "rolling_state.tar.gz"
+
+
+class DriveWatcherStoreConnection(WatcherStoreConnection):
+    """The DS's folder in /SyftBox: one checkpoint and one rolling state per DO."""
+
+    def __init__(self, cached_client: DriveCachedClient):
+        self._cached_client = cached_client
+
+    def clone_for_thread(self) -> "DriveWatcherStoreConnection":
+        """A connection over a new Drive client, for use from another thread."""
+        return DriveWatcherStoreConnection(self._cached_client.clone_for_thread())
+
+    @raise_unavailable("Could not save the checkpoint to Drive.")
+    def put_checkpoint(self, peer: str, data: bytes) -> None:
+        self._save_file(peer, CHECKPOINT_FILE_NAME, data)
+
+    @raise_unavailable("Could not read the checkpoint from Drive.")
+    def get_checkpoint(self, peer: str) -> bytes | None:
+        return self._read_file(peer, CHECKPOINT_FILE_NAME)
+
+    @raise_unavailable("Could not save the rolling state to Drive.")
+    def put_rolling_state(self, peer: str, data: bytes) -> None:
+        self._save_file(peer, ROLLING_STATE_FILE_NAME, data)
+
+    @raise_unavailable("Could not read the rolling state from Drive.")
+    def get_rolling_state(self, peer: str) -> bytes | None:
+        return self._read_file(peer, ROLLING_STATE_FILE_NAME)
+
+    @raise_unavailable("Could not remove the watcher state from Drive.")
+    def reset(self, peer: str) -> None:
+        peer_folder_id = self._peer_folder_id(peer)
+        if peer_folder_id is not None:
+            self._cached_client.client.delete_file_by_id(peer_folder_id)
+
+    @raise_unavailable("Could not list the watcher state on Drive.")
+    def peers_with_saved_state(self) -> list[str]:
+        watcher_folder_id = self._watcher_folder_id()
+        if watcher_folder_id is None:
+            return []
+        peer_folders = self._cached_client.client.find_folders(
+            parent_id=watcher_folder_id
+        )
+        return [name for _, name in peer_folders]
+
+    def _save_file(self, peer: str, name: str, data: bytes) -> None:
+        client = self._cached_client.client
+        watcher_folder_id = self._watcher_folder_id() or self._create_watcher_folder()
+        peer_folder_id = self._find_id(peer, watcher_folder_id) or client.create_folder(
+            peer, watcher_folder_id
+        )
+        payload, _ = client.create_file_payload(data)
+        file_id = self._find_id(name, peer_folder_id)
+        if file_id is None:
+            body = {"name": name, "parents": [peer_folder_id]}
+            client.create_syft_file(body, payload)
+        else:
+            execute_with_retries(
+                client.drive_service.files().update(fileId=file_id, media_body=payload)
+            )
+
+    def _read_file(self, peer: str, name: str) -> bytes | None:
+        peer_folder_id = self._peer_folder_id(peer)
+        if peer_folder_id is None:
+            return None
+        file_id = self._find_id(name, peer_folder_id)
+        if file_id is None:
+            return None
+        return self._cached_client.client.download_file(file_id)
+
+    def _peer_folder_id(self, peer: str) -> str | None:
+        watcher_folder_id = self._watcher_folder_id()
+        if watcher_folder_id is None:
+            return None
+        return self._find_id(peer, watcher_folder_id)
+
+    def _find_id(self, name: str, parent_id: str) -> str | None:
+        """The id of the file or folder ``name`` in ``parent_id``, or None.
+
+        Raises:
+            RuntimeError: if the folder holds more than one of that name.
+        """
+        query = f"name='{name}' and '{parent_id}' in parents and trashed=false"
+        files = self._cached_client.client.list_all_files(query, fields="id")
+        if len(files) > 1:
+            raise RuntimeError(
+                f"Found {len(files)} files named {name} in one folder of the "
+                "watcher state on Drive. Exactly one is expected."
+            )
+        return files[0]["id"] if files else None
+
+    def _watcher_folder_name(self) -> str:
+        return (
+            f"{self._cached_client.email}-{SYFT_VERSION}{WATCHER_STATE_FOLDER_SUFFIX}"
+        )
+
+    def _watcher_folder_id(self) -> str | None:
+        folders = self._cached_client.client.find_folders(
+            name_contains=[
+                f"{self._cached_client.email}-",
+                WATCHER_STATE_FOLDER_SUFFIX,
+            ],
+            parent_id=self._cached_client.get_syftbox_folder_id(),
+        )
+        return self._cached_client.find_or_adopt_versioned_folder(
+            folders, current_name=self._watcher_folder_name()
+        )
+
+    def _create_watcher_folder(self) -> str:
+        return self._cached_client.client.create_folder(
+            self._watcher_folder_name(), self._cached_client.get_syftbox_folder_id()
+        )