From 1c4a51d336b7d381ec1087c2cc998434be81d2e6 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 3 Nov 2026 17:39:27 -0300 Subject: [PATCH 01/12] feat: add the DS's storage on Drive for what it received from each DO - WatcherCheckpoint and WatcherRollingState in syft/sync/checkpoints/watcher_state.py hold the files received from one DO, the changes since then and the cursor in its outbox. They refuse a state that a later client wrote. - parse_versioned_json() checks the version of JSON text or bytes before the fields. read_versioned_json() now calls it for a file. - DriveWatcherStoreConnection keeps one folder per DO in /SyftBox/--watcher-state, which is shared with nobody. A save replaces the content of the checkpoint file or the rolling-state file in place. - put_checkpoint() and put_rolling_state() take no file name, because each DO has one file of each kind. --- syft/sync/checkpoints/watcher_state.py | 63 ++++++++++ .../connections/drive/drive_connections.py | 24 +--- .../connections/drive/drive_watcher_store.py | 116 ++++++++++++++++++ syft/sync/connections/protocols.py | 6 +- syft/sync/syftbox_manager.py | 10 +- syft/sync/utils/versioned_file.py | 34 +++-- tests/unit/test_drive_connections.py | 30 ++++- tests/unit/test_shared_router.py | 14 +-- tests/unit/test_watcher_state.py | 18 +++ 9 files changed, 259 insertions(+), 56 deletions(-) create mode 100644 syft/sync/checkpoints/watcher_state.py create mode 100644 syft/sync/connections/drive/drive_watcher_store.py create mode 100644 tests/unit/test_watcher_state.py diff --git a/syft/sync/checkpoints/watcher_state.py b/syft/sync/checkpoints/watcher_state.py new file mode 100644 index 00000000000..12b2dc26999 --- /dev/null +++ b/syft/sync/checkpoints/watcher_state.py @@ -0,0 +1,63 @@ +"""What a DS received from each DO, kept in the DS's own storage. + +A lost local disk, as after a Colab runtime reset, loses the files from a DO +and the cursor in its outbox. The DS restores both from the checkpoint and the +rolling state of that DO, so it does not apply an outbox message again. +""" + +from typing import ClassVar, Self + +from pydantic import BaseModel, ConfigDict + +from syft.sync.events.file_change_event import FileChangeEvent +from syft.sync.utils.syftbox_utils import compress_data, uncompress_data +from syft.sync.utils.versioned_file import parse_versioned_json + +# Format of the checkpoint and the rolling state. Raise it when the layout of +# either changes. A state of a later version is refused. +WATCHER_STATE_VERSION = 1 + + +class _WatcherState(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + # Names the state in errors. + kind: ClassVar[str] + + version: int + + def as_compressed_data(self) -> bytes: + return compress_data(self.model_dump_json().encode("utf-8")) + + @classmethod + def from_compressed_data(cls, data: bytes) -> Self: + """Raises: + ValueError: if a later client wrote the state, or it has another shape. + """ + return parse_versioned_json( + uncompress_data(data), cls, WATCHER_STATE_VERSION, cls.kind + ) + + +class WatcherCheckpoint(_WatcherState): + """The files received from one DO: the latest event of each path. A deleted + path has no event.""" + + kind: ClassVar[str] = "watcher checkpoint" + + events: list[FileChangeEvent] + + +class WatcherRollingState(_WatcherState): + """The changes received from one DO since its checkpoint, and the cursor in + its outbox.""" + + kind: ClassVar[str] = "watcher rolling state" + + # The latest event of each path since the checkpoint, deletes included. + events: list[FileChangeEvent] + # The msg_id of the newest events message received from the DO. + last_msg_id: str + # The ids of the events messages received. The outbox can list some of + # them again, and the DS skips those. + received_msg_ids: frozenset[str] 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..5661fb87639 --- /dev/null +++ b/syft/sync/connections/drive/drive_watcher_store.py @@ -0,0 +1,116 @@ +"""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 + + @raise_unavailable("Could not save the checkpoint to Drive.") + def put_checkpoint(self, peer: str, data: bytes) -> str: + return self._put(peer, CHECKPOINT_FILE_NAME, data) + + @raise_unavailable("Could not read the checkpoint from Drive.") + def latest_checkpoint(self, peer: str) -> bytes | None: + return self._get(peer, CHECKPOINT_FILE_NAME) + + @raise_unavailable("Could not save the rolling state to Drive.") + def put_rolling_state(self, peer: str, data: bytes) -> str: + return self._put(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._get(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) + + def _put(self, peer: str, name: str, data: bytes) -> str: + client = self._cached_client.client + folder_id = self._folder_id() or self._create_folder() + peer_folder_id = self._find_one(peer, folder_id) or client.create_folder( + peer, folder_id + ) + payload, _ = client.create_file_payload(data) + file_id = self._find_one(name, peer_folder_id) + if file_id is None: + body = {"name": name, "parents": [peer_folder_id]} + return client.create_syft_file(body, payload)["id"] + execute_with_retries( + client.drive_service.files().update(fileId=file_id, media_body=payload) + ) + return file_id + + def _get(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_one(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: + folder_id = self._folder_id() + if folder_id is None: + return None + return self._find_one(peer, folder_id) + + def _find_one(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 _folder_name(self) -> str: + return ( + f"{self._cached_client.email}-{SYFT_VERSION}{WATCHER_STATE_FOLDER_SUFFIX}" + ) + + def _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._folder_name() + ) + + def _create_folder(self) -> str: + return self._cached_client.client.create_folder( + self._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..aaa1a6edaf2 100644 --- a/syft/sync/connections/protocols.py +++ b/syft/sync/connections/protocols.py @@ -204,7 +204,7 @@ class WatcherStoreConnection(RouterConnection, ABC): """ @abstractmethod - def put_checkpoint(self, peer: str, name: str, data: bytes) -> str: + def put_checkpoint(self, peer: str, data: bytes) -> str: """Store a snapshot of every file received from ``peer``; it replaces the previous one.""" @@ -213,9 +213,9 @@ def latest_checkpoint(self, peer: str) -> bytes | None: """The newest 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) -> str: """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: diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index 9a36f0d6363..e3a10de8304 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -15,13 +15,13 @@ 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 +620,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(), diff --git a/syft/sync/utils/versioned_file.py b/syft/sync/utils/versioned_file.py index 65c803221f6..849997efe24 100644 --- a/syft/sync/utils/versioned_file.py +++ b/syft/sync/utils/versioned_file.py @@ -1,4 +1,4 @@ -"""Read a JSON file that carries a format version.""" +"""Read JSON that carries a format version.""" import json from pathlib import Path @@ -9,22 +9,34 @@ ModelT = TypeVar("ModelT", bound=BaseModel) -def read_versioned_json( - path: Path, model: type[ModelT], max_version: int, kind: str +def parse_versioned_json( + data: str | bytes, model: type[ModelT], max_version: int, kind: str ) -> ModelT: - """The file at ``path`` as ``model``. ``kind`` names the file in errors. + """``data`` as ``model``. ``kind`` names the data in errors. - The version is checked before the fields, so a file of a later layout is + The version is checked before the fields, so data of a later layout is refused by its version, not by a field error. Raises: - ValueError: if a later client wrote the file, or it has another shape. + ValueError: if a later client wrote the data, or it has another shape. """ - text = path.read_text() - version = json.loads(text)["version"] + version = json.loads(data)["version"] if version > max_version: raise ValueError( - f"The {kind} at {path} has version {version}, and this client " - f"reads up to version {max_version}." + f"The {kind} has version {version}, and this client reads up to " + f"version {max_version}." ) - return model.model_validate_json(text) + return model.model_validate_json(data) + + +def read_versioned_json( + path: Path, model: type[ModelT], max_version: int, kind: str +) -> ModelT: + """The file at ``path`` as ``model``. ``kind`` names the file in errors. + + Raises: + ValueError: if a later client wrote the file, or it has another shape. + """ + return parse_versioned_json( + path.read_text(), model, max_version, f"{kind} at {path}" + ) diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 175b80a8ca4..7c69f248a24 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,30 @@ 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.latest_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.latest_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") + + store.reset(DO) + + assert store.get_rolling_state(DO) is None + assert store.get_rolling_state(other_do) == b"from other do" # -- CollectionsConnection -------------------------------------------------------- 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_state.py b/tests/unit/test_watcher_state.py new file mode 100644 index 00000000000..5023ecc7ee0 --- /dev/null +++ b/tests/unit/test_watcher_state.py @@ -0,0 +1,18 @@ +import json + +import pytest + +from syft.sync.checkpoints.watcher_state import ( + WATCHER_STATE_VERSION, + WatcherCheckpoint, +) +from syft.sync.utils.syftbox_utils import compress_data + + +def test_watcher_state_of_later_version_refused(): + later_version = WATCHER_STATE_VERSION + 1 + state = {"version": later_version, "events": [], "new_setting": True} + data = compress_data(json.dumps(state).encode("utf-8")) + + with pytest.raises(ValueError, match=f"version {later_version}"): + WatcherCheckpoint.from_compressed_data(data) From ebca5a0ec07a8a4873c7310f7f5da4800d7ca210 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 3 Nov 2026 17:57:59 -0300 Subject: [PATCH 02/12] feat: save and load the DS's state for each DO through the router, encrypted with the DS key - watcher_save_checkpoint() and watcher_save_rolling_state() encrypt the state for the DS when encryption is required, then store it. - watcher_load_checkpoint() and watcher_load_rolling_state() return None only when nothing is saved. They raise when the state cannot be decrypted or a later client saved it. - When the DS key cannot decrypt the state, the error tells the DS to log in with the key_file of the key that saved it. --- syft/sync/connections/connection_router.py | 67 ++++++++++++++++++++++ tests/unit/test_encryption.py | 49 +++++++++++++++- 2 files changed, 115 insertions(+), 1 deletion(-) diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 851bf4ea828..425d61d7ad4 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -14,6 +14,7 @@ ) from syft.sync.checkpoints.checkpoint import Checkpoint, IncrementalCheckpoint from syft.sync.checkpoints.rolling_state import RollingState +from syft.sync.checkpoints.watcher_state import WatcherCheckpoint, WatcherRollingState from syft.sync.connections.base_connection import FileCollection from syft.sync.connections.naming import PEERS_META_KEY from syft.sync.connections.protocols import Connections @@ -575,3 +576,69 @@ 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 (self-encrypted at rest when encryption enabled) + # ========================================================================= + + def watcher_save_checkpoint( + self, peer_email: str, checkpoint: WatcherCheckpoint + ) -> None: + """Save the files received from ``peer_email``; it replaces the previous + checkpoint.""" + data = self.peer_store.encrypt_for_self_if_needed( + checkpoint.as_compressed_data() + ) + self.connections.watcher_store_con.put_checkpoint(peer_email, data) + + def watcher_load_checkpoint(self, peer_email: str) -> WatcherCheckpoint | None: + """The saved checkpoint of ``peer_email``, or None if there is none. + + Raises: + ValueError: if this DS key cannot decrypt it, or a later client + saved it. + """ + data = self.connections.watcher_store_con.latest_checkpoint(peer_email) + if data is None: + return None + return WatcherCheckpoint.from_compressed_data( + self._decrypt_watcher_state(data, WatcherCheckpoint.kind, peer_email) + ) + + def watcher_save_rolling_state( + self, peer_email: str, rolling_state: WatcherRollingState + ) -> None: + """Save the changes received from ``peer_email`` since its checkpoint, + and the cursor; it replaces the previous rolling state.""" + data = self.peer_store.encrypt_for_self_if_needed( + rolling_state.as_compressed_data() + ) + self.connections.watcher_store_con.put_rolling_state(peer_email, data) + + def watcher_load_rolling_state(self, peer_email: str) -> WatcherRollingState | None: + """The saved rolling state of ``peer_email``, or None if there is none. + + Raises: + ValueError: if this DS key cannot decrypt it, or a later client + saved it. + """ + data = self.connections.watcher_store_con.get_rolling_state(peer_email) + if data is None: + return None + return WatcherRollingState.from_compressed_data( + self._decrypt_watcher_state(data, WatcherRollingState.kind, peer_email) + ) + + def _decrypt_watcher_state(self, data: bytes, kind: str, peer_email: str) -> bytes: + """``data`` decrypted with this DS key. ``kind`` names the state in errors. + + Raises: + ValueError: if this DS key cannot decrypt it. + """ + try: + return self.peer_store.decrypt_and_verify_for_self_if_needed(data) + except ValueError as e: + raise ValueError( + f"This DS key cannot decrypt the {kind} of {peer_email}. Log in " + "with the key_file of the key that saved it." + ) from e diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index c35aef902dd..cd5ead7c46e 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -5,13 +5,19 @@ import pytest +from syft.sync.checkpoints.watcher_state import ( + WATCHER_STATE_VERSION, + WatcherCheckpoint, + WatcherRollingState, +) +from syft.sync.connections.connection_router import ConnectionRouter from syft.sync.encryption_policy import EncryptionPolicy from syft.sync.peers.exceptions import SyftPeerNotReadyError from syft.sync.peers.peer import Peer, PeerState from syft.sync.peers.peer_store import PeerStore, import_key_file 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 tests.unit.utils import grant_job_inbox_access, router_with_mock_connection from syft.sync.connections.drive.drive_connections import drive_cached_client_of @@ -638,3 +644,44 @@ def test_tampered_inbox_message_is_archived_with_warning(): with pytest.warns(UserWarning, match="unread"): do_manager.sync() assert do_cached_client.get_file_metadatas_from_folder(inbox_id) == [] + + +def _router_with_new_key(connections) -> ConnectionRouter: + """A DS router over ``connections`` that requires encryption, with a key of + its own.""" + peer_store = PeerStore( + email="ds@example.com", encryption_policy=EncryptionPolicy.REQUIRED + ) + peer_store.generate_keys() + return ConnectionRouter(connections=connections, peer_store=peer_store) + + +def test_watcher_state_encrypted_at_rest(): + router = _router_with_new_key( + router_with_mock_connection("ds@example.com").connections + ) + state = WatcherRollingState( + version=WATCHER_STATE_VERSION, + events=[], + last_msg_id="msg-1", + received_msg_ids=frozenset({"msg-1"}), + ) + + router.watcher_save_rolling_state("do@example.com", state) + + stored = router.connections.watcher_store_con.get_rolling_state("do@example.com") + assert stored is not None + plaintext = router.peer_store.decrypt_for_self(stored) + assert WatcherRollingState.from_compressed_data(plaintext) == state + assert router.watcher_load_rolling_state("do@example.com") == state + + +def test_watcher_state_of_other_key_names_key_file(): + connections = router_with_mock_connection("ds@example.com").connections + checkpoint = WatcherCheckpoint(version=WATCHER_STATE_VERSION, events=[]) + _router_with_new_key(connections).watcher_save_checkpoint( + "do@example.com", checkpoint + ) + + with pytest.raises(ValueError, match="key_file"): + _router_with_new_key(connections).watcher_load_checkpoint("do@example.com") From 330ead146398561272875527ff239102396cb7e1 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 3 Nov 2026 19:59:20 -0300 Subject: [PATCH 03/12] feat: restore a DS's files and cursor for each DO from its own Drive at startup - At startup, a DS in on-disk mode restores every DO that has saved state. It writes the files of the checkpoint and the rolling state, and records their hashes and the received message ids. - After a pull that receives messages, the DS saves the rolling state of that DO, then marks the messages processed. At 50 changed paths, it folds them into the checkpoint. - The DS keeps the received message ids per DO. An event that pre_write_filter refuses is not saved, so a restore does not write it. - WatcherRollingState has no last_msg_id, because the cursor is the newest received id. - peers_with_saved_state() lists the DOs that have saved state. --- syft/sync/checkpoints/watcher_state.py | 6 +- syft/sync/connections/connection_router.py | 4 + .../connections/drive/drive_watcher_store.py | 8 ++ syft/sync/connections/protocols.py | 4 + .../sync/caches/datasite_watcher_cache.py | 123 +++++++++++++++--- tests/unit/test_drive_connections.py | 2 + tests/unit/test_encryption.py | 1 - tests/unit/test_watcher_state.py | 102 +++++++++++++++ 8 files changed, 229 insertions(+), 21 deletions(-) diff --git a/syft/sync/checkpoints/watcher_state.py b/syft/sync/checkpoints/watcher_state.py index 12b2dc26999..0301ce08ea2 100644 --- a/syft/sync/checkpoints/watcher_state.py +++ b/syft/sync/checkpoints/watcher_state.py @@ -49,15 +49,13 @@ class WatcherCheckpoint(_WatcherState): class WatcherRollingState(_WatcherState): - """The changes received from one DO since its checkpoint, and the cursor in - its outbox.""" + """The changes received from one DO since its checkpoint, and the ids of the + messages received from it. The newest id is the cursor in its outbox.""" kind: ClassVar[str] = "watcher rolling state" # The latest event of each path since the checkpoint, deletes included. events: list[FileChangeEvent] - # The msg_id of the newest events message received from the DO. - last_msg_id: str # The ids of the events messages received. The outbox can list some of # them again, and the DS skips those. received_msg_ids: frozenset[str] diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 425d61d7ad4..b30f16c062c 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -629,6 +629,10 @@ def watcher_load_rolling_state(self, peer_email: str) -> WatcherRollingState | N self._decrypt_watcher_state(data, WatcherRollingState.kind, peer_email) ) + 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() + def _decrypt_watcher_state(self, data: bytes, kind: str, peer_email: str) -> bytes: """``data`` decrypted with this DS key. ``kind`` names the state in errors. diff --git a/syft/sync/connections/drive/drive_watcher_store.py b/syft/sync/connections/drive/drive_watcher_store.py index 5661fb87639..a451a9b7baa 100644 --- a/syft/sync/connections/drive/drive_watcher_store.py +++ b/syft/sync/connections/drive/drive_watcher_store.py @@ -47,6 +47,14 @@ def reset(self, peer: str) -> None: 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]: + folder_id = self._folder_id() + if folder_id is None: + return [] + peer_folders = self._cached_client.client.find_folders(parent_id=folder_id) + return [name for _, name in peer_folders] + def _put(self, peer: str, name: str, data: bytes) -> str: client = self._cached_client.client folder_id = self._folder_id() or self._create_folder() diff --git a/syft/sync/connections/protocols.py b/syft/sync/connections/protocols.py index aaa1a6edaf2..90fc3a95af9 100644 --- a/syft/sync/connections/protocols.py +++ b/syft/sync/connections/protocols.py @@ -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/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 9255d7abf91..49358adec8f 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -1,18 +1,25 @@ 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.watcher_state import ( + WATCHER_STATE_VERSION, + WatcherCheckpoint, + WatcherRollingState, +) +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 @@ -61,6 +68,9 @@ def _events_within_peer_datasite( SECONDS_BEFORE_SYNCING_DOWN = 0 +# The changes from a DO go into its checkpoint at this many paths, as the DO +# does with its own state. +WATCHER_CHECKPOINT_THRESHOLD = 50 class DataSiteWatcherCacheConfig(BaseModel): @@ -90,9 +100,14 @@ 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 + # peer key -> {path in syftbox: latest event since the peer's checkpoint} + rolling_events_per_peer: Dict[str, Dict[Path, FileChangeEvent]] = {} # Base syftbox folder syftbox_folder: Path | None = None # Collections to sync down (prefix + local subpath per collection type) @@ -118,6 +133,7 @@ def from_config( connection_router=connection_router, syftbox_folder=config.syftbox_folder, collection_specs=config.collection_specs, + saves_watcher_state=False, ) return res else: @@ -136,8 +152,10 @@ def from_config( connection_router=connection_router, syftbox_folder=config.syftbox_folder, collection_specs=config.collection_specs, + saves_watcher_state=True, ) cache._load_cached_state() + cache._restore_saved_state() return cache def _load_cached_state(self): @@ -159,9 +177,8 @@ def _load_saved_messages(self): 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: + msg_id = events_message.message_filepath.as_string() self._record_received(events_message.events[0].datasite_email, msg_id) def _restore_file_hashes(self, events_message: FileChangeEventsMessage) -> None: @@ -226,7 +243,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_events_per_peer = {} self.collection_hashes = {} @property @@ -239,11 +257,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 +273,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_state(peer_email) self.connection_router.mark_all_processed([raw for raw, _ in received]) self.last_sync = datetime.now() @@ -259,8 +283,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 @@ -285,12 +309,20 @@ def apply_event_message( 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_events(peer_email)[event.path_in_syftbox] = 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 +332,65 @@ 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_saved_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.""" + for peer_email in self.connection_router.watcher_peers_with_saved_state(): + checkpoint = self.connection_router.watcher_load_checkpoint(peer_email) + if checkpoint is not None: + for event in checkpoint.events: + self._apply_event(event) + rolling_state = self.connection_router.watcher_load_rolling_state( + peer_email + ) + if rolling_state is not None: + for event in rolling_state.events: + self._apply_event(event) + self._rolling_events(peer_email)[event.path_in_syftbox] = event + for msg_id in rolling_state.received_msg_ids: + self._record_received(peer_email, msg_id) + + def _save_state(self, peer_email: str) -> None: + """Save the rolling state of ``peer_email``. At + ``WATCHER_CHECKPOINT_THRESHOLD`` paths, fold it into the checkpoint.""" + self._save_rolling_state(peer_email) + if len(self._rolling_events(peer_email)) >= WATCHER_CHECKPOINT_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.""" + checkpoint = self.connection_router.watcher_load_checkpoint(peer_email) + latest = {} + if checkpoint is not None: + latest = {event.path_in_syftbox: event for event in checkpoint.events} + latest.update(self._rolling_events(peer_email)) + events = [event for event in latest.values() if not event.is_deleted] + # 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( + peer_email, + WatcherCheckpoint(version=WATCHER_STATE_VERSION, events=events), + ) + self._rolling_events(peer_email).clear() + self._save_rolling_state(peer_email) + + def _save_rolling_state(self, peer_email: str) -> None: + rolling_state = WatcherRollingState( + version=WATCHER_STATE_VERSION, + events=list(self._rolling_events(peer_email).values()), + received_msg_ids=frozenset( + self.received_msg_ids_per_peer[_peer_key(peer_email)] + ), + ) + self.connection_router.watcher_save_rolling_state(peer_email, rolling_state) + + def _rolling_events(self, peer_email: str) -> Dict[Path, FileChangeEvent]: + """The latest event of each path from ``peer_email`` since its checkpoint.""" + return self.rolling_events_per_peer.setdefault(_peer_key(peer_email), {}) + 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/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 7c69f248a24..1479841468e 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -568,11 +568,13 @@ def test_watcher_state_reset_keeps_other_peers(pair): 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 cd5ead7c46e..400dbf5be6c 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -663,7 +663,6 @@ def test_watcher_state_encrypted_at_rest(): state = WatcherRollingState( version=WATCHER_STATE_VERSION, events=[], - last_msg_id="msg-1", received_msg_ids=frozenset({"msg-1"}), ) diff --git a/tests/unit/test_watcher_state.py b/tests/unit/test_watcher_state.py index 5023ecc7ee0..6fd3a926083 100644 --- a/tests/unit/test_watcher_state.py +++ b/tests/unit/test_watcher_state.py @@ -1,4 +1,6 @@ import json +import time +import uuid import pytest @@ -6,7 +8,62 @@ WATCHER_STATE_VERSION, WatcherCheckpoint, ) +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 syft.sync.utils.syftbox_utils import compress_data +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_watcher_state_of_later_version_refused(): @@ -16,3 +73,48 @@ def test_watcher_state_of_later_version_refused(): with pytest.raises(ValueError, match=f"version {later_version}"): WatcherCheckpoint.from_compressed_data(data) + + +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, "WATCHER_CHECKPOINT_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() From 8402930423e04ac62e8ed4f78b6d27682f1277a3 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Tue, 3 Nov 2026 20:14:02 -0300 Subject: [PATCH 04/12] refactor: remove the DS's local -event-messages cache, now that Drive holds its state - In on-disk mode, the DS keeps its record of received messages in memory, as in-memory mode does. A restart restores only from the saved state on Drive. - The code that rebuilt the cursor and the file hashes from the local folder is gone, and so is the unused events_base_path. - The local cleanup no longer deletes a -event-messages folder. - The three tests that restarted from the local folder now check the state restored from Drive. --- syft/sync/syftbox_manager.py | 1 - .../sync/caches/datasite_watcher_cache.py | 47 ++----------------- syft/sync/utils/syftbox_utils.py | 5 +- .../test_cache_aware_sync.py | 18 +++---- tests/unit/test_cache_aware_sync.py | 18 ++----- .../unit/test_watcher_event_datasite_check.py | 31 +++++++----- 6 files changed, 35 insertions(+), 85 deletions(-) diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index e3a10de8304..a829ad1a198 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -1344,7 +1344,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(): diff --git a/syft/sync/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 49358adec8f..5f469159807 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -77,7 +77,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] = [] @@ -140,55 +139,17 @@ def from_config( 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_saved_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.""" - if events_message.events: - msg_id = events_message.message_filepath.as_string() - 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(): @@ -300,10 +261,8 @@ def apply_event_message( ``//``, so an event aimed anywhere else is dropped before it reaches the local file connection. - 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) saved = event_message.model_copy(update={"events": accepted}) 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/unit/test_cache_aware_sync.py b/tests/unit/test_cache_aware_sync.py index f9991f9c427..8851e112c1d 100644 --- a/tests/unit/test_cache_aware_sync.py +++ b/tests/unit/test_cache_aware_sync.py @@ -94,8 +94,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 +130,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 +156,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.""" 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(): From f54b85f606b926cc5e2d296edb2f7f6a1bf25688 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 10:51:23 -0300 Subject: [PATCH 05/12] refactor: keep the DS's watcher state unencrypted and reuse the DO's checkpoint models - Watcher state is not encrypted until key recovery exists, as the ticket says. With the DS key, an enclave could not boot twice: its new key could not read the state of the last boot. - The DS saves a Checkpoint and a RollingState from syft/sync/checkpoints/, as the ticket says. watcher_state.py and parse_versioned_json are gone. - CheckpointFile.content holds str or bytes, so a binary file survives a checkpoint. RollingState.received_msg_ids holds the ids that a DS received. - The DS saves and loads the state of a DO under its email without case, so two spellings share one Drive folder. A restore writes the checkpoint files under the email as the DO's events spell it. --- syft/sync/checkpoints/checkpoint.py | 39 ++++++- syft/sync/checkpoints/rolling_state.py | 3 + syft/sync/checkpoints/watcher_state.py | 61 ---------- syft/sync/connections/connection_router.py | 55 +++------ .../sync/caches/datasite_watcher_cache.py | 106 +++++++++++------- syft/sync/utils/versioned_file.py | 34 ++---- tests/unit/test_checkpoints.py | 21 +++- tests/unit/test_encryption.py | 54 +-------- tests/unit/test_watcher_state.py | 89 ++++++++++++--- 9 files changed, 228 insertions(+), 234 deletions(-) delete mode 100644 syft/sync/checkpoints/watcher_state.py diff --git a/syft/sync/checkpoints/checkpoint.py b/syft/sync/checkpoints/checkpoint.py index d7abb481bb9..5d8512a6e94 100644 --- a/syft/sync/checkpoints/checkpoint.py +++ b/syft/sync/checkpoints/checkpoint.py @@ -11,10 +11,17 @@ - After N incremental checkpoints: compact into single full Checkpoint """ +import base64 from pathlib import Path -from typing import TYPE_CHECKING, Dict, List - -from pydantic import BaseModel, Field +from typing import TYPE_CHECKING, Any, Dict, List, Literal + +from pydantic import ( + BaseModel, + Field, + computed_field, + field_serializer, + model_validator, +) from syft.sync.utils.syftbox_utils import ( compress_data, @@ -54,7 +61,31 @@ class CheckpointFile(BaseModel): 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) -> Literal["text", "binary"]: + """Saved with the file, so that a read gives bytes back as bytes.""" + return "binary" if isinstance(self.content, bytes) else "text" + + @field_serializer("content", when_used="json") + def _content_as_text(self, content: str | bytes) -> str: + """Bytes as base64, because JSON holds only text.""" + if isinstance(content, bytes): + return base64.b64encode(content).decode("ascii") + return content + + @model_validator(mode="before") + @classmethod + def _content_from_text(cls, data: Any) -> Any: + if ( + isinstance(data, dict) + and data.get("content_type") == "binary" + and isinstance(data.get("content"), str) + ): + return {**data, "content": base64.b64decode(data["content"])} + return data class Checkpoint(BaseModel): 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/checkpoints/watcher_state.py b/syft/sync/checkpoints/watcher_state.py deleted file mode 100644 index 0301ce08ea2..00000000000 --- a/syft/sync/checkpoints/watcher_state.py +++ /dev/null @@ -1,61 +0,0 @@ -"""What a DS received from each DO, kept in the DS's own storage. - -A lost local disk, as after a Colab runtime reset, loses the files from a DO -and the cursor in its outbox. The DS restores both from the checkpoint and the -rolling state of that DO, so it does not apply an outbox message again. -""" - -from typing import ClassVar, Self - -from pydantic import BaseModel, ConfigDict - -from syft.sync.events.file_change_event import FileChangeEvent -from syft.sync.utils.syftbox_utils import compress_data, uncompress_data -from syft.sync.utils.versioned_file import parse_versioned_json - -# Format of the checkpoint and the rolling state. Raise it when the layout of -# either changes. A state of a later version is refused. -WATCHER_STATE_VERSION = 1 - - -class _WatcherState(BaseModel): - model_config = ConfigDict(frozen=True, extra="forbid") - - # Names the state in errors. - kind: ClassVar[str] - - version: int - - def as_compressed_data(self) -> bytes: - return compress_data(self.model_dump_json().encode("utf-8")) - - @classmethod - def from_compressed_data(cls, data: bytes) -> Self: - """Raises: - ValueError: if a later client wrote the state, or it has another shape. - """ - return parse_versioned_json( - uncompress_data(data), cls, WATCHER_STATE_VERSION, cls.kind - ) - - -class WatcherCheckpoint(_WatcherState): - """The files received from one DO: the latest event of each path. A deleted - path has no event.""" - - kind: ClassVar[str] = "watcher checkpoint" - - events: list[FileChangeEvent] - - -class WatcherRollingState(_WatcherState): - """The changes received from one DO since its checkpoint, and the ids of the - messages received from it. The newest id is the cursor in its outbox.""" - - kind: ClassVar[str] = "watcher rolling state" - - # The latest event of each path since the checkpoint, deletes included. - events: list[FileChangeEvent] - # The ids of the events messages received. The outbox can list some of - # them again, and the DS skips those. - received_msg_ids: frozenset[str] diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index b30f16c062c..8f9882e28d8 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -14,7 +14,6 @@ ) from syft.sync.checkpoints.checkpoint import Checkpoint, IncrementalCheckpoint from syft.sync.checkpoints.rolling_state import RollingState -from syft.sync.checkpoints.watcher_state import WatcherCheckpoint, WatcherRollingState from syft.sync.connections.base_connection import FileCollection from syft.sync.connections.naming import PEERS_META_KEY from syft.sync.connections.protocols import Connections @@ -578,71 +577,49 @@ def delete_rolling_state(self) -> None: self.connections.owner_store_con.delete_rolling_state() # ========================================================================= - # WATCHER STATE (self-encrypted at rest when encryption enabled) + # 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: WatcherCheckpoint - ) -> None: + def watcher_save_checkpoint(self, peer_email: str, checkpoint: Checkpoint) -> None: """Save the files received from ``peer_email``; it replaces the previous checkpoint.""" - data = self.peer_store.encrypt_for_self_if_needed( - checkpoint.as_compressed_data() + self.connections.watcher_store_con.put_checkpoint( + peer_email, checkpoint.as_compressed_data() ) - self.connections.watcher_store_con.put_checkpoint(peer_email, data) - def watcher_load_checkpoint(self, peer_email: str) -> WatcherCheckpoint | None: + 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 this DS key cannot decrypt it, or a later client - saved it. + ValueError: if a later client saved it. """ data = self.connections.watcher_store_con.latest_checkpoint(peer_email) if data is None: return None - return WatcherCheckpoint.from_compressed_data( - self._decrypt_watcher_state(data, WatcherCheckpoint.kind, peer_email) - ) + return Checkpoint.from_compressed_data(data) def watcher_save_rolling_state( - self, peer_email: str, rolling_state: WatcherRollingState + self, peer_email: str, rolling_state: RollingState ) -> None: """Save the changes received from ``peer_email`` since its checkpoint, - and the cursor; it replaces the previous rolling state.""" - data = self.peer_store.encrypt_for_self_if_needed( - rolling_state.as_compressed_data() + 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() ) - self.connections.watcher_store_con.put_rolling_state(peer_email, data) - def watcher_load_rolling_state(self, peer_email: str) -> WatcherRollingState | None: + 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 this DS key cannot decrypt it, or a later client - saved it. + 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 WatcherRollingState.from_compressed_data( - self._decrypt_watcher_state(data, WatcherRollingState.kind, peer_email) - ) + 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() - - def _decrypt_watcher_state(self, data: bytes, kind: str, peer_email: str) -> bytes: - """``data`` decrypted with this DS key. ``kind`` names the state in errors. - - Raises: - ValueError: if this DS key cannot decrypt it. - """ - try: - return self.peer_store.decrypt_and_verify_for_self_if_needed(data) - except ValueError as e: - raise ValueError( - f"This DS key cannot decrypt the {kind} of {peer_email}. Log in " - "with the key_file of the key that saved it." - ) from e diff --git a/syft/sync/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 5f469159807..4ab761cc7db 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -6,11 +6,8 @@ from pydantic import BaseModel, Field -from syft.sync.checkpoints.watcher_state import ( - WATCHER_STATE_VERSION, - WatcherCheckpoint, - WatcherRollingState, -) +from syft.sync.checkpoints.checkpoint import 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, @@ -105,8 +102,10 @@ class DataSiteWatcherCache(BaseModel): 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 - # peer key -> {path in syftbox: latest event since the peer's checkpoint} - rolling_events_per_peer: Dict[str, Dict[Path, FileChangeEvent]] = {} + # 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) @@ -205,7 +204,7 @@ def clear_cache(self): self.current_check_point = None self.last_msg_id_per_peer = {} self.received_msg_ids_per_peer = {} - self.rolling_events_per_peer = {} + self.rolling_state_per_peer = {} self.collection_hashes = {} @property @@ -272,7 +271,7 @@ def apply_event_message( continue self._apply_event(event) if self.saves_watcher_state: - self._rolling_events(peer_email)[event.path_in_syftbox] = event + self._rolling_state(peer_email).add_event(event) def _write_denied(self, event: FileChangeEvent) -> bool: """Whether ``pre_write_filter`` refuses ``event``.""" @@ -295,60 +294,89 @@ def _apply_event(self, event: FileChangeEvent) -> None: def _restore_saved_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.""" - for peer_email in self.connection_router.watcher_peers_with_saved_state(): - checkpoint = self.connection_router.watcher_load_checkpoint(peer_email) + 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 event in checkpoint.events: - self._apply_event(event) - rolling_state = self.connection_router.watcher_load_rolling_state( - peer_email - ) + 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) - self._rolling_events(peer_email)[event.path_in_syftbox] = event for msg_id in rolling_state.received_msg_ids: - self._record_received(peer_email, msg_id) + 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_state(self, peer_email: str) -> None: """Save the rolling state of ``peer_email``. At ``WATCHER_CHECKPOINT_THRESHOLD`` paths, fold it into the checkpoint.""" self._save_rolling_state(peer_email) - if len(self._rolling_events(peer_email)) >= WATCHER_CHECKPOINT_THRESHOLD: + if self._rolling_state(peer_email).event_count >= WATCHER_CHECKPOINT_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.""" - checkpoint = self.connection_router.watcher_load_checkpoint(peer_email) - latest = {} + key = _peer_key(peer_email) + checkpoint = self.connection_router.watcher_load_checkpoint(key) + files: Dict[str, CheckpointFile] = {} if checkpoint is not None: - latest = {event.path_in_syftbox: event for event in checkpoint.events} - latest.update(self._rolling_events(peer_email)) - events = [event for event in latest.values() if not event.is_deleted] + files = {file.path: file for file in checkpoint.files} + rolling_state = self._rolling_state(peer_email) + for event in rolling_state.events: + path = str(event.path_in_datasite) + if event.is_deleted: + files.pop(path, None) + elif event.new_hash is None or event.content is None: + raise ValueError( + f"The change of {path} from {peer_email} has no content or no hash." + ) + else: + files[path] = CheckpointFile( + path=path, hash=event.new_hash, content=event.content + ) + folded = Checkpoint( + # Live events write under the email as the DO's events spell it, so + # a restore writes the checkpoint files there too. + email=rolling_state.events[-1].datasite_email, + last_event_timestamp=rolling_state.last_event_timestamp, + files=list(files.values()), + ) # 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( - peer_email, - WatcherCheckpoint(version=WATCHER_STATE_VERSION, events=events), + self.connection_router.watcher_save_checkpoint(key, folded) + self.rolling_state_per_peer[key] = RollingState( + email=peer_email, base_checkpoint_timestamp=folded.timestamp ) - self._rolling_events(peer_email).clear() self._save_rolling_state(peer_email) def _save_rolling_state(self, peer_email: str) -> None: - rolling_state = WatcherRollingState( - version=WATCHER_STATE_VERSION, - events=list(self._rolling_events(peer_email).values()), - received_msg_ids=frozenset( - self.received_msg_ids_per_peer[_peer_key(peer_email)] - ), + key = _peer_key(peer_email) + saved = self._rolling_state(peer_email).model_copy( + update={"received_msg_ids": sorted(self.received_msg_ids_per_peer[key])} ) - self.connection_router.watcher_save_rolling_state(peer_email, rolling_state) + self.connection_router.watcher_save_rolling_state(key, saved) - def _rolling_events(self, peer_email: str) -> Dict[Path, FileChangeEvent]: - """The latest event of each path from ``peer_email`` since its checkpoint.""" - return self.rolling_events_per_peer.setdefault(_peer_key(peer_email), {}) + 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() diff --git a/syft/sync/utils/versioned_file.py b/syft/sync/utils/versioned_file.py index 849997efe24..65c803221f6 100644 --- a/syft/sync/utils/versioned_file.py +++ b/syft/sync/utils/versioned_file.py @@ -1,4 +1,4 @@ -"""Read JSON that carries a format version.""" +"""Read a JSON file that carries a format version.""" import json from pathlib import Path @@ -9,34 +9,22 @@ ModelT = TypeVar("ModelT", bound=BaseModel) -def parse_versioned_json( - data: str | bytes, model: type[ModelT], max_version: int, kind: str +def read_versioned_json( + path: Path, model: type[ModelT], max_version: int, kind: str ) -> ModelT: - """``data`` as ``model``. ``kind`` names the data in errors. + """The file at ``path`` as ``model``. ``kind`` names the file in errors. - The version is checked before the fields, so data of a later layout is + The version is checked before the fields, so a file of a later layout is refused by its version, not by a field error. Raises: - ValueError: if a later client wrote the data, or it has another shape. + ValueError: if a later client wrote the file, or it has another shape. """ - version = json.loads(data)["version"] + text = path.read_text() + version = json.loads(text)["version"] if version > max_version: raise ValueError( - f"The {kind} has version {version}, and this client reads up to " - f"version {max_version}." + f"The {kind} at {path} has version {version}, and this client " + f"reads up to version {max_version}." ) - return model.model_validate_json(data) - - -def read_versioned_json( - path: Path, model: type[ModelT], max_version: int, kind: str -) -> ModelT: - """The file at ``path`` as ``model``. ``kind`` names the file in errors. - - Raises: - ValueError: if a later client wrote the file, or it has another shape. - """ - return parse_versioned_json( - path.read_text(), model, max_version, f"{kind} at {path}" - ) + return model.model_validate_json(text) diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index 8a6bbd84eec..cd61b712452 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,21 @@ 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()) diff --git a/tests/unit/test_encryption.py b/tests/unit/test_encryption.py index 400dbf5be6c..d89f2056702 100644 --- a/tests/unit/test_encryption.py +++ b/tests/unit/test_encryption.py @@ -5,21 +5,14 @@ import pytest -from syft.sync.checkpoints.watcher_state import ( - WATCHER_STATE_VERSION, - WatcherCheckpoint, - WatcherRollingState, -) -from syft.sync.connections.connection_router import ConnectionRouter +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 from syft.sync.peers.peer_store import PeerStore, import_key_file 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, router_with_mock_connection -from syft.sync.connections.drive.drive_connections import drive_cached_client_of - +from tests.unit.utils import grant_job_inbox_access # ========================================================================= # PeerStore unit tests @@ -453,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, @@ -644,43 +638,3 @@ def test_tampered_inbox_message_is_archived_with_warning(): with pytest.warns(UserWarning, match="unread"): do_manager.sync() assert do_cached_client.get_file_metadatas_from_folder(inbox_id) == [] - - -def _router_with_new_key(connections) -> ConnectionRouter: - """A DS router over ``connections`` that requires encryption, with a key of - its own.""" - peer_store = PeerStore( - email="ds@example.com", encryption_policy=EncryptionPolicy.REQUIRED - ) - peer_store.generate_keys() - return ConnectionRouter(connections=connections, peer_store=peer_store) - - -def test_watcher_state_encrypted_at_rest(): - router = _router_with_new_key( - router_with_mock_connection("ds@example.com").connections - ) - state = WatcherRollingState( - version=WATCHER_STATE_VERSION, - events=[], - received_msg_ids=frozenset({"msg-1"}), - ) - - router.watcher_save_rolling_state("do@example.com", state) - - stored = router.connections.watcher_store_con.get_rolling_state("do@example.com") - assert stored is not None - plaintext = router.peer_store.decrypt_for_self(stored) - assert WatcherRollingState.from_compressed_data(plaintext) == state - assert router.watcher_load_rolling_state("do@example.com") == state - - -def test_watcher_state_of_other_key_names_key_file(): - connections = router_with_mock_connection("ds@example.com").connections - checkpoint = WatcherCheckpoint(version=WATCHER_STATE_VERSION, events=[]) - _router_with_new_key(connections).watcher_save_checkpoint( - "do@example.com", checkpoint - ) - - with pytest.raises(ValueError, match="key_file"): - _router_with_new_key(connections).watcher_load_checkpoint("do@example.com") diff --git a/tests/unit/test_watcher_state.py b/tests/unit/test_watcher_state.py index 6fd3a926083..d1df7dd5524 100644 --- a/tests/unit/test_watcher_state.py +++ b/tests/unit/test_watcher_state.py @@ -1,18 +1,11 @@ -import json import time import uuid -import pytest - -from syft.sync.checkpoints.watcher_state import ( - WATCHER_STATE_VERSION, - WatcherCheckpoint, -) +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 syft.sync.utils.syftbox_utils import compress_data from tests.unit.utils import get_mock_event @@ -66,15 +59,6 @@ def _restart_on_empty_folder(ds_manager: SyftboxManager, tmp_path) -> SyftboxMan return restarted -def test_watcher_state_of_later_version_refused(): - later_version = WATCHER_STATE_VERSION + 1 - state = {"version": later_version, "events": [], "new_setting": True} - data = compress_data(json.dumps(state).encode("utf-8")) - - with pytest.raises(ValueError, match=f"version {later_version}"): - WatcherCheckpoint.from_compressed_data(data) - - 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") @@ -118,3 +102,74 @@ def test_file_refused_by_write_filter_not_restored(tmp_path): 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, "WATCHER_CHECKPOINT_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, "WATCHER_CHECKPOINT_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, "WATCHER_CHECKPOINT_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"] From 7fd99f747bb6dec40d9beb936f68f2f403ca832d Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 10:56:58 -0300 Subject: [PATCH 06/12] fix: drop a DO change that has no content or no hash when it arrives - Such a change blocked every later sync of the DS. A change with no hash made each fold raise, and a change with no content raised when it was applied. - The DS now drops it with a warning, as it does for an event aimed at another datasite. --- .../sync/caches/datasite_watcher_cache.py | 28 ++++++++++++++----- tests/unit/test_watcher_state.py | 20 +++++++++++++ 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/syft/sync/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 4ab761cc7db..67ac6e9c983 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -47,20 +47,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 @@ -258,12 +271,13 @@ 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 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: diff --git a/tests/unit/test_watcher_state.py b/tests/unit/test_watcher_state.py index d1df7dd5524..c0d6965441a 100644 --- a/tests/unit/test_watcher_state.py +++ b/tests/unit/test_watcher_state.py @@ -173,3 +173,23 @@ def test_do_state_saved_under_email_without_case(monkeypatch): 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, "WATCHER_CHECKPOINT_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() From 2d3823a3cb9d85e5c005caeac73bbd04ba621e7b Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 11:13:59 -0300 Subject: [PATCH 07/12] refactor: share the file content and checkpoint merge code between the DO and the DS - syft/sync/utils/file_content.py holds the text-or-bytes content code that ProposedFileChange, FileChangeEvent and CheckpointFile used to repeat. The JSON shape does not change. - Checkpoint.with_events() applies events to a checkpoint. The DO compaction and the DS fold both use it, and compact_incremental_checkpoints is gone. - The DO compaction now raises on a change without content or hash. Before, it dropped the file and could write "None" as a hash. --- syft/sync/checkpoints/checkpoint.py | 110 +++++++----------- syft/sync/events/file_change_event.py | 31 ++--- syft/sync/messages/proposed_filechange.py | 29 ++--- .../sync/caches/datasite_watcher_cache.py | 25 +--- syft/sync/sync/datasite_owner_syncer.py | 65 +---------- syft/sync/utils/file_content.py | 36 ++++++ tests/unit/test_checkpoints.py | 8 ++ 7 files changed, 111 insertions(+), 193 deletions(-) create mode 100644 syft/sync/utils/file_content.py diff --git a/syft/sync/checkpoints/checkpoint.py b/syft/sync/checkpoints/checkpoint.py index 5d8512a6e94..02cbe47848f 100644 --- a/syft/sync/checkpoints/checkpoint.py +++ b/syft/sync/checkpoints/checkpoint.py @@ -11,9 +11,8 @@ - After N incremental checkpoints: compact into single full Checkpoint """ -import base64 from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, List, Literal +from typing import TYPE_CHECKING, Any, Dict, List from pydantic import ( BaseModel, @@ -23,6 +22,12 @@ 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, @@ -65,27 +70,18 @@ class CheckpointFile(BaseModel): @computed_field @property - def content_type(self) -> Literal["text", "binary"]: + def content_type(self) -> ContentType | None: """Saved with the file, so that a read gives bytes back as bytes.""" - return "binary" if isinstance(self.content, bytes) else "text" + return content_type_of(self.content) @field_serializer("content", when_used="json") - def _content_as_text(self, content: str | bytes) -> str: - """Bytes as base64, because JSON holds only text.""" - if isinstance(content, bytes): - return base64.b64encode(content).decode("ascii") - return content + 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: - if ( - isinstance(data, dict) - and data.get("content_type") == "binary" - and isinstance(data.get("content"), str) - ): - return {**data, "content": base64.b64decode(data["content"])} - return data + return decode_content(data) class Checkpoint(BaseModel): @@ -164,6 +160,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")) @@ -235,59 +260,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/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/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 67ac6e9c983..ebdad69b53c 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -343,30 +343,13 @@ 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) - checkpoint = self.connection_router.watcher_load_checkpoint(key) - files: Dict[str, CheckpointFile] = {} - if checkpoint is not None: - files = {file.path: file for file in checkpoint.files} rolling_state = self._rolling_state(peer_email) - for event in rolling_state.events: - path = str(event.path_in_datasite) - if event.is_deleted: - files.pop(path, None) - elif event.new_hash is None or event.content is None: - raise ValueError( - f"The change of {path} from {peer_email} has no content or no hash." - ) - else: - files[path] = CheckpointFile( - path=path, hash=event.new_hash, content=event.content - ) - folded = Checkpoint( + 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. - email=rolling_state.events[-1].datasite_email, - last_event_timestamp=rolling_state.last_event_timestamp, - files=list(files.values()), - ) + 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) diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 3ffad11d4e7..926b172ef33 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -15,9 +15,7 @@ from syft.sync.checkpoints.checkpoint import ( DEFAULT_COMPACTING_THRESHOLD, Checkpoint, - CheckpointFile, IncrementalCheckpoint, - compact_incremental_checkpoints, ) from syft.sync.checkpoints.rolling_state import ( RollingState, @@ -1267,67 +1265,16 @@ 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: + ordered = sorted(incremental_cps, key=lambda c: c.sequence_number) + events = [event for checkpoint in ordered 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..cc54be0e811 --- /dev/null +++ b/syft/sync/utils/file_content.py @@ -0,0 +1,36 @@ +"""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"``.""" + if ( + isinstance(data, dict) + and data.get("content_type") == "binary" + and isinstance(data.get("content"), str) + ): + return {**data, "content": base64.b64decode(data["content"])} + return data diff --git a/tests/unit/test_checkpoints.py b/tests/unit/test_checkpoints.py index cd61b712452..a3183d55227 100644 --- a/tests/unit/test_checkpoints.py +++ b/tests/unit/test_checkpoints.py @@ -848,3 +848,11 @@ def test_rolling_state_of_later_version_refused(): 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]) From ff407d16a6cc165b44a498295babb9647d102126 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 11:18:20 -0300 Subject: [PATCH 08/12] fix: refuse file content that is not valid base64 - decode_content() dropped characters outside the base64 alphabet and returned other bytes, with no error. It now raises, so the model refuses the content. --- syft/sync/utils/file_content.py | 8 ++++++-- .../unit/test_file_change_events_serialization.py | 10 ++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/syft/sync/utils/file_content.py b/syft/sync/utils/file_content.py index cc54be0e811..fd994a1b3c2 100644 --- a/syft/sync/utils/file_content.py +++ b/syft/sync/utils/file_content.py @@ -26,11 +26,15 @@ def content_to_json(content: str | bytes | None) -> str | None: def decode_content(data: Any) -> Any: """The input of a model, with its base64 ``content`` decoded to bytes when - its ``content_type`` is ``"binary"``.""" + 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"])} + return {**data, "content": base64.b64decode(data["content"], validate=True)} return data 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. From 30347d856adc4bbe02fbeecfef566c2ddbe46671 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 11:32:14 -0300 Subject: [PATCH 09/12] refactor: read the checkpoint threshold from one constant - DEFAULT_CHECKPOINT_EVENT_THRESHOLD moved to syft/sync/checkpoints/checkpoint.py. The DO syncer, the DS cache and SyftboxManager all read it. - WATCHER_CHECKPOINT_THRESHOLD and three literal defaults of 50 in SyftboxManager are gone. --- syft/sync/checkpoints/checkpoint.py | 4 ++++ syft/sync/syftbox_manager.py | 11 ++++++++--- syft/sync/sync/caches/datasite_watcher_cache.py | 16 ++++++++++------ syft/sync/sync/datasite_owner_syncer.py | 4 +--- tests/unit/test_watcher_state.py | 10 +++++----- 5 files changed, 28 insertions(+), 17 deletions(-) diff --git a/syft/sync/checkpoints/checkpoint.py b/syft/sync/checkpoints/checkpoint.py index 02cbe47848f..b949845b5ba 100644 --- a/syft/sync/checkpoints/checkpoint.py +++ b/syft/sync/checkpoints/checkpoint.py @@ -60,6 +60,10 @@ 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.""" diff --git a/syft/sync/syftbox_manager.py b/syft/sync/syftbox_manager.py index a829ad1a198..81488eee766 100644 --- a/syft/sync/syftbox_manager.py +++ b/syft/sync/syftbox_manager.py @@ -11,6 +11,7 @@ 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 @@ -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, @@ -1396,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. @@ -1410,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 ebdad69b53c..52256d696a7 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -6,7 +6,11 @@ from pydantic import BaseModel, Field -from syft.sync.checkpoints.checkpoint import Checkpoint, CheckpointFile +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 ( @@ -78,9 +82,6 @@ def _accepted_events( SECONDS_BEFORE_SYNCING_DOWN = 0 -# The changes from a DO go into its checkpoint at this many paths, as the DO -# does with its own state. -WATCHER_CHECKPOINT_THRESHOLD = 50 class DataSiteWatcherCacheConfig(BaseModel): @@ -334,9 +335,12 @@ def _restore_file(self, datasite_email: str, file: CheckpointFile) -> None: def _save_state(self, peer_email: str) -> None: """Save the rolling state of ``peer_email``. At - ``WATCHER_CHECKPOINT_THRESHOLD`` paths, fold it into the checkpoint.""" + ``DEFAULT_CHECKPOINT_EVENT_THRESHOLD`` paths, fold it into the checkpoint.""" self._save_rolling_state(peer_email) - if self._rolling_state(peer_email).event_count >= WATCHER_CHECKPOINT_THRESHOLD: + 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: diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 926b172ef33..98c06f3d2f3 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -13,6 +13,7 @@ from syft.sync.callback_mixin import BaseModelCallbackMixin from syft.sync.checkpoints.checkpoint import ( + DEFAULT_CHECKPOINT_EVENT_THRESHOLD, DEFAULT_COMPACTING_THRESHOLD, Checkpoint, IncrementalCheckpoint, @@ -49,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. diff --git a/tests/unit/test_watcher_state.py b/tests/unit/test_watcher_state.py index c0d6965441a..e124cbd6b67 100644 --- a/tests/unit/test_watcher_state.py +++ b/tests/unit/test_watcher_state.py @@ -74,7 +74,7 @@ def test_restart_on_empty_folder_restores_files_and_cursor(tmp_path): def test_restart_after_checkpoint_keeps_deleted_file_deleted(tmp_path, monkeypatch): - monkeypatch.setattr(datasite_watcher_cache, "WATCHER_CHECKPOINT_THRESHOLD", 2) + 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") @@ -132,7 +132,7 @@ def test_restart_with_new_key_restores_files_and_cursor(tmp_path): def test_binary_file_restored_from_checkpoint(tmp_path, monkeypatch): - monkeypatch.setattr(datasite_watcher_cache, "WATCHER_CHECKPOINT_THRESHOLD", 1) + 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"} @@ -151,7 +151,7 @@ def test_binary_file_restored_from_checkpoint(tmp_path, monkeypatch): 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, "WATCHER_CHECKPOINT_THRESHOLD", 1) + 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]) @@ -163,7 +163,7 @@ def test_restore_writes_files_where_events_write(tmp_path, monkeypatch): def test_do_state_saved_under_email_without_case(monkeypatch): - monkeypatch.setattr(datasite_watcher_cache, "WATCHER_CHECKPOINT_THRESHOLD", 1) + 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 ) @@ -176,7 +176,7 @@ def test_do_state_saved_under_email_without_case(monkeypatch): def test_change_without_hash_or_content_dropped(monkeypatch): - monkeypatch.setattr(datasite_watcher_cache, "WATCHER_CHECKPOINT_THRESHOLD", 1) + 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} From 2e13bb9c0b422c94e99df8323dabf3fc6013a051 Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 12:45:45 -0300 Subject: [PATCH 10/12] fix: give the Drive watcher store its own Drive client in another thread - DriveWatcherStoreConnection.clone_for_thread() returned the same object, so another thread would share one Drive client. It now returns a store over a new client, as the other Drive connections do. --- syft/sync/connections/drive/drive_watcher_store.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/syft/sync/connections/drive/drive_watcher_store.py b/syft/sync/connections/drive/drive_watcher_store.py index a451a9b7baa..188fbf0889f 100644 --- a/syft/sync/connections/drive/drive_watcher_store.py +++ b/syft/sync/connections/drive/drive_watcher_store.py @@ -25,6 +25,10 @@ class DriveWatcherStoreConnection(WatcherStoreConnection): 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) -> str: return self._put(peer, CHECKPOINT_FILE_NAME, data) From 43196331fe0f54496126c82e998cf0dacd60033e Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 13:40:56 -0300 Subject: [PATCH 11/12] refactor: name the watcher store and cache helpers after what they handle - WatcherStoreConnection.latest_checkpoint() is now get_checkpoint(), because each DO has one checkpoint. put_checkpoint() and put_rolling_state() return nothing, because no caller read the file id. - DriveWatcherStoreConnection names its helpers after the watcher folder and the files: _watcher_folder_id(), _save_file(), _read_file() and _find_id(). - DataSiteWatcherCache has _save_watcher_state() and _restore_watcher_state(). The DO compaction names its sorted checkpoints incrementals_in_order. --- syft/sync/connections/connection_router.py | 2 +- .../connections/drive/drive_watcher_store.py | 64 ++++++++++--------- syft/sync/connections/protocols.py | 8 +-- .../sync/caches/datasite_watcher_cache.py | 11 ++-- syft/sync/sync/datasite_owner_syncer.py | 6 +- tests/unit/test_drive_connections.py | 4 +- 6 files changed, 50 insertions(+), 45 deletions(-) diff --git a/syft/sync/connections/connection_router.py b/syft/sync/connections/connection_router.py index 8f9882e28d8..6d3b22d449e 100644 --- a/syft/sync/connections/connection_router.py +++ b/syft/sync/connections/connection_router.py @@ -594,7 +594,7 @@ def watcher_load_checkpoint(self, peer_email: str) -> Checkpoint | None: Raises: ValueError: if a later client saved it. """ - data = self.connections.watcher_store_con.latest_checkpoint(peer_email) + data = self.connections.watcher_store_con.get_checkpoint(peer_email) if data is None: return None return Checkpoint.from_compressed_data(data) diff --git a/syft/sync/connections/drive/drive_watcher_store.py b/syft/sync/connections/drive/drive_watcher_store.py index 188fbf0889f..fcf90a44b79 100644 --- a/syft/sync/connections/drive/drive_watcher_store.py +++ b/syft/sync/connections/drive/drive_watcher_store.py @@ -30,20 +30,20 @@ def clone_for_thread(self) -> "DriveWatcherStoreConnection": 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) -> str: - return self._put(peer, CHECKPOINT_FILE_NAME, data) + 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 latest_checkpoint(self, peer: str) -> bytes | None: - return self._get(peer, CHECKPOINT_FILE_NAME) + 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) -> str: - return self._put(peer, ROLLING_STATE_FILE_NAME, data) + 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._get(peer, ROLLING_STATE_FILE_NAME) + 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: @@ -53,44 +53,46 @@ def reset(self, peer: str) -> None: @raise_unavailable("Could not list the watcher state on Drive.") def peers_with_saved_state(self) -> list[str]: - folder_id = self._folder_id() - if folder_id is None: + watcher_folder_id = self._watcher_folder_id() + if watcher_folder_id is None: return [] - peer_folders = self._cached_client.client.find_folders(parent_id=folder_id) + peer_folders = self._cached_client.client.find_folders( + parent_id=watcher_folder_id + ) return [name for _, name in peer_folders] - def _put(self, peer: str, name: str, data: bytes) -> str: + def _save_file(self, peer: str, name: str, data: bytes) -> None: client = self._cached_client.client - folder_id = self._folder_id() or self._create_folder() - peer_folder_id = self._find_one(peer, folder_id) or client.create_folder( - peer, folder_id + 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_one(name, peer_folder_id) + file_id = self._find_id(name, peer_folder_id) if file_id is None: body = {"name": name, "parents": [peer_folder_id]} - return client.create_syft_file(body, payload)["id"] - execute_with_retries( - client.drive_service.files().update(fileId=file_id, media_body=payload) - ) - return file_id + client.create_syft_file(body, payload) + else: + execute_with_retries( + client.drive_service.files().update(fileId=file_id, media_body=payload) + ) - def _get(self, peer: str, name: str) -> bytes | None: + 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_one(name, peer_folder_id) + 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: - folder_id = self._folder_id() - if folder_id is None: + watcher_folder_id = self._watcher_folder_id() + if watcher_folder_id is None: return None - return self._find_one(peer, folder_id) + return self._find_id(peer, watcher_folder_id) - def _find_one(self, name: str, parent_id: str) -> str | None: + 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: @@ -105,12 +107,12 @@ def _find_one(self, name: str, parent_id: str) -> str | None: ) return files[0]["id"] if files else None - def _folder_name(self) -> str: + def _watcher_folder_name(self) -> str: return ( f"{self._cached_client.email}-{SYFT_VERSION}{WATCHER_STATE_FOLDER_SUFFIX}" ) - def _folder_id(self) -> str | None: + def _watcher_folder_id(self) -> str | None: folders = self._cached_client.client.find_folders( name_contains=[ f"{self._cached_client.email}-", @@ -119,10 +121,10 @@ def _folder_id(self) -> str | None: parent_id=self._cached_client.get_syftbox_folder_id(), ) return self._cached_client.find_or_adopt_versioned_folder( - folders, current_name=self._folder_name() + folders, current_name=self._watcher_folder_name() ) - def _create_folder(self) -> str: + def _create_watcher_folder(self) -> str: return self._cached_client.client.create_folder( - self._folder_name(), self._cached_client.get_syftbox_folder_id() + 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 90fc3a95af9..25a0497dad3 100644 --- a/syft/sync/connections/protocols.py +++ b/syft/sync/connections/protocols.py @@ -204,16 +204,16 @@ class WatcherStoreConnection(RouterConnection, ABC): """ @abstractmethod - def put_checkpoint(self, peer: 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, 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; it replaces the previous one.""" diff --git a/syft/sync/sync/caches/datasite_watcher_cache.py b/syft/sync/sync/caches/datasite_watcher_cache.py index 52256d696a7..09cd3c803d7 100644 --- a/syft/sync/sync/caches/datasite_watcher_cache.py +++ b/syft/sync/sync/caches/datasite_watcher_cache.py @@ -160,7 +160,7 @@ def from_config( saves_watcher_state=True, ) cache._load_collection_hashes_from_disk() - cache._restore_saved_state() + cache._restore_watcher_state() return cache def _load_collection_hashes_from_disk(self): @@ -248,7 +248,7 @@ def sync_down(self, peer_email: str) -> int: event_count += len(event_message.events) self._record_received(peer_email, raw.msg_id) if received and self.saves_watcher_state: - self._save_state(peer_email) + self._save_watcher_state(peer_email) self.connection_router.mark_all_processed([raw for raw, _ in received]) self.last_sync = datetime.now() @@ -307,7 +307,7 @@ def _apply_event(self, event: FileChangeEvent) -> None: # --- State saved in the DS's storage --- - def _restore_saved_state(self) -> None: + 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. @@ -333,7 +333,7 @@ def _restore_file(self, datasite_email: str, file: CheckpointFile) -> None: self.file_connection.write_file(str(path_in_syftbox), file.content) self.file_hashes[path_in_syftbox] = file.hash - def _save_state(self, peer_email: str) -> None: + 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) @@ -364,8 +364,9 @@ def _fold_into_checkpoint(self, peer_email: str) -> None: 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": sorted(self.received_msg_ids_per_peer[key])} + update={"received_msg_ids": received_msg_ids} ) self.connection_router.watcher_save_rolling_state(key, saved) diff --git a/syft/sync/sync/datasite_owner_syncer.py b/syft/sync/sync/datasite_owner_syncer.py index 98c06f3d2f3..1f058ebb9fc 100644 --- a/syft/sync/sync/datasite_owner_syncer.py +++ b/syft/sync/sync/datasite_owner_syncer.py @@ -1263,8 +1263,10 @@ def compact_checkpoints(self) -> Checkpoint: print(f"Found {len(incremental_cps)} incremental checkpoints to compact") - ordered = sorted(incremental_cps, key=lambda c: c.sequence_number) - events = [event for checkpoint in ordered for event in checkpoint.events] + 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: diff --git a/tests/unit/test_drive_connections.py b/tests/unit/test_drive_connections.py index 1479841468e..02ab0373820 100644 --- a/tests/unit/test_drive_connections.py +++ b/tests/unit/test_drive_connections.py @@ -552,13 +552,13 @@ def test_owner_store_files_carry_the_syft_tag(pair): def test_watcher_state_save_replaces_previous(pair): _, ds_cached_client = pair store = DriveWatcherStoreConnection(ds_cached_client) - assert store.latest_checkpoint(DO) is None + 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.latest_checkpoint(DO) == b"checkpoint 2" + assert store.get_checkpoint(DO) == b"checkpoint 2" assert store.get_rolling_state(DO) == b"rolling state" From d369c28f425f3ed47f01575c91633a1c24e294ab Mon Sep 17 00:00:00 2001 From: Pedro Werneck Date: Wed, 4 Nov 2026 16:58:36 -0300 Subject: [PATCH 12/12] test: restore the DS from its own saved state in the deletion-reload test - test_ds_cache_handles_deletions_correctly reloaded over an empty mock Drive, so it passed even with no restore. It now reloads over the DS's own router. - The test also checks that a kept file comes back, so a restore that does nothing fails it. --- tests/unit/test_cache_aware_sync.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/unit/test_cache_aware_sync.py b/tests/unit/test_cache_aware_sync.py index 8851e112c1d..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, @@ -208,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 ) @@ -219,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() @@ -258,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" )