From c4af646bdfc5ff340807f6609bbb8ff7cf44f526 Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Sat, 17 Oct 2026 17:46:41 +0000 Subject: [PATCH] Read event contexts that came from another worker in Rust A context sent to the event persister over replication holds its state maps as `immutabledict`s, which Rust's persist context conversion refused, as it took only a `dict`. So on a deployment with workers, preparing a chunk in Rust failed for every event another worker sent, and fell back to Python with an error logged. It now takes any mapping. Two worker tests with the native driver showed it, and then showed a test harness problem. `make_worker_hs` points a worker's Python pool at the master's database once the worker is built, but the worker's Rust handlers keep the native pool they were built with, on the worker's own database. It now points that pool at the master's database too. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LXMkcZVkcZhi21augBff4a --- rust/src/persist/context.rs | 23 ++++++++++++++-------- tests/replication/_base.py | 8 ++++++++ tests/storage/test_rust_persist_prereqs.py | 23 ++++++++++++++++++++++ 3 files changed, 46 insertions(+), 8 deletions(-) diff --git a/rust/src/persist/context.rs b/rust/src/persist/context.rs index 18d2755165..ccb5144d89 100644 --- a/rust/src/persist/context.rs +++ b/rust/src/persist/context.rs @@ -21,7 +21,7 @@ use std::{collections::HashMap, sync::Arc}; use pyo3::{ intern, prelude::*, - types::{PyDict, PyTuple}, + types::{PyMapping, PyTuple}, }; use crate::state::{StateKey, StateMap}; @@ -62,15 +62,19 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PersistContext { None } else { Some(state_map_from_py( - state_delta_due_to_event.cast::()?, + state_delta_due_to_event.cast::()?, )?) }; let mut state_group_deltas = HashMap::new(); let deltas = context.getattr(intern!(py, "state_group_deltas"))?; - for (key, delta) in deltas.cast::()?.iter() { + for item in deltas.cast::()?.items()?.iter() { + let (key, delta): (Bound<'_, PyAny>, Bound<'_, PyAny>) = item.extract()?; let (prev, current): (i64, i64) = key.extract()?; - state_group_deltas.insert((prev, current), state_map_from_py(delta.cast::()?)?); + state_group_deltas.insert( + (prev, current), + state_map_from_py(delta.cast::()?)?, + ); } Ok(PersistContext { @@ -97,10 +101,13 @@ impl PersistContext { } } -/// Convert a dict of `(type, state_key)` to event ID. -fn state_map_from_py(dict: &Bound<'_, PyDict>) -> PyResult { - let mut map = StateMap::with_capacity(dict.len()); - for (key, event_id) in dict.iter() { +/// Convert a mapping of `(type, state_key)` to event ID: a dict, or the +/// `immutabledict` of a context that came over replication. +fn state_map_from_py(mapping: &Bound<'_, PyMapping>) -> PyResult { + let items = mapping.items()?; + let mut map = StateMap::with_capacity(items.len()); + for item in items.iter() { + let (key, event_id): (Bound<'_, PyAny>, Bound<'_, PyAny>) = item.extract()?; let key = key.cast_into::()?; let event_type: String = key.get_item(0)?.extract()?; let state_key: String = key.get_item(1)?.extract()?; diff --git a/tests/replication/_base.py b/tests/replication/_base.py index 3bb0c9fa5d..43e7f2ef65 100644 --- a/tests/replication/_base.py +++ b/tests/replication/_base.py @@ -38,6 +38,7 @@ from synapse.replication.tcp.protocol import ( ) from synapse.replication.tcp.resource import ReplicationStreamProtocolFactory from synapse.server import HomeServer +from synapse.storage.rust_pool import RustConnectionPool from synapse.util.clock import Clock from synapse.util.duration import Duration @@ -382,7 +383,14 @@ class BaseMultiWorkerStreamTestCase(unittest.HomeserverTestCase): ) store = worker_hs.get_datastores().main + worker_pool = store.db_pool._db_pool store.db_pool._db_pool = self.database_pool._db_pool + # The worker's Rust handlers hold its own native pool. Point that at + # the master's database too, so Rust reads and writes what Python does. + if isinstance(worker_pool, RustConnectionPool) and isinstance( + self.database_pool._db_pool, RustConnectionPool + ): + worker_pool._pool = self.database_pool._db_pool._pool # Set up a resource for the worker resource = ReplicationRestResource(worker_hs) diff --git a/tests/storage/test_rust_persist_prereqs.py b/tests/storage/test_rust_persist_prereqs.py index bb3b9957f7..861efbb86e 100644 --- a/tests/storage/test_rust_persist_prereqs.py +++ b/tests/storage/test_rust_persist_prereqs.py @@ -25,6 +25,7 @@ from typing import Any, Collection from unittest import SkipTest from unittest.mock import Mock, patch +from immutabledict import immutabledict from prometheus_client import REGISTRY from twisted.internet.testing import MemoryReactor @@ -235,6 +236,28 @@ class PersistContextTestCase(unittest.HomeserverTestCase): ) self._assert_converts(context) + def test_context_from_replication(self) -> None: + """A context sent from another worker holds `immutabledict`s.""" + event, context = self.get_success( + create_event( + self.hs, + room_id=self.room_id, + type=EventTypes.Topic, + state_key="", + sender=self.user, + content={"topic": "a topic"}, + ) + ) + serialized = self.get_success( + context.serialize(event, self.hs.get_datastores().main) + ) + received = EventContext.deserialize( + self.hs.get_storage_controllers(), serialized + ) + self.assertIsInstance(received._state_delta_due_to_event, immutabledict) + self.assertTrue(received.state_group_deltas) + self._assert_converts(received) + def test_outlier_and_rejected(self) -> None: outlier = EventContext.for_outlier(self.hs.get_storage_controllers()) self._assert_converts(outlier)