# # This file is licensed under the Affero General Public License (AGPL) version 3. # # Copyright (C) 2026 Element Creations Ltd # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU Affero General Public License as # published by the Free Software Foundation, either version 3 of the # License, or (at your option) any later version. # # See the GNU Affero General Public License for more details: # . # """The Rust state group writer against Python's, on the same batches.""" from typing import Any from twisted.internet import defer from twisted.internet.testing import MemoryReactor from synapse.api.room_versions import RoomVersions from synapse.events import EventBase, make_event_from_dict from synapse.events.snapshot import UnpersistedEventContext from synapse.logging.context import ( LoggingContext, current_context, make_deferred_yieldable, ) from synapse.server import HomeServer from synapse.types import StateMap from synapse.util.clock import Clock from tests import unittest from tests.utils import USE_POSTGRES_FOR_TESTS, USE_RUST_DRIVER_FOR_TESTS # The tables a chain of state groups writes to, and their columns. TABLES = { "state_groups": ("id", "room_id", "event_id"), "state_group_edges": ("state_group", "prev_state_group"), "state_groups_state": ("state_group", "room_id", "type", "state_key", "event_id"), } class RustStateGroupWriterTestCase(unittest.HomeserverTestCase): if not (USE_POSTGRES_FOR_TESTS and USE_RUST_DRIVER_FOR_TESTS): skip = "Needs Postgres with the native driver" def prepare( self, reactor: MemoryReactor, clock: Clock, homeserver: HomeServer ) -> None: self.state_store = homeserver.get_datastores().state self.rust_writer = self.state_store._rust_state_group_writer assert self.rust_writer is not None self.storage = homeserver.get_storage_controllers() def _batch( self, room_id: str, state_before: StateMap[str] | None, prev_group: int | None ) -> list[tuple[EventBase, UnpersistedEventContext]]: """Three events: two state events with a message between them.""" batch = [] for i, (event_type, state_key) in enumerate( [("m.room.topic", ""), ("m.room.message", None), ("m.room.member", "@a:t")] ): event_dict: dict[str, Any] = { "type": event_type, "sender": "@a:t", "room_id": room_id, "content": {"i": i}, "auth_events": [], "prev_events": [], "depth": 1 + i, "hashes": {"sha256": "x"}, "signatures": {}, "origin_server_ts": 1 + i, # The same in both twins, unlike a hashed event ID. "event_id": f"$e{i}:t", } if state_key is not None: event_dict["state_key"] = state_key context = UnpersistedEventContext( storage=self.storage, state_group_before_event=prev_group if i == 0 else None, state_group_after_event=None, state_delta_due_to_event=None, prev_group_for_state_group_before_event=None, delta_ids_to_state_group_before_event=None, partial_state=False, state_map_before_event=state_before if i == 0 else None, ) batch.append((make_event_from_dict(event_dict, RoomVersions.V1), context)) return batch def _store( self, batch: list[tuple[EventBase, UnpersistedEventContext]], room_id: str, prev_group: int | None, rust: bool, ) -> None: self.state_store._rust_state_group_writer = self.rust_writer if rust else None try: self.get_success( self.state_store.store_state_deltas_for_batched( batch, # type: ignore[arg-type] room_id, prev_group, ) ) finally: self.state_store._rust_state_group_writer = self.rust_writer def _rows(self, room_id: str) -> dict[str, list[tuple]]: """The rows of `room_id`'s state groups, with each group ID replaced by its rank among them, so two rooms can be compared.""" groups = self.get_success( self.state_store.db_pool.simple_select_onecol( "state_groups", {"room_id": room_id}, "id" ) ) rank = {group: i for i, group in enumerate(sorted(groups))} def norm(value: Any) -> Any: if isinstance(value, int): return rank.get(value, ("outside", value)) if isinstance(value, str): return value.replace(room_id, "!room") return value rows: dict[str, list[tuple]] = {} for table, columns in TABLES.items(): if table == "state_group_edges": found = self.get_success( self.state_store.db_pool.simple_select_many_batch( table, "state_group", groups, columns ) ) else: found = self.get_success( self.state_store.db_pool.simple_select_list( table, {"room_id": room_id}, columns ) ) rows[table] = sorted(tuple(norm(v) for v in row) for row in found) return rows def _contexts( self, batch: list[tuple[EventBase, UnpersistedEventContext]] ) -> list[tuple]: """Each context's groups, relative to the group before the batch, and its state delta.""" base = batch[0][1].state_group_before_event assert base is not None result = [] for _, context in batch: assert context.state_group_before_event is not None assert context.state_group_after_event is not None result.append( ( context.state_group_before_event - base, context.state_group_after_event - base, context.state_delta_due_to_event, ) ) return result def test_new_room(self) -> None: """A batch that starts a room stores the state before it as a full group, then the chain.""" state_before = {("m.room.create", ""): "$create", ("m.room.name", ""): "$name"} batches = {} for rust, room_id in ((False, "!python:t"), (True, "!rust:t")): batch = self._batch(room_id, state_before, None) self._store(batch, room_id, None, rust) batches[rust] = batch self.assertEqual(self._rows("!python:t"), self._rows("!rust:t")) self.assertEqual(len(self._rows("!rust:t")["state_groups"]), 3) # The contexts get the same groups, relative to the first. self.assertEqual(self._contexts(batches[False]), self._contexts(batches[True])) # The state after the batch reads back. group = batches[True][-1][1].state_group_after_event assert group is not None expected = dict(state_before) expected[("m.room.topic", "")] = batches[True][0][0].event_id expected[("m.room.member", "@a:t")] = batches[True][2][0].event_id self.assertEqual( self.get_success(self.storage.state.get_state_ids_for_group(group)), expected, ) def test_follows_existing_group(self) -> None: """A batch after an existing group chains on from it.""" groups = {} for rust, room_id in ((False, "!python:t"), (True, "!rust:t")): prev_group = self.get_success( self.state_store.store_state_group( "$prev", room_id, None, None, {("m.room.create", ""): "$create"} ) ) batch = self._batch(room_id, None, prev_group) self._store(batch, room_id, prev_group, rust) groups[rust] = (prev_group, batch) self.assertEqual(self._rows("!python:t"), self._rows("!rust:t")) for rust in (False, True): prev_group, batch = groups[rust] self.assertEqual(batch[0][1].state_group_before_event, prev_group) self.assertEqual( batch[1][1].state_group_after_event, batch[0][1].state_group_after_event ) def test_missing_prev_group(self) -> None: """A deleted `prev_group` fails the write, with Python's message.""" batch = self._batch("!rust:t", None, 999999) failure = self.get_failure( self.state_store.store_state_deltas_for_batched( batch, # type: ignore[arg-type] "!rust:t", 999999, ), Exception, ) self.assertIn( "Trying to persist state with unpersisted prev_group: 999999", str(failure.value), ) self.assertEqual(self._rows("!rust:t")["state_groups"], []) def test_bumps_pending_deletion(self) -> None: """Following a group that is pending deletion pushes its deletion back.""" prev_group = self.get_success( self.state_store.store_state_group( "$prev", "!rust:t", None, None, {("m.room.create", ""): "$create"} ) ) self.get_success( self.state_store.db_pool.simple_insert( "state_groups_pending_deletion", {"state_group": prev_group, "insertion_ts": 0}, ) ) before = self.get_success( self.state_store.db_pool.simple_select_one( "state_groups_pending_deletion", {"state_group": prev_group}, ("sequence_number", "insertion_ts"), ) ) self.reactor.advance(10) self._store( self._batch("!rust:t", None, prev_group), "!rust:t", prev_group, True ) sequence_number, insertion_ts = self.get_success( self.state_store.db_pool.simple_select_one( "state_groups_pending_deletion", {"state_group": prev_group}, ("sequence_number", "insertion_ts"), ) ) self.assertGreater(sequence_number, before[0]) self.assertEqual(insertion_ts, 10_000) def test_insert_failing_at_commit(self) -> None: """A row that fails to insert, which only shows at the COMMIT, writes nothing.""" prev_group = self.get_success( self.state_store.store_state_group( "$prev", "!rust:t", None, None, {("m.room.create", ""): "$create"} ) ) rows_before = self._rows("!rust:t") # Take the next state group ID from the sequence, give it to a row in # another room, and wind the sequence back so the writer gets it too. def clash(txn: Any) -> None: txn.execute("SELECT nextval('state_group_id_seq')") (next_id,) = txn.fetchone() txn.execute( "INSERT INTO state_groups (id, room_id, event_id) VALUES (?, ?, ?)", (next_id, "!other:t", "$other"), ) txn.execute("SELECT setval('state_group_id_seq', ?, false)", (next_id,)) self.get_success(self.state_store.db_pool.runInteraction("clash", clash)) batch = self._batch("!rust:t", None, prev_group) self.state_store._rust_state_group_writer = self.rust_writer failure = self.get_failure( self.state_store.store_state_deltas_for_batched( batch, # type: ignore[arg-type] "!rust:t", prev_group, ), Exception, ) self.assertIn("state_groups", str(failure.value)) self.assertEqual(self._rows("!rust:t"), rows_before) def test_keeps_logcontext(self) -> None: """Waiting for the Rust write leaves the caller in its own logcontext. The test homeserver makes Rust calls block, which would hide a mistake here, so this uses the Rust store directly. """ writer = self.rust_writer self.state_store._rust_state_group_writer = getattr(writer, "_handler", writer) batch = self._batch("!rust:t", {("m.room.create", ""): "$create"}, None) async def store() -> str: await self.state_store.store_state_deltas_for_batched( batch, # type: ignore[arg-type] "!rust:t", None, ) return str(current_context()) with LoggingContext(name="request", server_name="test"): context_after = self.get_success( make_deferred_yieldable(defer.ensureDeferred(store())) ) # The code after the write runs in the caller's context. self.assertEqual(context_after, "request")