diff --git a/synapse/handlers/sliding_sync/__init__.py b/synapse/handlers/sliding_sync/__init__.py index b7affef811..f79056d3fa 100644 --- a/synapse/handlers/sliding_sync/__init__.py +++ b/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, diff --git a/synapse/handlers/sliding_sync/extensions.py b/synapse/handlers/sliding_sync/extensions.py index 0348f96115..f009b0d330 100644 --- a/synapse/handlers/sliding_sync/extensions.py +++ b/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: """ diff --git a/synapse/rest/client/sync.py b/synapse/rest/client/sync.py index b970d742fc..4ac3ea7a1b 100644 --- a/synapse/rest/client/sync.py +++ b/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 diff --git a/synapse/synapse_rust/sliding_sync.pyi b/synapse/synapse_rust/sliding_sync.pyi index 5656c8e8e2..537b0efba0 100644 --- a/synapse/synapse_rust/sliding_sync.pyi +++ b/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.""" diff --git a/synapse/types/handlers/sliding_sync.py b/synapse/types/handlers/sliding_sync.py index fadeb24222..b139dce376 100644 --- a/synapse/types/handlers/sliding_sync.py +++ b/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`.