synapse/tests/state/test_rust_handler.py

#
# 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:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
#
#

"""Tests for resolving state with the Rust state resolution
(`experimental_features.rust_state_res`). These need the native Rust Postgres
driver:

    SYNAPSE_TEST_RUST_DRIVER=1 POSTGRES=1 run-trial.sh tests.state.test_rust_handler
"""

from unittest import SkipTest
from unittest.mock import patch

from twisted.internet.testing import MemoryReactor

from synapse.api.room_versions import RoomVersions
from synapse.rest import admin
from synapse.rest.client import login, room
from synapse.server import HomeServer
from synapse.state import _rust_shadow_counter
from synapse.storage.rust_pool import RustConnectionPool
from synapse.types import create_requester
from synapse.util.clock import Clock

from tests import unittest


class RustStateResHandlerTestCase(unittest.HomeserverTestCase):
    servlets = [
        admin.register_servlets,
        room.register_servlets,
        login.register_servlets,
    ]

    def make_homeserver(self, reactor: MemoryReactor, clock: Clock) -> HomeServer:
        hs = self.setup_test_homeserver()
        if not isinstance(
            hs.get_datastores().main.db_pool._db_pool, RustConnectionPool
        ):
            raise SkipTest("needs the native Rust Postgres driver")
        return hs

    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
        self.store = hs.get_datastores().main
        self.user_id = self.register_user("user", "pass")
        self.tok = self.login("user", "pass")
        self.handler = hs.get_state_resolution_handler()

    def _fork(self, room_id: str) -> tuple[str, str]:
        """Send two topics with the same prev events, so persisting the second
        resolves state across the fork. Returns their event IDs."""
        event_handler = self.hs.get_event_creation_handler()
        requester = create_requester(self.user_id)
        prev_event_ids = self.get_success(self.store.get_prev_events_for_room(room_id))

        event_ids = []
        for topic in ("A", "B"):
            event, unpersisted_context = self.get_success(
                event_handler.create_event(
                    requester,
                    {
                        "type": "m.room.topic",
                        "state_key": "",
                        "content": {"topic": topic},
                        "room_id": room_id,
                        "sender": self.user_id,
                    },
                    prev_event_ids=prev_event_ids,
                )
            )
            context = self.get_success(unpersisted_context.persist(event))
            self.get_success(
                event_handler.handle_new_client_event(
                    requester, events_and_context=[(event, context)]
                )
            )
            event_ids.append(event.event_id)
        return event_ids[0], event_ids[1]

    def _shadow_count(self, result: str) -> float:
        return _rust_shadow_counter.labels(
            result=result, server_name=self.hs.hostname
        )._value.get()

    def _current_topic(self, room_id: str) -> str | None:
        state = self.get_success(self.store.get_partial_current_state_ids(room_id))
        return state.get(("m.room.topic", ""))

    def test_shadow_mode_compares(self) -> None:
        """In shadow mode, resolutions run in both and agree."""
        self.handler._rust_state_res_mode = "shadow"
        self.handler._rust_state_res = self.hs.get_rust_handlers().state_resolution
        assert self.handler._rust_state_res is not None

        matches = self._shadow_count("match")
        mismatches = self._shadow_count("mismatch")
        errors = self._shadow_count("error")

        room_id = self.helper.create_room_as(self.user_id, tok=self.tok)
        topics = self._fork(room_id)

        self.assertGreater(self._shadow_count("match"), matches)
        self.assertEqual(self._shadow_count("mismatch"), mismatches)
        self.assertEqual(self._shadow_count("error"), errors)
        self.assertIn(self._current_topic(room_id), topics)

    def test_on_mode_uses_rust(self) -> None:
        """In "on" mode, v2 and v2.1 resolutions run in Rust and the metrics
        come from its stats."""
        self.handler._rust_state_res_mode = "on"
        self.handler._rust_state_res = self.hs.get_rust_handlers().state_resolution
        assert self.handler._rust_state_res is not None

        for room_version in (RoomVersions.V10, RoomVersions.V12):
            room_id = self.helper.create_room_as(
                self.user_id, room_version=room_version.identifier, tok=self.tok
            )

            with patch(
                "synapse.state.v2.resolve_events_with_store",
                side_effect=AssertionError("Python state res should not run"),
            ):
                topics = self._fork(room_id)

            self.assertIn(self._current_topic(room_id), topics)
            self.assertIn(room_id, self.handler._state_res_metrics)

    def test_on_mode_falls_back_without_chain_cover_index(self) -> None:
        """A room without a chain cover index is resolved in Python."""
        self.handler._rust_state_res_mode = "on"
        self.handler._rust_state_res = self.hs.get_rust_handlers().state_resolution

        room_id = self.helper.create_room_as(
            self.user_id, room_version=RoomVersions.V10.identifier, tok=self.tok
        )
        self.get_success(
            self.store.db_pool.simple_update(
                table="rooms",
                keyvalues={"room_id": room_id},
                updatevalues={"has_auth_chain_index": False},
                desc="test",
            )
        )

        with self.assertLogs("synapse.state", level="INFO") as logs:
            topics = self._fork(room_id)
        self.assertTrue(
            any("Falling back to Python state res" in line for line in logs.output)
        )
        self.assertIn(self._current_topic(room_id), topics)