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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01LXMkcZVkcZhi21augBff4a

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::<PyDict>()?,
+                state_delta_due_to_event.cast::<PyMapping>()?,
             )?)
         };
 
         let mut state_group_deltas = HashMap::new();
         let deltas = context.getattr(intern!(py, "state_group_deltas"))?;
-        for (key, delta) in deltas.cast::<PyDict>()?.iter() {
+        for item in deltas.cast::<PyMapping>()?.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::<PyDict>()?)?);
+            state_group_deltas.insert(
+                (prev, current),
+                state_map_from_py(delta.cast::<PyMapping>()?)?,
+            );
         }
 
         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<StateMap> {
-    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<StateMap> {
+    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::<PyTuple>()?;
         let event_type: String = key.get_item(0)?.extract()?;
         let state_key: String = key.get_item(1)?.extract()?;

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)

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)