diff --git a/synapse/storage/databases/main/event_federation.py b/synapse/storage/databases/main/event_federation.py index 415926eb0a..bf08824eb4 100644 --- a/synapse/storage/databases/main/event_federation.py +++ b/synapse/storage/databases/main/event_federation.py @@ -37,7 +37,7 @@ from prometheus_client import Counter, Gauge from synapse.api.constants import MAX_DEPTH from synapse.api.errors import StoreError from synapse.api.room_versions import EventFormatVersions, RoomVersion -from synapse.events import EventBase, make_event_from_dict +from synapse.events import EventBase, FrozenEventVMSC4242, make_event_from_dict from synapse.logging.opentracing import tag_args, trace from synapse.metrics import SERVER_NAME_LABEL from synapse.metrics.background_process_metrics import wrap_as_background_process @@ -1194,6 +1194,116 @@ class EventFederationWorkerStore( # Return all events where not all sets can reach them. return {eid for eid, n in event_to_missing_sets.items() if n} + # 22/04/2026: Unused currently, but will be used in future MSC4242 PRs. + # Remove comment when used. + async def get_state_dag( + self, room_id: str, forward_extrems: set[str] + ) -> dict[str, FrozenEventVMSC4242]: + """Get the current state DAG for the given room. + + This function is called when calculating a /send_join response. + This does not check that the room is an state DAG room, so check this + before calling this function! + + This functions guarantees that the returned state DAG is connected. + + Args: + room_id: The room to get the state dag for + forward_extrems: latest event IDs in the room. The state DAG is all events reachable from these events. + Returns: + A map of event_id => event + """ + + def _get_state_events_txn(txn: LoggingTransaction, room_id: str) -> list[str]: + sql = """ + SELECT event_id FROM msc4242_state_dag_edges WHERE room_id = ? + """ + txn.execute(sql, (room_id,)) + event_ids = [ev_id for (ev_id,) in txn] + return event_ids + + # Pull out all state events for this room using the events_by_room_and_type index. + event_ids = await self.db_pool.runInteraction( + "_get_state_events_txn", + _get_state_events_txn, + room_id, + ) + event_map = await self.get_events(event_ids) + # Filter the returned state events to only include ones on the paths back from the forward + # extremities. + result: dict[str, FrozenEventVMSC4242] = {} + next = forward_extrems + seen: set[str] = set() + while len(next) > 0: + # Pull the event and add the prev_state_events. + # We must have the event. + event_id = next.pop() + if event_id in seen: + continue + seen.add(event_id) + ev = event_map[event_id] + assert isinstance(ev, FrozenEventVMSC4242) + result[event_id] = ev + for prev_state_event_id in ev.prev_state_events: + next.add(prev_state_event_id) + return result + + # 22/04/2026: Unused currently, but will be used in future MSC4242 PRs. + # Remove comment when used. + async def get_missing_events_state_dag( + self, + room_id: str, + earliest_events: list[str], + latest_events: list[str], + limit: int, + ) -> list[EventBase]: + ids = await self.db_pool.runInteraction( + "get_missing_events_state_dag", + self._get_missing_events_state_dag_txn, + room_id, + earliest_events, + latest_events, + limit, + ) + return await self.get_events_as_list(ids) + + # 22/04/2026: Unused currently, but will be used in future MSC4242 PRs. + # Remove comment when used. + def _get_missing_events_state_dag_txn( + self, + txn: LoggingTransaction, + room_id: str, + earliest_events: list[str], + latest_events: list[str], + limit: int, + ) -> list[str]: + seen_events = set(earliest_events) + # lexicographical sort + front_queue = sorted(set(latest_events) - seen_events) + event_results: list[str] = [] + # TODO(kegan): use a recursive CTE? + # The limit is usually pretty low, so it's cheaper to select the events we need via querying + # rather than selecting all events and filtering python-side. + query = ( + "SELECT prev_state_event_id FROM msc4242_state_dag_edges " + "WHERE room_id = ? AND event_id = ? " + "ORDER BY prev_state_event_id ASC " + "LIMIT ?" + ) + while front_queue and len(event_results) < limit: + event_id = front_queue.pop(0) + txn.execute(query, (room_id, event_id, limit - len(event_results))) + # None check because the m.room.create event has NULL prev_state_events + new_results = [ + t[0] for t in txn if t[0] is not None and t[0] not in seen_events + ] + for next in new_results: + front_queue.append(next) + seen_events |= set(new_results) + event_results.extend(new_results) + + return event_results + @trace @tag_args async def get_backfill_points_in_room( diff --git a/tests/storage/test_event_federation.py b/tests/storage/test_event_federation.py index 59e914ca8b..aa29f467c7 100644 --- a/tests/storage/test_event_federation.py +++ b/tests/storage/test_event_federation.py @@ -19,6 +19,7 @@ # import datetime +from collections import namedtuple from typing import ( Collection, Iterable, @@ -39,13 +40,14 @@ from synapse.api.room_versions import ( EventFormatVersions, RoomVersion, ) -from synapse.events import EventBase +from synapse.events import EventBase, FrozenEventVMSC4242 from synapse.rest import admin from synapse.rest.client import login, room from synapse.server import HomeServer from synapse.storage.database import LoggingTransaction from synapse.storage.types import Cursor from synapse.synapse_rust.events import EventInternalMetadata +from synapse.synapse_rust.room_versions import RoomVersions from synapse.types import JsonDict from synapse.util.clock import Clock from synapse.util.json import json_encoder @@ -1414,6 +1416,178 @@ class EventFederationWorkerStoreTestCase(tests.unittest.HomeserverTestCase): # elapsed past the backoff range so there is no events to backoff from. self.assertEqual(event_ids_with_backoff, {}) + @tests.unittest.override_config( + {"experimental_features": {"msc4242_enabled": True}} + ) + def test_get_state_dag(self) -> None: + """ + Test that MSC4242 state dag rooms can return the complete state dag on request. + """ + # Create the room + user_id = self.register_user("alice", "test") + tok = self.login("alice", "test") + room_id = self.helper.create_room_as( + room_creator=user_id, + tok=tok, + room_version=RoomVersions.MSC4242v12.identifier, + ) + resp = self.helper.send_state( + room_id, + "m.room.join_rules", + {"join_rule": "knock"}, + tok=tok, + ) + latest = resp["event_id"] + state_dag = self.get_success( + self.store.get_state_dag(room_id, {latest}), + ) + # create <- member <- pl <- join_rules <- his vis <- join_rules + self.assertEquals(len(state_dag), 6) + want_types = [ + EventTypes.Create, + EventTypes.Member, + EventTypes.PowerLevels, + EventTypes.JoinRules, + EventTypes.RoomHistoryVisibility, + EventTypes.JoinRules, + ] + curr = {latest} + while len(curr) > 0: + event_id = curr.pop() + ev = state_dag[event_id] + want_type = want_types.pop() + self.assertEqual(ev.type, want_type) + curr.update(ev.prev_state_events) + + @tests.unittest.override_config( + {"experimental_features": {"msc4242_enabled": True}} + ) + def test_get_missing_events_state_dag(self) -> None: + # Primarily testing to make sure that we sort events + # correctly when there are multiple prev_state_events + # .- C -- D ---. + # A <- B E + # `- R -- W --` + # `-- T -` + graph = { + "A": [], + "B": ["A"], + "C": ["B"], + "R": ["B"], + "D": ["C"], + "W": ["R"], + "T": ["R"], + "E": ["W", "D", "T"], + } + room_id = "@state_dag:local" + events = [ + cast( + FrozenEventVMSC4242, + StateDAGFakeEvent( + event_id, + room_id, + EventTypes.Create if len(graph[event_id]) == 0 else "foo", + graph[event_id], + ), + ) + for event_id in graph + ] + + def insert(txn: LoggingTransaction) -> None: + # store these to satisfy fk constraints + self.persist_events.db_pool.simple_insert_many_txn( + txn, + table="events", + keys=( + "instance_name", + "stream_ordering", + "topological_ordering", + "depth", + "event_id", + "room_id", + "type", + "processed", + "outlier", + "origin_server_ts", + "received_ts", + "sender", + "contains_url", + "state_key", + "rejection_reason", + ), + values=[ + ( + "test", + event.internal_metadata.stream_ordering, + event.depth, # topological_ordering + event.depth, # depth + event.event_id, + event.room_id, + event.type, + True, # processed + False, # outlier + 1337, + 1741622420, + event.sender, + False, # contains url + event.state_key, + False, # rejected + ) + for event in events + ], + ) + for ev in events: + self.persist_events._store_state_dag_edges( + txn, + ev, + ) + + # satisfy fk constraints + self.get_success( + self.store.store_room(room_id, "foo", False, RoomVersions.MSC4242v12) + ) + self.get_success( + self.store.db_pool.runInteraction( + "_store_state_dag_edges", + insert, + ) + ) + # .- C -- D ---. + # A <- B E + # `- R -- W --` + # `-- T -` + TestCase = namedtuple("TestCase", "latest want limit") + test_cases = [ + TestCase(latest=["E"], want=["D", "T", "W"], limit=3), + TestCase(latest=["W", "T", "D"], want=["C", "R"], limit=2), + # breadth first and new entries are added to the end, sorted lexicographically + TestCase(latest=["E"], want=["D", "T", "W", "C", "R", "B", "A"], limit=100), + # we should sort the latest values initially + TestCase(latest=["E", "C"], want=["B", "D", "T", "W"], limit=4), + TestCase(latest=["C", "E"], want=["B", "D", "T", "W"], limit=4), + # dupes are ignored + TestCase( + latest=["E", "E", "C", "C", "C"], want=["B", "D", "T", "W"], limit=4 + ), + # include latest events in response. W included because reachable from E. + TestCase(latest=["W", "E"], want=["D", "T", "W", "R"], limit=4), + ] + for test_case in test_cases: + + def do_test(txn: LoggingTransaction) -> None: + got = self.store._get_missing_events_state_dag_txn( + txn, + room_id, + [], + test_case.latest, + test_case.limit, + ) + self.assertEquals(got, test_case.want) + + self.get_success( + self.store.db_pool.runInteraction("test_case", do_test), + ) + @attr.s(auto_attribs=True) class FakeEvent: @@ -1431,3 +1605,19 @@ class FakeEvent: def is_state(self) -> bool: return True + + +@attr.s(auto_attribs=True) +class StateDAGFakeEvent: + event_id: str + room_id: str + type: str + prev_state_events: list[str] + state_key = "foo" + depth = 1 + sender = "foo" + + internal_metadata = EventInternalMetadata({}) + + def is_state(self) -> bool: + return True