diff --git a/syft/sync/checkpoints/checkpoint.py b/syft/sync/checkpoints/checkpoint.py index d7abb481bb9..b949845b5ba 100644 --- a/syft/sync/checkpoints/checkpoint.py +++ b/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 diff --git a/syft/sync/checkpoints/rolling_state.py b/syft/sync/checkpoints/rolling_state.py index 145c9c8d725..6ebeb6869a9 100644 --- a/syft/sync/checkpoints/rolling_state.py +++ b/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: diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 851bf4ea828..6d3b22d449e 100644 --- a/syft/sync/connections/connection_router.py +++ b/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() diff --git a/syft/sync/connections/drive/drive_connections.py b/syft/sync/connections/drive/drive_connections.py index 12b785b96e8..0fe4e0758b7 100644 --- a/syft/sync/connections/drive/drive_connections.py +++ b/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: diff --git a/syft/sync/connections/drive/drive_watcher_store.py b/syft/sync/connections/drive/drive_watcher_store.py new file mode 100644 index 00000000000..fcf90a44b79 --- /dev/null +++ b/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() + ) diff --git a/syft/sync/connections/protocols.py b/syft/sync/connections/protocols.py index f833cbeda43..25a0497dad3 100644 --- a/syft/sync/connections/protocols.py +++ b/syft/sync/connections/protocols.py @@ -204,18 +204,18 @@ class WatcherStoreConnection(RouterConnection, ABC): """ @abstractmethod - def put_checkpoint(self, peer: str, name: str, data: bytes) -> str: + def put_checkpoint(self, peer: str, data: bytes) -> None: """Store a snapshot of every file received from ``peer``; it replaces the previous one.""" @abstractmethod - def latest_checkpoint(self, peer: str) -> bytes | None: - """The newest snapshot of files received from ``peer``, or None.""" + def get_checkpoint(self, peer: str) -> bytes | None: + """The snapshot of files received from ``peer``, or None.""" @abstractmethod - def put_rolling_state(self, peer: str, name: str, data: bytes) -> str: + def put_rolling_state(self, peer: str, data: bytes) -> None: """Store the changes received from ``peer`` since the snapshot, and the - ids of the messages already processed.""" + ids of the messages already processed; it replaces the previous one.""" @abstractmethod def get_rolling_state(self, peer: str) -> bytes | None: @@ -226,6 +226,10 @@ def reset(self, peer: str) -> None: """Forget everything received from ``peer``, before asking it to send everything again.""" + @abstractmethod + def peers_with_saved_state(self) -> list[str]: + """The peers that have a stored snapshot or rolling state.""" + class CollectionsConnection(RouterConnection, ABC): """Sets of files a data owner shares with named readers (datasets).""" diff --git a/syft/sync/events/file_change_event.py b/syft/sync/events/file_change_event.py index 13ee2644621..f4f30aa383b 100644 --- a/syft/sync/events/file_change_event.py +++ b/syft/sync/events/file_change_event.py @@ -1,7 +1,6 @@ from typing import Any, List, Literal from pathlib import Path from uuid import UUID, uuid4 -import base64 import json from pydantic import ( BaseModel, @@ -17,6 +16,11 @@ from syft.sync.utils.syftbox_utils import create_event_timestamp from syft.sync.utils.syftbox_utils import compress_data from syft.sync.utils.syftbox_utils import uncompress_data +from syft.sync.utils.file_content import ( + content_to_json, + content_type_of, + decode_content, +) FILE_CHANGE_FILENAME_PREFIX = "syfteventsmessagev3" @@ -74,21 +78,12 @@ class FileChangeEventV1(BaseModel): @property def content_type(self) -> Literal["text", "binary"] | None: """Computed field that stores the content type for proper deserialization.""" - if self.content is None: - return None - elif isinstance(self.content, bytes): - return "binary" - else: - return "text" + return content_type_of(self.content) @field_serializer("content", when_used="json") def serialize_content(self, value: str | bytes | None) -> str | None: """Serialize bytes as base64-encoded string for JSON.""" - if value is None: - return None - if isinstance(value, bytes): - return base64.b64encode(value).decode("utf-8") - return value + return content_to_json(value) @property def path_in_syftbox(self) -> Path: @@ -97,17 +92,7 @@ def path_in_syftbox(self) -> Path: @model_validator(mode="before") @classmethod def pre_init(cls, data: dict[str, Any]) -> dict[str, Any]: - # Deserialize content based on content_type metadata - content_type = data.get("content_type", None) - content = data.get("content") - - if content is not None and isinstance(content, str): - if content_type == "binary": - # Definitively decode base64 since we know it was binary - data["content"] = base64.b64decode(content) - # If content_type is "text" or None (legacy), keep as string - - return data + return decode_content(data) def eventfile_filepath(self) -> str: # TODO: remove diff --git a/syft/sync/messages/proposed_filechange.py b/syft/sync/messages/proposed_filechange.py index 6d0b7b64b60..7057a7cd6a7 100644 --- a/syft/sync/messages/proposed_filechange.py +++ b/syft/sync/messages/proposed_filechange.py @@ -4,7 +4,6 @@ import json import uuid import time -import base64 from pydantic import Field, model_validator, field_serializer, computed_field from pydantic.main import BaseModel from syft_migration import MigratableObject @@ -13,6 +12,11 @@ from syft.sync.utils.syftbox_utils import compress_data, uncompress_data from syft.sync.utils.syftbox_utils import create_event_timestamp from syft.sync.utils.syftbox_utils import get_event_hash_from_content +from syft.sync.utils.file_content import ( + content_to_json, + content_type_of, + decode_content, +) MESSAGE_FILENAME_PREFIX = "msgv2" @@ -36,34 +40,17 @@ class ProposedFileChangeV1(BaseModel): @property def content_type(self) -> Literal["text", "binary"] | None: """Computed field that stores the content type for proper deserialization.""" - if self.content is None: - return None - elif isinstance(self.content, bytes): - return "binary" - else: - return "text" + return content_type_of(self.content) @field_serializer("content", when_used="json") def serialize_content(self, value: str | bytes | None) -> str | None: """Serialize bytes as base64-encoded string for JSON.""" - if value is None: - return None - if isinstance(value, bytes): - return base64.b64encode(value).decode("utf-8") - return value + return content_to_json(value) @model_validator(mode="before") @classmethod def pre_init(cls, data: dict[str, Any]) -> dict[str, Any]: - # Deserialize content based on content_type metadata - content_type = data.get("content_type") - content = data.get("content") - - if content is not None and isinstance(content, str): - if content_type == "binary": - # Definitively decode base64 since we know it was binary - data["content"] = base64.b64decode(content) - # If content_type is "text" or None (legacy), keep as string + data = decode_content(data) # Generate hash if needed if "new_hash" not in data and not data.get("is_deleted", False): diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index 9a36f0d6363..81488eee766 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -11,17 +11,18 @@ from pydantic import BaseModel, ConfigDict, PrivateAttr from syft.sync.callback_mixin import BaseModelCallbackMixin +from syft.sync.checkpoints.checkpoint import DEFAULT_CHECKPOINT_EVENT_THRESHOLD from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.connections.drive import mock_drive_service from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.drive_collections import DriveCollectionsConnection -from syft.sync.connections.drive.drive_connections import ( - DriveWatcherStoreConnection, - drive_cached_client_of, -) +from syft.sync.connections.drive.drive_connections import drive_cached_client_of from syft.sync.connections.drive.drive_owner_store import DriveOwnerStoreConnection from syft.sync.connections.drive.drive_p2p import DriveP2PConnection from syft.sync.connections.drive.drive_peering import DrivePeeringConnection +from syft.sync.connections.drive.drive_watcher_store import ( + DriveWatcherStoreConnection, +) from syft.sync.connections.drive.grdrive_config import GdriveConnectionConfig from syft.sync.events.file_change_event import ( FileChangeEvent, @@ -620,7 +621,7 @@ def _drive_connections(cached_client: DriveCachedClient) -> Connections: return Connections( p2p_con=DriveP2PConnection(cached_client), owner_store_con=DriveOwnerStoreConnection(cached_client), - watcher_store_con=DriveWatcherStoreConnection(), + watcher_store_con=DriveWatcherStoreConnection(cached_client), collections_con=DriveCollectionsConnection(cached_client), peering_con=DrivePeeringConnection(cached_client), platform=GdriveFilesPlatform(), @@ -1074,7 +1075,7 @@ def sync_file_lock(self): def sync( self, auto_checkpoint: bool = True, - checkpoint_threshold: int = 50, + checkpoint_threshold: int = DEFAULT_CHECKPOINT_EVENT_THRESHOLD, auto_compact: bool = True, compact_threshold: int = MIN_MESSAGES_COMPACT, force_download_peer_state: bool = False, @@ -1344,7 +1345,6 @@ def _delete_local_dirs(self): dirs_to_delete = [ self.syftbox_folder, # main syftbox folder (datasets, private, etc.) syftbox_parent / f"{syftbox_name}-events", # DO event cache - syftbox_parent / f"{syftbox_name}-event-messages", # DS event cache ] for d in dirs_to_delete: if d.exists(): @@ -1397,7 +1397,9 @@ def compact_outboxes_if_needed( for peer in self.peer_manager.approved_peers } - def should_create_checkpoint(self, threshold: int = 50) -> bool: + def should_create_checkpoint( + self, threshold: int = DEFAULT_CHECKPOINT_EVENT_THRESHOLD + ) -> bool: """ Check if a checkpoint should be created based on event count. @@ -1411,7 +1413,9 @@ def should_create_checkpoint(self, threshold: int = 50) -> bool: return False return self.datasite_owner_syncer.should_create_checkpoint(threshold) - def try_create_checkpoint(self, threshold: int = 50): + def try_create_checkpoint( + self, threshold: int = DEFAULT_CHECKPOINT_EVENT_THRESHOLD + ): """ Try to create a checkpoint if the event count exceeds the threshold. diff --git a/syft/sync/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 9255d7abf91..09cd3c803d7 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -1,18 +1,26 @@ import logging from concurrent.futures import ThreadPoolExecutor -from typing import Callable, Dict, List, Set -from syft.sync.sync.caches.cache_file_writer_connection import FSFileConnection +from datetime import datetime, timedelta from pathlib import Path +from typing import Callable, Dict, List, Set + from pydantic import BaseModel, Field -from datetime import datetime, timedelta + +from syft.sync.checkpoints.checkpoint import ( + DEFAULT_CHECKPOINT_EVENT_THRESHOLD, + Checkpoint, + CheckpointFile, +) +from syft.sync.checkpoints.rolling_state import RollingState +from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.events.file_change_event import ( FileChangeEvent, FileChangeEventsMessage, FileChangeEventsMessageFileName, ) -from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.sync.caches.cache_file_writer_connection import ( CacheFileConnection, + FSFileConnection, InMemoryCacheFileConnection, ) from syft.sync.sync.collection_spec import CollectionSyncSpec @@ -43,20 +51,33 @@ def _sent_at(msg_id: str) -> float: return FileChangeEventsMessageFileName.from_string(msg_id).timestamp -def _events_within_peer_datasite( +def _is_complete(event: FileChangeEvent) -> bool: + """Whether ``event`` is a delete, or has the content and the hash of its file.""" + return event.is_deleted or ( + event.content is not None and event.new_hash is not None + ) + + +def _accepted_events( events: List[FileChangeEvent], peer_email: str ) -> List[FileChangeEvent]: - """The events of ``events`` that target ``peer_email``'s datasite; log the rest.""" + """The events of ``events`` that target ``peer_email``'s datasite and are + complete; log the rest.""" accepted = [] for event in events: - if _is_within_peer_datasite(event, peer_email): - accepted.append(event) - else: + if not _is_within_peer_datasite(event, peer_email): logger.warning( f"Dropping event from {peer_email} aimed at " f"{event.path_in_syftbox}: a peer may only change files " "under its own datasite" ) + elif not _is_complete(event): + logger.warning( + f"Dropping event from {peer_email} for {event.path_in_syftbox}: " + "it changes the file but has no content or no hash" + ) + else: + accepted.append(event) return accepted @@ -67,7 +88,6 @@ class DataSiteWatcherCacheConfig(BaseModel): email: str = "" use_in_memory_cache: bool = True syftbox_folder: Path | None = None - events_base_path: Path | None = None # Collections to sync down (prefix + local subpath per collection type) collection_specs: list[CollectionSyncSpec] = [] @@ -90,9 +110,16 @@ class DataSiteWatcherCache(BaseModel): # The msg_id of the newest events message received from each peer, keyed # by ``_peer_key``. last_msg_id_per_peer: Dict[str, str] = {} - # The msg_id of each events message received. The connection lists back a - # bit before the last one for late messages, and skips these. - received_msg_ids: Set[str] = set() + # The msg_ids of the events messages received from each peer, keyed by + # ``_peer_key``. The connection lists back a bit before the last one for + # late messages, and skips these. + received_msg_ids_per_peer: Dict[str, Set[str]] = {} + # Whether the state received from each peer is saved in the DS's storage. + saves_watcher_state: bool + # The changes from each peer since its checkpoint, keyed by ``_peer_key``. + # The received ids stay in ``received_msg_ids_per_peer``; a saved copy + # carries them. + rolling_state_per_peer: Dict[str, RollingState] = {} # Base syftbox folder syftbox_folder: Path | None = None # Collections to sync down (prefix + local subpath per collection type) @@ -118,60 +145,24 @@ def from_config( connection_router=connection_router, syftbox_folder=config.syftbox_folder, collection_specs=config.collection_specs, + saves_watcher_state=False, ) return res else: if config.syftbox_folder is None: raise ValueError("syftbox_folder is required for non-in-memory cache") - syftbox_folder_name = Path(config.syftbox_folder).name - syftbox_parent = Path(config.syftbox_folder).parent - events_folder = syftbox_parent / f"{syftbox_folder_name}-event-messages" - cache = cls( - events_connection=FSFileConnection( - base_dir=events_folder, dtype=FileChangeEventsMessage - ), file_connection=FSFileConnection(base_dir=config.syftbox_folder), connection_router=connection_router, syftbox_folder=config.syftbox_folder, collection_specs=config.collection_specs, + saves_watcher_state=True, ) - cache._load_cached_state() + cache._load_collection_hashes_from_disk() + cache._restore_watcher_state() return cache - def _load_cached_state(self): - """Load cached state from disk: file hashes, last messages, and collection hashes.""" - self._load_saved_messages() - self._load_collection_hashes_from_disk() - - def _load_saved_messages(self): - """Restore the received messages, last messages and file hashes from - the saved events messages, oldest first.""" - try: - saved_messages = self.events_connection.get_all() - except Exception: - saved_messages = [] - for events_message in sorted(saved_messages, key=lambda m: m.timestamp): - self._restore_received(events_message) - self._restore_file_hashes(events_message) - - def _restore_received(self, events_message: FileChangeEventsMessage) -> None: - """Record a saved message as received. It does not store its sender, but - its events name it: only events for the sender's datasite are saved.""" - msg_id = events_message.message_filepath.as_string() - self.received_msg_ids.add(msg_id) - if events_message.events: - self._record_received(events_message.events[0].datasite_email, msg_id) - - def _restore_file_hashes(self, events_message: FileChangeEventsMessage) -> None: - for event in events_message.events: - path_key = Path(event.path_in_syftbox) - if event.is_deleted: - self.file_hashes.pop(path_key, None) - else: - self.file_hashes[path_key] = event.new_hash - def _load_collection_hashes_from_disk(self): """Scan local collection directories and compute hashes to populate collection_hashes.""" for collection_path in self._get_local_collection_folders(): @@ -226,7 +217,8 @@ def clear_cache(self): self.peers = [] self.current_check_point = None self.last_msg_id_per_peer = {} - self.received_msg_ids = set() + self.received_msg_ids_per_peer = {} + self.rolling_state_per_peer = {} self.collection_hashes = {} @property @@ -239,11 +231,15 @@ def sync_down(self, peer_email: str) -> int: """Apply the events messages of ``peer_email`` not received yet: those after the last one, and any that landed late. Returns the number of events applied. A message that can never be read is passed over, so it - does not hold back the later ones.""" + does not hold back the later ones. + + A cache that saves its state saves the state of ``peer_email`` before + it marks the messages processed.""" + key = _peer_key(peer_email) received = self.connection_router.watcher_receive_events_messages( peer_email, - after=self.last_msg_id_per_peer.get(_peer_key(peer_email)), - already_received=self.received_msg_ids, + after=self.last_msg_id_per_peer.get(key), + already_received=self.received_msg_ids_per_peer.get(key, frozenset()), ) event_count = 0 for raw, event_message in received: @@ -251,6 +247,8 @@ def sync_down(self, peer_email: str) -> int: self.apply_event_message(event_message, peer_email) event_count += len(event_message.events) self._record_received(peer_email, raw.msg_id) + if received and self.saves_watcher_state: + self._save_watcher_state(peer_email) self.connection_router.mark_all_processed([raw for raw, _ in received]) self.last_sync = datetime.now() @@ -259,8 +257,8 @@ def sync_down(self, peer_email: str) -> int: def _record_received(self, peer_email: str, msg_id: str) -> None: """Remember ``msg_id``, and make it ``peer_email``'s last message when it is newer. A message that landed late is older than the last one.""" - self.received_msg_ids.add(msg_id) key = _peer_key(peer_email) + self.received_msg_ids_per_peer.setdefault(key, set()).add(msg_id) last = self.last_msg_id_per_peer.get(key) if last is None or _sent_at(msg_id) > _sent_at(last): self.last_msg_id_per_peer[key] = msg_id @@ -274,23 +272,30 @@ def apply_event_message( ``path_in_datasite``), and both fields come from the peer. A peer may only change files under its own datasite, mirrored locally at ``//``, so an event aimed anywhere else is dropped - before it reaches the local file connection. + before it reaches the local file connection. A change without its + content or its hash is dropped too. - The saved copy of the message holds only the accepted events. - ``_load_saved_messages`` reads saved messages back on start, - and a saved message does not record its sender, so a dropped event - cannot be checked again there. + The record of the message and the saved state of ``peer_email`` hold + only the accepted events, so a restore never writes a dropped event. """ - accepted = _events_within_peer_datasite(event_message.events, peer_email) + accepted = _accepted_events(event_message.events, peer_email) saved = event_message.model_copy(update={"events": accepted}) self.events_connection.write_file(saved.message_filepath.as_string(), saved) for event in accepted: + if self._write_denied(event): + continue self._apply_event(event) + if self.saves_watcher_state: + self._rolling_state(peer_email).add_event(event) + + def _write_denied(self, event: FileChangeEvent) -> bool: + """Whether ``pre_write_filter`` refuses ``event``.""" + if self.pre_write_filter is None: + return False + return not self.pre_write_filter(str(event.path_in_syftbox), event.is_deleted) def _apply_event(self, event: FileChangeEvent) -> None: path = str(event.path_in_syftbox) - if self.pre_write_filter and not self.pre_write_filter(path, event.is_deleted): - return # Normalize path to Path object for consistency in file_hashes dict path_key = Path(event.path_in_syftbox) if event.is_deleted: @@ -300,6 +305,81 @@ def _apply_event(self, event: FileChangeEvent) -> None: self.file_connection.write_file(path, event.content) self.file_hashes[path_key] = event.new_hash + # --- State saved in the DS's storage --- + + def _restore_watcher_state(self) -> None: + """Restore the files, their hashes and the received message ids of each + peer with saved state. A Colab runtime reset loses the local copy. + + The state of a peer is saved under its ``_peer_key``, so two spellings + of one email share it.""" + for peer_key in self.connection_router.watcher_peers_with_saved_state(): + checkpoint = self.connection_router.watcher_load_checkpoint(peer_key) + if checkpoint is not None: + for file in checkpoint.files: + self._restore_file(checkpoint.email, file) + rolling_state = self.connection_router.watcher_load_rolling_state(peer_key) + if rolling_state is not None: + for event in rolling_state.events: + self._apply_event(event) + for msg_id in rolling_state.received_msg_ids: + self._record_received(peer_key, msg_id) + self.rolling_state_per_peer[peer_key] = rolling_state.model_copy( + update={"received_msg_ids": []} + ) + + def _restore_file(self, datasite_email: str, file: CheckpointFile) -> None: + path_in_syftbox = Path(datasite_email) / file.path + self.file_connection.write_file(str(path_in_syftbox), file.content) + self.file_hashes[path_in_syftbox] = file.hash + + def _save_watcher_state(self, peer_email: str) -> None: + """Save the rolling state of ``peer_email``. At + ``DEFAULT_CHECKPOINT_EVENT_THRESHOLD`` paths, fold it into the checkpoint.""" + self._save_rolling_state(peer_email) + if ( + self._rolling_state(peer_email).event_count + >= DEFAULT_CHECKPOINT_EVENT_THRESHOLD + ): + self._fold_into_checkpoint(peer_email) + + def _fold_into_checkpoint(self, peer_email: str) -> None: + """Merge the changes since the checkpoint into it, then save an empty + rolling state. A deleted path leaves the checkpoint.""" + key = _peer_key(peer_email) + rolling_state = self._rolling_state(peer_email) + checkpoint = self.connection_router.watcher_load_checkpoint(key) + if checkpoint is None: + # Live events write under the email as the DO's events spell it, so + # a restore writes the checkpoint files there too. + checkpoint = Checkpoint(email=rolling_state.events[-1].datasite_email) + folded = checkpoint.with_events(rolling_state.events) + # The checkpoint goes first. After a stop between the two saves, the old + # rolling state only repeats changes that the checkpoint has. + self.connection_router.watcher_save_checkpoint(key, folded) + self.rolling_state_per_peer[key] = RollingState( + email=peer_email, base_checkpoint_timestamp=folded.timestamp + ) + self._save_rolling_state(peer_email) + + def _save_rolling_state(self, peer_email: str) -> None: + key = _peer_key(peer_email) + received_msg_ids = sorted(self.received_msg_ids_per_peer[key]) + saved = self._rolling_state(peer_email).model_copy( + update={"received_msg_ids": received_msg_ids} + ) + self.connection_router.watcher_save_rolling_state(key, saved) + + def _rolling_state(self, peer_email: str) -> RollingState: + """The changes from ``peer_email`` since its checkpoint.""" + key = _peer_key(peer_email) + if key not in self.rolling_state_per_peer: + # 0.0 means no checkpoint yet, as for the DO's own rolling state. + self.rolling_state_per_peer[key] = RollingState( + email=peer_email, base_checkpoint_timestamp=0.0 + ) + return self.rolling_state_per_peer[key] + def get_cached_events(self) -> List[FileChangeEvent]: messages = self.events_connection.get_all() return [event for message in messages for event in message.events] diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 3ffad11d4e7..1f058ebb9fc 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -13,11 +13,10 @@ from syft.sync.callback_mixin import BaseModelCallbackMixin from syft.sync.checkpoints.checkpoint import ( + DEFAULT_CHECKPOINT_EVENT_THRESHOLD, DEFAULT_COMPACTING_THRESHOLD, Checkpoint, - CheckpointFile, IncrementalCheckpoint, - compact_incremental_checkpoints, ) from syft.sync.checkpoints.rolling_state import ( RollingState, @@ -51,9 +50,6 @@ logger = logging.getLogger(__name__) -# Default threshold for creating incremental checkpoint from rolling state -DEFAULT_CHECKPOINT_EVENT_THRESHOLD = 50 - # save_owner_state() uploads the rolling state once this many events were added # since the last upload. At 1 the rolling state on Drive holds every change after # each save_owner_state(), so the owner's event log is not written. @@ -1267,67 +1263,18 @@ def compact_checkpoints(self) -> Checkpoint: print(f"Found {len(incremental_cps)} incremental checkpoints to compact") - # Start with existing checkpoint files (if any) - if existing_checkpoint: + incrementals_in_order = sorted(incremental_cps, key=lambda c: c.sequence_number) + events = [ + event for checkpoint in incrementals_in_order for event in checkpoint.events + ] + if existing_checkpoint is None: + compacted = Checkpoint(email=self.email).with_events(events) + else: print( f"Merging with existing full checkpoint " f"({len(existing_checkpoint.files)} files)..." ) - # Convert existing checkpoint files to events for merging - merged_events: dict[str, FileChangeEvent] = {} - - # Add existing checkpoint files as events - for checkpoint_file in existing_checkpoint.files: - # Create a FileChangeEvent from CheckpointFile - event = FileChangeEvent( - id=uuid4(), - path_in_datasite=Path(checkpoint_file.path), - datasite_email=self.email, - content=checkpoint_file.content, - old_hash=None, - new_hash=checkpoint_file.hash, - is_deleted=False, - submitted_timestamp=existing_checkpoint.timestamp, - timestamp=existing_checkpoint.timestamp, - ) - merged_events[checkpoint_file.path] = event - - # Merge incremental checkpoints on top (later events overwrite) - for inc_cp in sorted(incremental_cps, key=lambda c: c.sequence_number): - for event in inc_cp.events: - merged_events[str(event.path_in_datasite)] = event - - # Create full checkpoint from merged events - 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, - ) - ) - - # Find latest timestamp - last_event_timestamp = existing_checkpoint.last_event_timestamp - for inc_cp in incremental_cps: - for event in inc_cp.events: - if event.timestamp is not None: - if ( - last_event_timestamp is None - or event.timestamp > last_event_timestamp - ): - last_event_timestamp = event.timestamp - - compacted = Checkpoint( - email=self.email, - files=files, - last_event_timestamp=last_event_timestamp, - ) - else: - # No existing checkpoint, just merge incremental checkpoints - compacted = compact_incremental_checkpoints(incremental_cps, self.email) + compacted = existing_checkpoint.with_events(events) # Upload the compacted checkpoint (this deletes old full checkpoints) print(f"Uploading compacted checkpoint with {len(compacted.files)} files...") diff --git a/syft/sync/utils/file_content.py b/syft/sync/utils/file_content.py new file mode 100644 index 00000000000..fd994a1b3c2 --- /dev/null +++ b/syft/sync/utils/file_content.py @@ -0,0 +1,40 @@ +"""File content in a model: text as it is, bytes as base64 in JSON. + +A model keeps ``content`` next to a computed ``content_type``, so that a read +gives bytes back as bytes. +""" + +import base64 +from typing import Any, Literal + +ContentType = Literal["text", "binary"] + + +def content_type_of(content: str | bytes | None) -> ContentType | None: + """``"binary"`` for bytes, ``"text"`` for a str, None for no content.""" + if content is None: + return None + return "binary" if isinstance(content, bytes) else "text" + + +def content_to_json(content: str | bytes | None) -> str | None: + """``content`` as JSON text: bytes go as base64.""" + if isinstance(content, bytes): + return base64.b64encode(content).decode("ascii") + return content + + +def decode_content(data: Any) -> Any: + """The input of a model, with its base64 ``content`` decoded to bytes when + its ``content_type`` is ``"binary"``. + + Raises: + ValueError: if the content is not valid base64. + """ + if ( + isinstance(data, dict) + and data.get("content_type") == "binary" + and isinstance(data.get("content"), str) + ): + return {**data, "content": base64.b64decode(data["content"], validate=True)} + return data diff --git a/syft/sync/utils/syftbox_utils.py b/syft/sync/utils/syftbox_utils.py index b96c34cf53f..cd540ff35b1 100644 --- a/syft/sync/utils/syftbox_utils.py +++ b/syft/sync/utils/syftbox_utils.py @@ -137,7 +137,6 @@ def _delete_local_syftbox_dirs(local_syftbox_path: Path, verbose: bool = True) - dirs_to_delete = [ local_syftbox_path, syftbox_parent / f"{syftbox_name}-events", - syftbox_parent / f"{syftbox_name}-event-messages", ] for d in dirs_to_delete: if d.exists(): @@ -153,8 +152,8 @@ def delete_local_syftbox( ) -> None: """Delete local SyftBox directories. - Deletes the SyftBox directory and its companion cache directories - (``-events``, ``-event-messages``). + Deletes the SyftBox directory and its companion cache directory + (``-events``). Args: email: Google account email. Used to resolve default path if diff --git a/tests/integration/with_unit_coverage/test_cache_aware_sync.py b/tests/integration/with_unit_coverage/test_cache_aware_sync.py index 3f1b703f273..8cec9240b69 100644 --- a/tests/integration/with_unit_coverage/test_cache_aware_sync.py +++ b/tests/integration/with_unit_coverage/test_cache_aware_sync.py @@ -125,11 +125,11 @@ def test_do_cache_aware_sync_gdrive(): @pytest.mark.flaky(reruns=3, reruns_delay=2) @pytest.mark.usefixtures("setup_delete_syftboxes") def test_ds_cache_aware_sync_gdrive(): - """Real GDrive test: Create DS, sync 3 files, create 2 more, restart with same cache, - verify cache has 3 events, sync and end up with 5. + """Real GDrive test: Create DS, sync 3 files, create 2 more, restart the DS, + verify its state is restored, sync and end up with 5. - This test verifies that when a DS manager is restarted with a persisted filesystem cache, - it uses the cached last_msg_id_per_peer to only download new events. + This test verifies that a restarted DS restores last_msg_id_per_peer from + its saved state, so it downloads only new events. """ # Create initial pair with filesystem cache ds_manager1, do_manager1 = ( @@ -197,17 +197,13 @@ def test_ds_cache_aware_sync_gdrive(): ) new_ds_manager = SyftboxManager.from_config(new_ds_config) - # Verify the new manager's cache loaded existing events from disk + # Verify the new manager restored its state from the DS's storage new_cache = new_ds_manager.datasite_watcher_syncer.datasite_watcher_cache - assert len(new_cache.events_connection) == initial_cache_count, ( - f"New manager cache should have loaded {initial_cache_count} events from disk, " - f"got {len(new_cache.events_connection)}" - ) assert len(new_cache.last_msg_id_per_peer) > 0, ( - "last_msg_id_per_peer should be pre-populated from disk cache" + "last_msg_id_per_peer should be restored from the saved state" ) assert len(new_cache.file_hashes) >= 3, ( - f"New manager should have at least 3 file hashes pre-populated, got {len(new_cache.file_hashes)}" + f"New manager should have at least 3 file hashes restored, got {len(new_cache.file_hashes)}" ) # Sync the new manager - should only download the new events diff --git a/tests/migrations/unit/test_file_change_events_serialization.py b/tests/migrations/unit/test_file_change_events_serialization.py index 751e6542a8f..f341de6e8c7 100644 --- a/tests/migrations/unit/test_file_change_events_serialization.py +++ b/tests/migrations/unit/test_file_change_events_serialization.py @@ -4,6 +4,8 @@ import json from uuid import uuid4 +import pytest + from syft.migrations import client_registry from syft.sync.events.file_change_event import ( FileChangeEvent, @@ -44,6 +46,14 @@ def test_round_trip_text_binary_and_deletion(): assert restored.message_filepath == original.message_filepath +def test_corrupt_binary_content_refused(): + data = json.loads(_make_event(b"\x00\xff").model_dump_json()) + data["content"] += "!" + + with pytest.raises(ValueError): + FileChangeEvent.model_validate(data) + + def test_identity_fields_on_the_wire(): # load_as_latest would setdefault them back, so the round-trip alone # cannot catch a silently broken serialization of the identity fields. diff --git a/tests/unit/test_cache_aware_sync.py b/tests/unit/test_cache_aware_sync.py index f9991f9c427..2e64ad09a93 100644 --- a/tests/unit/test_cache_aware_sync.py +++ b/tests/unit/test_cache_aware_sync.py @@ -17,9 +17,7 @@ ) from syft.sync.syftbox_manager import SyftboxManager from syft.sync.sync.collection_spec import CollectionSyncSpec - from tests.unit.utils import ( - router_with_mock_connection, TEST_COLLECTION_PREFIX, TEST_COLLECTION_SUBPATH, get_mock_events_messages, @@ -94,8 +92,8 @@ def counted_download(events_message_id): def test_ds_incremental_sync_downloads_only_new_events(monkeypatch): - """Pre-populate cache with first n events, create new manager with same cache dir, - sync with n+m in outbox, verify only m new events are downloaded.""" + """Sync n events, add m more to the outbox, restart the DS, and verify that + the restored state makes it download only the m new events.""" import time # Create initial pair with filesystem cache @@ -130,17 +128,14 @@ def test_ds_incremental_sync_downloads_only_new_events(monkeypatch): # Create a NEW SyftboxManager with the same cache directory (simulating restart) new_ds_manager = ds_manager._copy() - # Verify the new manager's cache loaded existing events from disk + # Verify the new manager restored its state from the DS's storage new_cache = new_ds_manager.datasite_watcher_syncer.datasite_watcher_cache - assert len(new_cache.events_connection) == initial_cache_count, ( - f"New manager cache should have loaded {initial_cache_count} events from disk" - ) assert ( new_cache.last_msg_id_per_peer[do_manager.email] == ds_cache.last_msg_id_per_peer[do_manager.email] - ), "last_msg_id_per_peer should be pre-populated from disk cache" + ), "last_msg_id_per_peer should be restored from the saved state" assert len(new_cache.file_hashes) == 3, ( - "file_hashes should be pre-populated from disk cache" + "file_hashes should be restored from the saved state" ) # Sync the new manager - should only download the 2 new events @@ -159,11 +154,6 @@ def counted_receive(*args, **kwargs): f"Should receive only 2 new messages, got {len(received)}" ) - # Verify total events in cache - assert len(new_cache.events_connection) == 5, ( - "Should have 5 total events after syncing new events" - ) - def test_do_cache_handles_deletions_correctly(): """Test that file_hashes properly reflects deletions when loading from cache.""" @@ -216,7 +206,8 @@ def test_do_cache_handles_deletions_correctly(): def test_ds_cache_handles_deletions_correctly(): - """Test that DS file_hashes properly reflects deletions when loading from cache.""" + """Test that DS file_hashes properly reflects deletions when restoring its + saved state.""" ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( use_in_memory_cache=False ) @@ -227,10 +218,11 @@ def test_ds_cache_handles_deletions_correctly(): ctx = do_manager.datasite_owner_syncer.perm_context ctx.open(".").grant_read_access(ds_manager.email) - # Create a file on DO side + # Create two files on DO side: one is deleted later, one is kept test_file = datasite_dir_do / "test_file.txt" test_file.parent.mkdir(parents=True, exist_ok=True) test_file.write_text("Hello") + (datasite_dir_do / "kept_file.txt").write_text("Kept") do_manager.sync() ds_manager.sync() @@ -266,13 +258,15 @@ def test_ds_cache_handles_deletions_correctly(): CollectionSyncSpec.public(TEST_COLLECTION_PREFIX, TEST_COLLECTION_SUBPATH) ], ) - new_cache = DataSiteWatcherCache.from_config( - config, router_with_mock_connection(do_manager.email) - ) + new_cache = DataSiteWatcherCache.from_config(config, ds_manager._connection_router) - # Verify deleted file is not in reloaded cache + # Verify the restore kept one file and not the deleted one + kept_path = Path(do_manager.email) / "kept_file.txt" + assert kept_path in new_cache.file_hashes, ( + "Kept file should be restored from the saved state" + ) assert expected_path not in new_cache.file_hashes, ( - "Deleted file should not appear in file_hashes after cache reload" + "Deleted file should not appear in file_hashes after the restore" ) diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index 8a6bbd84eec..a3183d55227 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -16,11 +16,12 @@ from syft.sync.connections.drive.drive_owner_store import DriveOwnerStoreConnection from syft.sync.checkpoints.checkpoint import ( CHECKPOINT_FILENAME_PREFIX, + CHECKPOINT_VERSION, INCREMENTAL_CHECKPOINT_PREFIX, Checkpoint, IncrementalCheckpoint, ) -from syft.sync.checkpoints.rolling_state import RollingState +from syft.sync.checkpoints.rolling_state import ROLLING_STATE_VERSION, RollingState from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.mock_drive_service import ( http_error, @@ -829,3 +830,29 @@ def test_incremental_left_after_compaction_does_not_revert_a_file(): compacted = syncer.compact_checkpoints() [file1] = [f for f in compacted.files if f.path == "file1.txt"] assert file1.content == "v2" + + +def test_checkpoint_of_later_version_refused(): + later_version = CHECKPOINT_VERSION + 1 + checkpoint = Checkpoint(email="do@example.com", version=later_version) + + with pytest.raises(ValueError, match=f"version {later_version}"): + Checkpoint.from_compressed_data(checkpoint.as_compressed_data()) + + +def test_rolling_state_of_later_version_refused(): + later_version = ROLLING_STATE_VERSION + 1 + state = RollingState( + email="do@example.com", base_checkpoint_timestamp=0.0, version=later_version + ) + + with pytest.raises(ValueError, match=f"version {later_version}"): + RollingState.from_compressed_data(state.as_compressed_data()) + + +def test_checkpoint_with_change_without_content_raises(): + checkpoint = Checkpoint(email="do@example.com") + change = get_mock_event("do@example.com/a.txt").model_copy(update={"content": None}) + + with pytest.raises(ValueError, match="no content or no hash"): + checkpoint.with_events([change]) diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 175b80a8ca4..02ab0373820 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -22,13 +22,15 @@ CollectionFolder, DriveCollectionsConnection, ) -from syft.sync.connections.drive.drive_connections import DriveWatcherStoreConnection from syft.sync.connections.drive.drive_owner_store import DriveOwnerStoreConnection from syft.sync.connections.drive.drive_p2p import ( OUTBOX_LOOKBACK_SECONDS, DriveP2PConnection, ) from syft.sync.connections.drive.drive_peering import DrivePeeringConnection +from syft.sync.connections.drive.drive_watcher_store import ( + DriveWatcherStoreConnection, +) from syft.sync.connections.drive.gdrive_client import GDriveClient from syft.sync.connections.drive.mock_drive_service import ( MockDeleteRequest, @@ -547,10 +549,32 @@ def test_owner_store_files_carry_the_syft_tag(pair): # -- WatcherStoreConnection ------------------------------------------------------- -def test_watcher_storage_on_drive_is_not_built_yet(pair): +def test_watcher_state_save_replaces_previous(pair): + _, ds_cached_client = pair + store = DriveWatcherStoreConnection(ds_cached_client) + assert store.get_checkpoint(DO) is None + + store.put_checkpoint(DO, b"checkpoint 1") + store.put_rolling_state(DO, b"rolling state") + store.put_checkpoint(DO, b"checkpoint 2") + + assert store.get_checkpoint(DO) == b"checkpoint 2" + assert store.get_rolling_state(DO) == b"rolling state" + + +def test_watcher_state_reset_keeps_other_peers(pair): _, ds_cached_client = pair - with pytest.raises(NotImplementedError): - DriveWatcherStoreConnection().latest_checkpoint(DO) + store = DriveWatcherStoreConnection(ds_cached_client) + other_do = "other-do@example.com" + store.put_rolling_state(DO, b"from do") + store.put_rolling_state(other_do, b"from other do") + assert set(store.peers_with_saved_state()) == {DO, other_do} + + store.reset(DO) + + assert store.get_rolling_state(DO) is None + assert store.get_rolling_state(other_do) == b"from other do" + assert store.peers_with_saved_state() == [other_do] # -- CollectionsConnection -------------------------------------------------------- diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index c35aef902dd..d89f2056702 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -5,6 +5,7 @@ import pytest +from syft.sync.connections.drive.drive_connections import drive_cached_client_of from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.peers.exceptions import SyftPeerNotReadyError from syft.sync.peers.peer import Peer, PeerState @@ -12,8 +13,6 @@ from syft.sync.syftbox_manager import SyftboxManager from tests.unit.test_sync_manager import path_for_job from tests.unit.utils import grant_job_inbox_access -from syft.sync.connections.drive.drive_connections import drive_cached_client_of - # ========================================================================= # PeerStore unit tests @@ -447,12 +446,13 @@ def test_encrypted_checkpoint_roundtrip(): def test_encrypted_rolling_state_roundtrip(): """Rolling state is encrypted at rest and decrypted on read.""" + from uuid import uuid4 + from syft.sync.checkpoints.rolling_state import RollingState from syft.sync.events.file_change_event import ( FileChangeEvent, FileChangeEventsMessage, ) - from uuid import uuid4 ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( encryption=True, diff --git a/tests/unit/test_shared_router.py b/tests/unit/test_shared_router.py index 6402c8f00fe..8f0d0cc2763 100644 --- a/tests/unit/test_shared_router.py +++ b/tests/unit/test_shared_router.py @@ -6,10 +6,7 @@ import pytest from syft.sync.connections.base_connection import ConnectionConfig -from syft.sync.connections.drive.drive_connections import ( - DriveWatcherStoreConnection, - drive_cached_client_of, -) +from syft.sync.connections.drive.drive_connections import drive_cached_client_of from syft.sync.connections.drive.drive_cached_client import DriveCachedClient from syft.sync.connections.drive.grdrive_config import GdriveConnectionConfig from syft.sync.connections.drive.mock_drive_service import pair_with_mock_service @@ -50,13 +47,8 @@ def test_every_part_of_manager_shares_one_router(): def _drive_cached_clients(connections) -> list[DriveCachedClient]: - """The cached client of each connection that uses Drive. The watcher store - does not use it yet.""" - return [ - connection._cached_client - for connection in connections.all() - if not isinstance(connection, DriveWatcherStoreConnection) - ] + """The cached client of each connection.""" + return [connection._cached_client for connection in connections.all()] def test_manager_builds_one_drive_connection_from_its_config(monkeypatch): diff --git a/tests/unit/test_watcher_event_datasite_check.py b/tests/unit/test_watcher_event_datasite_check.py index 0bda6c93916..326cf852505 100644 --- a/tests/unit/test_watcher_event_datasite_check.py +++ b/tests/unit/test_watcher_event_datasite_check.py @@ -122,23 +122,28 @@ def test_saved_message_holds_only_accepted_events(): assert saved.message_filepath == message.message_filepath -def test_dropped_event_is_not_restored_after_restart(tmp_path): - """On start, the watcher rebuilds its state from saved messages, with no sender to check.""" - config = DataSiteWatcherCacheConfig( - email=VICTIM, use_in_memory_cache=False, syftbox_folder=tmp_path / "SyftBox" +def test_dropped_event_is_not_restored_after_restart(): + """Over the mock Drive: on start, the DS restores its saved state, which + holds only the accepted events.""" + ds, do = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=False ) - forged = _event(VICTIM, "syft.pub.yaml") - legitimate = _event(PEER, "data/file.txt") - router = router_with_mock_connection(VICTIM) - DataSiteWatcherCache.from_config(config, router).apply_event_message( - FileChangeEventsMessage(events=[forged, legitimate]), PEER + forged = _event(ds.email, "syft.pub.yaml") + legitimate = _event(do.email, "data/file.txt") + do._connection_router.owner_write_event_messages_to_outbox( + ds.email, FileChangeEventsMessage(events=[forged, legitimate]) ) + ds.sync() - restarted = DataSiteWatcherCache.from_config(config, router) + restarted_syncer = ds._copy().datasite_watcher_syncer + assert restarted_syncer is not None + restarted = restarted_syncer.datasite_watcher_cache - assert Path(f"{VICTIM}/syft.pub.yaml") not in restarted.file_hashes - assert set(restarted.last_msg_id_per_peer) == {PEER} - assert restarted.file_hashes[Path(f"{PEER}/data/file.txt")] == legitimate.new_hash + assert Path(f"{ds.email}/syft.pub.yaml") not in restarted.file_hashes + assert set(restarted.last_msg_id_per_peer) == {do.email} + assert ( + restarted.file_hashes[Path(f"{do.email}/data/file.txt")] == legitimate.new_hash + ) def test_forged_outbox_message_cannot_rewrite_ds_root_permissions(): diff --git a/tests/unit/test_watcher_state.py b/tests/unit/test_watcher_state.py new file mode 100644 index 00000000000..e124cbd6b67 --- /dev/null +++ b/tests/unit/test_watcher_state.py @@ -0,0 +1,195 @@ +import time +import uuid + +from syft.sync.encryption_policy import EncryptionPolicy +from syft.sync.events.file_change_event import FileChangeEvent, FileChangeEventsMessage +from syft.sync.syftbox_manager import SyftboxManager, SyftboxManagerConfig +from syft.sync.sync.caches import datasite_watcher_cache +from syft.sync.sync.caches.datasite_watcher_cache import DataSiteWatcherCache +from tests.unit.utils import get_mock_event + + +def _pair_on_disk() -> tuple[SyftboxManager, SyftboxManager]: + return SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=False + ) + + +def _send_to_ds( + do_manager: SyftboxManager, + ds_manager: SyftboxManager, + events: list[FileChangeEvent], +) -> None: + do_manager._connection_router.owner_write_event_messages_to_outbox( + ds_manager.email, FileChangeEventsMessage(events=events) + ) + + +def _deletion_of(event: FileChangeEvent) -> FileChangeEvent: + return FileChangeEvent( + id=uuid.uuid4(), + path_in_datasite=event.path_in_datasite, + datasite_email=event.datasite_email, + old_hash=event.new_hash, + is_deleted=True, + submitted_timestamp=time.time(), + timestamp=time.time(), + ) + + +def _watcher_cache(manager: SyftboxManager) -> DataSiteWatcherCache: + syncer = manager.datasite_watcher_syncer + assert syncer is not None + return syncer.datasite_watcher_cache + + +def _restart_on_empty_folder(ds_manager: SyftboxManager, tmp_path) -> SyftboxManager: + """A new DS manager on the Drive of ``ds_manager`` and an empty folder, as + after a Colab runtime reset. It loads its peers, as a login does.""" + config = SyftboxManagerConfig._base_config_for_testing( + email=ds_manager.email, + syftbox_folder=tmp_path / "SyftBox", + has_ds_role=True, + use_in_memory_cache=False, + ) + restarted = SyftboxManager.from_config( + config, ds_manager._connection_router.connections + ) + restarted.load_peers() + return restarted + + +def test_restart_on_empty_folder_restores_files_and_cursor(tmp_path): + ds_manager, do_manager = _pair_on_disk() + event = get_mock_event(f"{do_manager.email}/result.txt") + _send_to_ds(do_manager, ds_manager, [event]) + ds_manager.sync() + + restarted = _restart_on_empty_folder(ds_manager, tmp_path) + + restored = restarted.syftbox_folder / event.path_in_syftbox + assert restored.read_text() == event.content + cache = _watcher_cache(restarted) + assert cache.sync_down(do_manager.email) == 0 + + +def test_restart_after_checkpoint_keeps_deleted_file_deleted(tmp_path, monkeypatch): + monkeypatch.setattr(datasite_watcher_cache, "DEFAULT_CHECKPOINT_EVENT_THRESHOLD", 2) + ds_manager, do_manager = _pair_on_disk() + kept = get_mock_event(f"{do_manager.email}/kept.txt") + deleted = get_mock_event(f"{do_manager.email}/deleted.txt") + _send_to_ds(do_manager, ds_manager, [kept, deleted]) + ds_manager.sync() + router = ds_manager._connection_router + assert router.watcher_load_checkpoint(do_manager.email) is not None + _send_to_ds(do_manager, ds_manager, [_deletion_of(deleted)]) + ds_manager.sync() + + restarted = _restart_on_empty_folder(ds_manager, tmp_path) + + assert (restarted.syftbox_folder / kept.path_in_syftbox).exists() + assert not (restarted.syftbox_folder / deleted.path_in_syftbox).exists() + + +def test_file_refused_by_write_filter_not_restored(tmp_path): + ds_manager, do_manager = _pair_on_disk() + refused = get_mock_event(f"{do_manager.email}/refused.txt") + cache = _watcher_cache(ds_manager) + cache.pre_write_filter = lambda path, is_delete: not path.endswith("refused.txt") + _send_to_ds(do_manager, ds_manager, [refused]) + ds_manager.sync() + + restarted = _restart_on_empty_folder(ds_manager, tmp_path) + + assert not (restarted.syftbox_folder / refused.path_in_syftbox).exists() + + +def test_restart_with_new_key_restores_files_and_cursor(tmp_path): + """An enclave gets a new key at each boot. Its watcher state stays + readable, because it is not encrypted.""" + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + use_in_memory_cache=False, encryption=True + ) + event = get_mock_event(f"{do_manager.email}/result.txt") + _send_to_ds(do_manager, ds_manager, [event]) + ds_manager.sync() + config = SyftboxManagerConfig._base_config_for_testing( + email=ds_manager.email, + syftbox_folder=tmp_path / "SyftBox", + has_ds_role=True, + use_in_memory_cache=False, + encryption_policy=EncryptionPolicy.REQUIRED, + ) + + rebooted = SyftboxManager.from_config( + config, ds_manager._connection_router.connections + ) + + restored = rebooted.syftbox_folder / event.path_in_syftbox + assert restored.read_text() == event.content + cursor = _watcher_cache(ds_manager).last_msg_id_per_peer + assert _watcher_cache(rebooted).last_msg_id_per_peer == cursor + + +def test_binary_file_restored_from_checkpoint(tmp_path, monkeypatch): + monkeypatch.setattr(datasite_watcher_cache, "DEFAULT_CHECKPOINT_EVENT_THRESHOLD", 1) + ds_manager, do_manager = _pair_on_disk() + event = get_mock_event(f"{do_manager.email}/image.bin").model_copy( + update={"content": b"\x00\xff"} + ) + _send_to_ds(do_manager, ds_manager, [event]) + ds_manager.sync() + router = ds_manager._connection_router + assert router.watcher_load_checkpoint(do_manager.email) is not None + + restarted = _restart_on_empty_folder(ds_manager, tmp_path) + + restored = restarted.syftbox_folder / event.path_in_syftbox + assert restored.read_bytes() == b"\x00\xff" + + +def test_restore_writes_files_where_events_write(tmp_path, monkeypatch): + """The DO's events may spell its email in another case than the DS's peer + entry does. A restored file goes where the live event put it.""" + monkeypatch.setattr(datasite_watcher_cache, "DEFAULT_CHECKPOINT_EVENT_THRESHOLD", 1) + ds_manager, do_manager = _pair_on_disk() + event = get_mock_event(f"{do_manager.email.upper()}/result.txt") + _send_to_ds(do_manager, ds_manager, [event]) + ds_manager.sync() + + restarted = _restart_on_empty_folder(ds_manager, tmp_path) + + assert (restarted.syftbox_folder / event.path_in_syftbox).exists() + + +def test_do_state_saved_under_email_without_case(monkeypatch): + monkeypatch.setattr(datasite_watcher_cache, "DEFAULT_CHECKPOINT_EVENT_THRESHOLD", 1) + ds_manager, do_manager = SyftboxManager.pair_with_mock_drive_service_connection( + email1="Owner@Example.com", use_in_memory_cache=False + ) + _send_to_ds(do_manager, ds_manager, [get_mock_event("Owner@Example.com/a.txt")]) + + ds_manager.sync() + + router = ds_manager._connection_router + assert router.watcher_peers_with_saved_state() == ["owner@example.com"] + + +def test_change_without_hash_or_content_dropped(monkeypatch): + monkeypatch.setattr(datasite_watcher_cache, "DEFAULT_CHECKPOINT_EVENT_THRESHOLD", 1) + ds_manager, do_manager = _pair_on_disk() + no_hash = get_mock_event(f"{do_manager.email}/no_hash.txt").model_copy( + update={"new_hash": None} + ) + no_content = get_mock_event(f"{do_manager.email}/no_content.txt").model_copy( + update={"content": None} + ) + good = get_mock_event(f"{do_manager.email}/good.txt") + _send_to_ds(do_manager, ds_manager, [no_hash, no_content, good]) + + ds_manager.sync() + + datasite = ds_manager.syftbox_folder / do_manager.email + assert (datasite / "good.txt").exists() + assert not (datasite / "no_hash.txt").exists() + assert not (datasite / "no_content.txt").exists()