MORE HACKS

synapse/replication/tcp/streams/partial_state.py

@@ -71,6 +71,6 @@ def __init__(self, hs: "HomeServer"):
         super().__init__(
             hs.get_instance_name(),
             # TODO(faster_joins, multiple writers): we need to account for instance names
-            current_token_without_instance(store.get_un_partial_stated_events_token),
+            store.get_un_partial_stated_events_token,
             store.get_un_partial_stated_events_from_stream,
         )

synapse/storage/databases/main/events_worker.py

@@ -314,11 +314,12 @@ def get_chain_id_txn(txn: Cursor) -> int:
                 db_conn, "un_partial_stated_event_stream", "stream_id"
             )
 
-    def get_un_partial_stated_events_token(self) -> int:
-        # TODO(faster_joins, multiple writers): This is inappropriate if there are multiple
-        #     writers because workers that don't write often will hold all
-        #     readers up.
-        return self._un_partial_stated_events_stream_id_gen.get_current_token()
+    def get_un_partial_stated_events_token(self, instance_name: str) -> int:
+        return (
+            self._un_partial_stated_events_stream_id_gen.get_current_token_for_writer(
+                instance_name
+            )
+        )
 
     async def get_un_partial_stated_events_from_stream(
         self, instance_name: str, last_id: int, current_id: int, limit: int
@@ -408,6 +409,11 @@ def process_replication_position(
             self._stream_id_gen.advance(instance_name, token)
         elif stream_name == BackfillStream.NAME:
             self._backfill_id_gen.advance(instance_name, -token)
+        elif stream_name == UnPartialStatedEventStream.NAME:
+            logger.info(
+                "Advancing %s token to %s", UnPartialStatedEventStream.NAME, token
+            )
+            self._un_partial_stated_events_stream_id_gen.advance(instance_name, token)
         super().process_replication_position(stream_name, instance_name, token)
 
     async def have_censored_event(self, event_id: str) -> bool: