From a368d30c1cfe7457fca4fcdd03ae481ba65a226c Mon Sep 17 00:00:00 2001
From: Shay <hillerys@element.io>
Date: Tue, 7 Mar 2023 13:54:39 -0800
Subject: [PATCH] More speedups/fixes to creating batched events (#15195)

---
 changelog.d/15195.misc         |  1 +
 synapse/event_auth.py          | 23 +++++++++++++++++------
 synapse/events/snapshot.py     |  1 +
 synapse/handlers/event_auth.py | 13 +++++++++++--
 synapse/handlers/room.py       |  4 +++-
 5 files changed, 33 insertions(+), 9 deletions(-)
 create mode 100644 changelog.d/15195.misc

diff --git a/changelog.d/15195.misc b/changelog.d/15195.misc
new file mode 100644
index 000000000..d8beea917
--- /dev/null
+++ b/changelog.d/15195.misc
@@ -0,0 +1 @@
+Improve performance of creating and authenticating events.
\ No newline at end of file
diff --git a/synapse/event_auth.py b/synapse/event_auth.py
index 4d6d1b8eb..af55874b5 100644
--- a/synapse/event_auth.py
+++ b/synapse/event_auth.py
@@ -168,13 +168,24 @@ async def check_state_independent_auth_rules(
         return
 
     # 2. Reject if event has auth_events that: ...
-    auth_events = await store.get_events(
-        event.auth_event_ids(),
-        redact_behaviour=EventRedactBehaviour.as_is,
-        allow_rejected=True,
-    )
     if batched_auth_events:
-        auth_events.update(batched_auth_events)
+        # Copy the batched auth events to avoid mutating them.
+        auth_events = dict(batched_auth_events)
+        needed_auth_event_ids = set(event.auth_event_ids()) - batched_auth_events.keys()
+        if needed_auth_event_ids:
+            auth_events.update(
+                await store.get_events(
+                    needed_auth_event_ids,
+                    redact_behaviour=EventRedactBehaviour.as_is,
+                    allow_rejected=True,
+                )
+            )
+    else:
+        auth_events = await store.get_events(
+            event.auth_event_ids(),
+            redact_behaviour=EventRedactBehaviour.as_is,
+            allow_rejected=True,
+        )
 
     room_id = event.room_id
     auth_dict: MutableStateMap[str] = {}
diff --git a/synapse/events/snapshot.py b/synapse/events/snapshot.py
index a91a5d1e3..c04ad08cb 100644
--- a/synapse/events/snapshot.py
+++ b/synapse/events/snapshot.py
@@ -293,6 +293,7 @@ class EventContext(UnpersistedEventContextBase):
             Maps a (type, state_key) to the event ID of the state event matching
             this tuple.
         """
+
         assert self.state_group_before_event is not None
         return await self._storage.state.get_state_ids_for_group(
             self.state_group_before_event, state_filter
diff --git a/synapse/handlers/event_auth.py b/synapse/handlers/event_auth.py
index c508861b6..0db0bd730 100644
--- a/synapse/handlers/event_auth.py
+++ b/synapse/handlers/event_auth.py
@@ -63,9 +63,18 @@ class EventAuthHandler:
             self._store, event, batched_auth_events
         )
         auth_event_ids = event.auth_event_ids()
-        auth_events_by_id = await self._store.get_events(auth_event_ids)
+
         if batched_auth_events:
-            auth_events_by_id.update(batched_auth_events)
+            # Copy the batched auth events to avoid mutating them.
+            auth_events_by_id = dict(batched_auth_events)
+            needed_auth_event_ids = set(auth_event_ids) - set(batched_auth_events)
+            if needed_auth_event_ids:
+                auth_events_by_id.update(
+                    await self._store.get_events(needed_auth_event_ids)
+                )
+        else:
+            auth_events_by_id = await self._store.get_events(auth_event_ids)
+
         check_state_dependent_auth_rules(event, auth_events_by_id.values())
 
     def compute_auth_events(
diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index b1784638f..32451670f 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -1123,7 +1123,9 @@ class RoomCreationHandler:
                 event_dict,
                 prev_event_ids=prev_event,
                 depth=depth,
-                state_map=state_map,
+                # Take a copy to ensure each event gets a unique copy of
+                # state_map since it is modified below.
+                state_map=dict(state_map),
                 for_batch=for_batch,
             )