Serve sliding sync rooms from Rust in rust_sliding_sync: on mode

current_sync_for_user hands the rooms to send to Rust and applies the
connection-state changes it returns. Python computes the rooms Rust
leaves to it or fails on. The servlet encodes Rust rooms in Rust, and the
extensions read rooms through accessors both kinds of room have.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01L51DoxB8VnH8UjXJFDgpNG

synapse/handlers/sliding_sync/__init__.py

@@ -21,6 +21,8 @@ import attr
 from prometheus_client import Counter, Histogram
 from typing_extensions import assert_never
 
+from twisted.internet.defer import CancelledError
+
 from synapse.api.constants import Direction, EventTypes, Membership
 from synapse.events import EventBase
 from synapse.events.utils import FilteredEvent, strip_event
@@ -28,6 +30,7 @@ from synapse.handlers.relations import BundledAggregations
 from synapse.handlers.sliding_sync.extensions import SlidingSyncExtensionHandler
 from synapse.handlers.sliding_sync.room_lists import (
     RoomsForUserType,
+    SlidingSyncInterestedRooms,
     SlidingSyncRoomLists,
 )
 from synapse.handlers.sliding_sync.store import SlidingSyncConnectionStore
@@ -47,6 +50,7 @@ from synapse.storage.roommember import (
     MemberSummary,
 )
 from synapse.synapse_rust.sliding_sync import (
+    RoomInput as RustRoomInput,
     SlidingSyncHandler as RustSlidingSyncHandler,
     SlidingSyncRequest as RustSlidingSyncRequest,
     SlidingSyncStats as RustSlidingSyncStats,
@@ -66,6 +70,7 @@ from synapse.types import (
 )
 from synapse.types.handlers import SLIDING_SYNC_DEFAULT_BUMP_EVENT_TYPES
 from synapse.types.handlers.sliding_sync import (
+    AnySlidingSyncRoomResult,
     HaveSentRoom,
     HaveSentRoomFlag,
     MutablePerConnectionState,
@@ -444,7 +449,7 @@ class SlidingSyncHandler:
         relevant_rooms_to_send_map = interested_rooms.relevant_rooms_to_send_map
 
         # Fetch room data
-        rooms: dict[str, SlidingSyncResult.RoomResult] = {}
+        rooms: dict[str, AnySlidingSyncRoomResult] = {}
 
         new_connection_state = previous_connection_state.get_mutable()
 
@@ -473,7 +478,20 @@ class SlidingSyncHandler:
 
         if relevant_rooms_to_send_map:
             with start_active_span("sliding_sync.generate_room_entries"):
-                await concurrently_execute(handle_room, relevant_rooms_to_send_map, 20)
+                python_room_ids: StrCollection = list(relevant_rooms_to_send_map)
+                if self._rust_sliding_sync is not None:
+                    python_room_ids = await self._generate_room_entries_in_rust(
+                        self._rust_sliding_sync,
+                        sync_config=sync_config,
+                        previous_connection_state=previous_connection_state,
+                        new_connection_state=new_connection_state,
+                        interested_rooms=interested_rooms,
+                        from_token=from_token,
+                        to_token=to_token,
+                        rooms=rooms,
+                    )
+                if python_room_ids:
+                    await concurrently_execute(handle_room, python_room_ids, 20)
 
         extensions = await self.extensions.get_extensions_response(
             sync_config=sync_config,
@@ -567,6 +585,137 @@ class SlidingSyncHandler:
 
         return sliding_sync_result
 
+    async def _generate_room_entries_in_rust(
+        self,
+        rust: RustSlidingSyncHandler,
+        *,
+        sync_config: SlidingSyncConfig,
+        previous_connection_state: PerConnectionState,
+        new_connection_state: MutablePerConnectionState,
+        interested_rooms: SlidingSyncInterestedRooms,
+        from_token: SlidingSyncStreamToken | None,
+        to_token: StreamToken,
+        rooms: dict[str, AnySlidingSyncRoomResult],
+    ) -> list[str]:
+        """`get_room_sync_data` in Rust for every room of
+        `relevant_rooms_to_send_map`. Adds the entries to `rooms` and the
+        rooms' changes to `new_connection_state`.
+
+        Returns:
+            The rooms Python must compute: those Rust leaves to Python, and
+            those it failed on (outside strict mode). All of them if the Rust
+            call as a whole failed.
+        """
+        rooms_to_send = interested_rooms.relevant_rooms_to_send_map
+        labels = {SERVER_NAME_LABEL: self.server_name}
+        try:
+            request = await self._build_rust_request(sync_config, from_token, to_token)
+        except ValueError as e:
+            rust_sliding_sync_requests_counter.labels(result="fallback", **labels).inc()
+            logger.debug("Computing the sliding sync rooms in Python: %s", e)
+            return list(rooms_to_send)
+
+        # The writers that the rooms' `prev_batch` tokens may name: those of
+        # the tokens the timelines are bounded by.
+        instance_names = set(to_token.room_key.instance_map)
+        if from_token is not None:
+            instance_names.update(from_token.stream_token.room_key.instance_map)
+
+        room_inputs = []
+        for room_id, room_sync_config in rooms_to_send.items():
+            membership = interested_rooms.room_membership_for_user_map[room_id]
+            prev_room_sync_config = previous_connection_state.room_configs.get(room_id)
+            room_status = previous_connection_state.rooms.have_sent_room(room_id)
+            if room_status.last_token is not None:
+                instance_names.update(room_status.last_token.instance_map)
+            room_inputs.append(
+                RustRoomInput(
+                    room_id=room_id,
+                    sender=membership.sender,
+                    membership=membership.membership,
+                    event_id=membership.event_id,
+                    instance_name=membership.event_pos.instance_name,
+                    stream=membership.event_pos.stream,
+                    room_version_id=membership.room_version_id,
+                    timeline_limit=room_sync_config.timeline_limit,
+                    required_state_map=room_sync_config.required_state_map,
+                    prev_room_sync_config=(
+                        (
+                            prev_room_sync_config.timeline_limit,
+                            prev_room_sync_config.required_state_map,
+                        )
+                        if prev_room_sync_config is not None
+                        else None
+                    ),
+                    room_status=room_status.status.value,
+                    room_status_token=room_status.last_token,
+                    newly_joined=room_id in interested_rooms.newly_joined_rooms,
+                    newly_left=room_id in interested_rooms.newly_left_rooms,
+                    is_dm=room_id in interested_rooms.dm_room_ids,
+                )
+            )
+        instance_ids = {
+            name: await self.store.get_id_for_instance(name) for name in instance_names
+        }
+
+        try:
+            result = await rust.generate_room_entries(
+                request, room_inputs, instance_ids=instance_ids
+            )
+        except CancelledError:
+            raise
+        except Exception:
+            rust_sliding_sync_requests_counter.labels(result="error", **labels).inc()
+            if self._rust_sliding_sync_strict:
+                raise
+            logger.exception(
+                "Rust sliding sync failed for %s. Computing the rooms in Python",
+                sync_config.user.to_string(),
+            )
+            return list(rooms_to_send)
+
+        rust_sliding_sync_requests_counter.labels(result="rust", **labels).inc()
+        observe_rust_sliding_sync_stats(result.stats, self.server_name)
+
+        for room_id, timeline_limit, required_state_map in result.room_config_updates:
+            new_connection_state.room_configs[room_id] = RoomSyncConfig(
+                timeline_limit=timeline_limit, required_state_map=required_state_map
+            )
+        for room_id, returned, invalidated in result.lazy_memberships:
+            new_connection_state.room_lazy_membership[room_id] = (
+                RoomLazyMembershipChanges(
+                    returned_user_id_to_last_seen_ts_map=returned,
+                    invalidated_user_ids=invalidated,
+                )
+            )
+
+        entries = result.take_entries()
+        for entry in entries:
+            rooms[entry.room_id] = entry
+        rust_sliding_sync_rooms_counter.labels(result="rust", **labels).inc(
+            len(entries) + len(result.empty)
+        )
+
+        python_room_ids = []
+        for room_id, reason in result.fallbacks:
+            rust_sliding_sync_rooms_counter.labels(result="fallback", **labels).inc()
+            logger.debug("Computing %s in Python: %s", room_id, reason)
+            python_room_ids.append(room_id)
+        for room_id, error in result.errors:
+            rust_sliding_sync_rooms_counter.labels(result="error", **labels).inc()
+            if self._rust_sliding_sync_strict:
+                raise RuntimeError(f"Rust sliding sync failed for {room_id}: {error}")
+            # Rust only reads, and its changes to the connection state are
+            # for the rooms it finished, so Python can compute the room from
+            # scratch.
+            logger.error(
+                "Rust sliding sync failed for %s: %s. Computing it in Python",
+                room_id,
+                error,
+            )
+            python_room_ids.append(room_id)
+        return python_room_ids
+
     @trace
     async def get_current_state_ids_at(
         self,

synapse/handlers/sliding_sync/extensions.py

@@ -52,6 +52,7 @@ from synapse.types import (
     UserID,
 )
 from synapse.types.handlers.sliding_sync import (
+    AnySlidingSyncRoomResult,
     HaveSentRoomFlag,
     MutablePerConnectionState,
     OperationType,
@@ -103,7 +104,7 @@ class SlidingSyncExtensionHandler:
         all_interested_room_ids: set[str],
         actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList],
         actual_room_ids: set[str],
-        actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult],
+        actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult],
         to_token: StreamToken,
         from_token: SlidingSyncStreamToken | None,
     ) -> SlidingSyncResult.Extensions:
@@ -690,7 +691,7 @@ class SlidingSyncExtensionHandler:
         new_connection_state: "MutablePerConnectionState",
         actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList],
         actual_room_ids: set[str],
-        actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult],
+        actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult],
         receipts_request: SlidingSyncConfig.Extensions.ReceiptsExtension,
         to_token: StreamToken,
         from_token: SlidingSyncStreamToken | None,
@@ -815,10 +816,10 @@ class SlidingSyncExtensionHandler:
                 # in the timeline to avoid bloating and blowing up the sync response
                 # as the number of users in the room increases. (this behavior is part of the spec)
                 initial_rooms_and_event_ids = [
-                    (room_id, event.event.event_id)
+                    (room_id, event_id)
                     for room_id in initial_rooms
                     if room_id in actual_room_response_map
-                    for event in actual_room_response_map[room_id].timeline_events
+                    for event_id in actual_room_response_map[room_id].timeline_event_ids
                 ]
                 initial_receipts = await self.store.get_linearized_receipts_for_events(
                     room_and_event_ids=initial_rooms_and_event_ids,
@@ -894,7 +895,7 @@ class SlidingSyncExtensionHandler:
         sync_config: SlidingSyncConfig,
         actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList],
         actual_room_ids: set[str],
-        actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult],
+        actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult],
         typing_request: SlidingSyncConfig.Extensions.TypingExtension,
         to_token: StreamToken,
         from_token: SlidingSyncStreamToken | None,
@@ -1129,7 +1130,7 @@ class SlidingSyncExtensionHandler:
         user_id: str,
         actual_room_ids: set[str],
         sync_config: SlidingSyncConfig,
-        actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult],
+        actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult],
         actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList],
     ) -> tuple[set[str], set[str]]:
         """
@@ -1186,16 +1187,11 @@ class SlidingSyncExtensionHandler:
                 if room_id not in lazy_rooms:
                     continue
                 # Include users from timeline events
-                for timeline_event in room_data.timeline_events:
-                    lazy_profile_user_ids.add(timeline_event.event.sender)
+                lazy_profile_user_ids.update(room_data.timeline_senders)
                 # Include users from required state
-                for state_event in room_data.required_state:
-                    if state_event.type == EventTypes.Member:
-                        lazy_profile_user_ids.add(state_event.state_key)
+                lazy_profile_user_ids.update(room_data.required_state_member_keys)
                 # Include heroes
-                if room_data.heroes:
-                    for hero in room_data.heroes:
-                        lazy_profile_user_ids.add(hero.user_id)
+                lazy_profile_user_ids.update(room_data.hero_user_ids)
 
         non_lazy_rooms = actual_room_ids.difference(lazy_rooms)
         # If we still have non-lazy rooms, get their members.
@@ -1290,7 +1286,7 @@ class SlidingSyncExtensionHandler:
         actual_room_ids: set[str],
         to_token: StreamToken,
         from_token: SlidingSyncStreamToken | None,
-        actual_room_response_map: Mapping[str, SlidingSyncResult.RoomResult],
+        actual_room_response_map: Mapping[str, AnySlidingSyncRoomResult],
         actual_lists: Mapping[str, SlidingSyncResult.SlidingWindowList],
     ) -> SlidingSyncResult.Extensions.ProfilesExtension | None:
         """

synapse/rest/client/sync.py

@@ -63,6 +63,10 @@ from synapse.http.servlet import (
 from synapse.http.site import SynapseRequest
 from synapse.logging.opentracing import log_kv, set_tag, trace_with_opname
 from synapse.rest.admin.experimental_features import ExperimentalFeature
+from synapse.synapse_rust.sliding_sync import (
+    RoomResult as RustSlidingSyncRoomResult,
+    encode_rooms as encode_rust_sliding_sync_rooms,
+)
 from synapse.synapse_rust.sync import RoomsEncoder as RustRoomsEncoder
 from synapse.types import (
     JsonDict,
@@ -71,6 +75,7 @@ from synapse.types import (
     SlidingSyncStreamToken,
     StreamToken,
 )
+from synapse.types.handlers.sliding_sync import AnySlidingSyncRoomResult
 from synapse.types.rest.client import SlidingSyncBody
 from synapse.util.caches.lrucache import LruCache
 from synapse.util.cancellation import cancellable
@@ -1106,7 +1111,7 @@ class SlidingSyncRestServlet(RestServlet):
     async def encode_rooms(
         self,
         requester: Requester,
-        rooms: dict[str, SlidingSyncResult.RoomResult],
+        rooms: Mapping[str, AnySlidingSyncRoomResult],
     ) -> JsonDict:
         time_now = self.clock.time_msec()
 
@@ -1115,8 +1120,37 @@ class SlidingSyncRestServlet(RestServlet):
             requester=requester,
         )
 
+        # The rooms the Rust handler computed are encoded in Rust.
+        rust_rooms = [
+            room
+            for room in rooms.values()
+            if isinstance(room, RustSlidingSyncRoomResult)
+        ]
+        encoded_in_rust: dict[str, JsonDict] = {}
+        if rust_rooms:
+            # Module callbacks take Python events, so only then do the rooms'
+            # events become Python objects.
+            unsigned_additions: Mapping[str, JsonDict] = {}
+            if self.event_serializer.has_unsigned_callbacks:
+                unsigned_additions = await self.event_serializer.unsigned_additions(
+                    event for room in rust_rooms for event in room.events_to_serialize()
+                )
+            encoded = encode_rust_sliding_sync_rooms(
+                rust_rooms,
+                time_now,
+                serialize_options,
+                unsigned_additions=unsigned_additions,
+            )
+            encoded_in_rust = {
+                room.room_id: entry for room, entry in zip(rust_rooms, encoded)
+            }
+
         serialized_rooms: dict[str, JsonDict] = {}
         for room_id, room_result in rooms.items():
+            if isinstance(room_result, RustSlidingSyncRoomResult):
+                serialized_rooms[room_id] = encoded_in_rust[room_id]
+                continue
+
             serialized_rooms[room_id] = {
                 "notification_count": room_result.notification_count,
                 "highlight_count": room_result.highlight_count,
@@ -1228,7 +1262,7 @@ class SlidingSyncRestServlet(RestServlet):
         self,
         requester: Requester,
         extensions: SlidingSyncResult.Extensions,
-        ref_rooms_results: Mapping[str, SlidingSyncResult.RoomResult],
+        ref_rooms_results: Mapping[str, AnySlidingSyncRoomResult],
     ) -> JsonDict:
         """
         Args:
@@ -1343,7 +1377,7 @@ class SlidingSyncRestServlet(RestServlet):
         self,
         requester: Requester,
         sticky_events: SlidingSyncResult.Extensions.StickyEventsExtension,
-        ref_rooms_results: Mapping[str, SlidingSyncResult.RoomResult],
+        ref_rooms_results: Mapping[str, AnySlidingSyncRoomResult],
     ) -> JsonDict:
         """
         Serialise the sticky events extension response.
@@ -1385,9 +1419,7 @@ class SlidingSyncRestServlet(RestServlet):
                 # Nothing to deduplicate
                 sticky_events_to_write = possibly_duplicated_sticky_events
             else:
-                sent_event_ids_in_room_section = {
-                    ev.event.event_id for ev in room_result.timeline_events
-                }
+                sent_event_ids_in_room_section = set(room_result.timeline_event_ids)
                 sticky_events_to_write = [
                     ev
                     for ev in possibly_duplicated_sticky_events

synapse/synapse_rust/sliding_sync.pyi

@@ -12,7 +12,14 @@
 
 """Simplified sliding sync (MSC4186) in Rust. See `rust/src/sliding_sync/`."""
 
+from collections.abc import Mapping
+from typing import AbstractSet
+
+from twisted.internet.defer import Deferred
+
+from synapse.events import EventBase
 from synapse.synapse_rust.database.postgres import ConnectionPool
+from synapse.synapse_rust.events import SerializeEventConfig
 from synapse.types import JsonDict, JsonMapping, RoomStreamToken
 from synapse.types.state import StateFilter
 
@@ -20,6 +27,107 @@ class SlidingSyncHandler:
     """Computes sliding sync responses in Rust, reading the database through
     the native Postgres driver. Built by `RustHandlers`."""
 
+    def generate_room_entries(
+        self,
+        request: SlidingSyncRequest,
+        rooms: list[RoomInput],
+        *,
+        instance_ids: Mapping[str, int],
+    ) -> Deferred[RoomsResult]:
+        """`get_room_sync_data` for each of `rooms`, twenty at a time.
+        `instance_ids` are `get_id_for_instance` for the writers the rooms'
+        tokens may name."""
+    def generate_room_entries_blocking(
+        self,
+        request: SlidingSyncRequest,
+        rooms: list[RoomInput],
+        *,
+        instance_ids: Mapping[str, int],
+    ) -> RoomsResult: ...
+
+class RoomInput:
+    """The input of one room's `get_room_sync_data`."""
+
+    def __init__(
+        self,
+        *,
+        room_id: str,
+        sender: str | None,
+        membership: str,
+        event_id: str | None,
+        instance_name: str,
+        stream: int,
+        room_version_id: str | None,
+        timeline_limit: int,
+        required_state_map: Mapping[str, AbstractSet[str]],
+        prev_room_sync_config: tuple[int, Mapping[str, AbstractSet[str]]] | None,
+        room_status: str,
+        room_status_token: RoomStreamToken | None,
+        newly_joined: bool,
+        newly_left: bool,
+        is_dm: bool,
+    ) -> None:
+        """`room_status` is a `HaveSentRoomFlag` value, with its token for
+        "previously"."""
+    @property
+    def room_id(self) -> str: ...
+
+class RoomResult:
+    """One room's entry, computed in Rust: what `current_sync_for_user` and the
+    extensions read of a `SlidingSyncResult.RoomResult`."""
+
+    def __bool__(self) -> bool: ...
+    @property
+    def room_id(self) -> str: ...
+    @property
+    def initial(self) -> bool: ...
+    @property
+    def unstable_expanded_timeline(self) -> bool: ...
+    @property
+    def timeline_event_ids(self) -> list[str]: ...
+    @property
+    def timeline_senders(self) -> list[str]: ...
+    @property
+    def required_state_member_keys(self) -> list[str]: ...
+    @property
+    def hero_user_ids(self) -> list[str]: ...
+    def events_to_serialize(self) -> list[EventBase]:
+        """Every event the room's encoding serializes, for module callbacks."""
+
+class RoomsResult:
+    """What `generate_room_entries` made of the rooms."""
+
+    def take_entries(self) -> list[RoomResult]:
+        """The entries to send. A second call returns an empty list."""
+    @property
+    def empty(self) -> list[str]:
+        """The rooms computed with nothing to send."""
+    @property
+    def room_config_updates(self) -> list[tuple[str, int, dict[str, set[str]]]]:
+        """`(room_id, timeline_limit, required_state_map)` for each room whose
+        config the connection records."""
+    @property
+    def lazy_memberships(
+        self,
+    ) -> list[tuple[str, dict[str, int | None], set[str]]]:
+        """`(room_id, returned_user_id_to_last_seen_ts_map,
+        invalidated_user_ids)` for each room computed."""
+    @property
+    def fallbacks(self) -> list[tuple[str, str]]: ...
+    @property
+    def errors(self) -> list[tuple[str, str]]: ...
+    @property
+    def stats(self) -> SlidingSyncStats: ...
+
+def encode_rooms(
+    rooms: list[RoomResult],
+    time_now_ms: int,
+    config: SerializeEventConfig,
+    *,
+    unsigned_additions: Mapping[str, JsonDict] | None = None,
+) -> list[JsonDict]:
+    """The servlet's `encode_rooms` for rooms computed in Rust, in order."""
+
 class SlidingSyncStats:
     """What the Rust implementation did for one request."""
 

synapse/types/handlers/sliding_sync.py

@@ -25,6 +25,7 @@ from typing import (
     Mapping,
     MutableMapping,
     Sequence,
+    TypeAlias,
     TypeVar,
     cast,
 )
@@ -35,6 +36,7 @@ from pydantic import ConfigDict
 from synapse.api.constants import EventTypes
 from synapse.events import EventBase
 from synapse.events.utils import FilteredEvent
+from synapse.synapse_rust.sliding_sync import RoomResult as RustRoomResult
 from synapse.types import (
     DeviceListUpdates,
     JsonDict,
@@ -205,6 +207,31 @@ class SlidingSyncResult:
         notification_count: int
         highlight_count: int
 
+        @property
+        def timeline_event_ids(self) -> list[str]:
+            """The IDs of the timeline's events. The extensions read this,
+            which a room computed in Rust has too."""
+            return [event.event.event_id for event in self.timeline_events]
+
+        @property
+        def timeline_senders(self) -> list[str]:
+            """The senders of the timeline's events."""
+            return [event.event.sender for event in self.timeline_events]
+
+        @property
+        def required_state_member_keys(self) -> list[str]:
+            """The state keys of the membership events in `required_state`."""
+            return [
+                event.state_key
+                for event in self.required_state
+                if event.type == EventTypes.Member
+            ]
+
+        @property
+        def hero_user_ids(self) -> list[str]:
+            """The user IDs of the heroes."""
+            return [hero.user_id for hero in self.heroes or ()]
+
         def __bool__(self) -> bool:
             """Are there any updates that should be returned immediately to
             the client?
@@ -483,7 +510,8 @@ class SlidingSyncResult:
 
     next_pos: SlidingSyncStreamToken
     lists: Mapping[str, SlidingWindowList]
-    rooms: dict[str, RoomResult]
+    # A room computed in Rust is a `synapse_rust.sliding_sync.RoomResult`.
+    rooms: dict[str, "AnySlidingSyncRoomResult"]
     extensions: Extensions
 
     def __bool__(self) -> bool:
@@ -513,6 +541,11 @@ class SlidingSyncResult:
         )
 
 
+# A room's entry in a response: computed in Python, or in Rust, which has the
+# properties of the Python one that the extensions read.
+AnySlidingSyncRoomResult: TypeAlias = SlidingSyncResult.RoomResult | RustRoomResult
+
+
 class StateValues:
     """
     Understood values of the (type, state_key) tuple in `required_state`.