Limit size of get_auth_chain_ids query (#6947)

changelog.d/6947.misc

@@ -0,0 +1 @@
+Increase perf of `get_auth_chain_ids` used in state res v2.

synapse/storage/data_stores/main/event_federation.py

@@ -62,8 +62,20 @@ class EventFederationWorkerStore(EventsWorkerStore, SignatureWorkerStore, SQLBas
         )
 
     def _get_auth_chain_ids_txn(self, txn, event_ids, include_given):
+        if include_given:
+            results = set(event_ids)
+        else:
+            results = set()
+
         if isinstance(self.database_engine, PostgresEngine):
             # For efficiency we make the database do this if we can.
+
+            # We need to be a little careful with querying large amounts at
+            # once, for some reason postgres really doesn't like it. We do this
+            # by only asking for auth chain of 500 events at a time.
+            event_ids = list(event_ids)
+            chunks = [event_ids[x : x + 500] for x in range(0, len(event_ids), 500)]
+            for chunk in chunks:
                 sql = """
                     WITH RECURSIVE auth_chain(event_id) AS (
                         SELECT auth_id FROM event_auth WHERE event_id = ANY(?)
@@ -73,21 +85,14 @@ class EventFederationWorkerStore(EventsWorkerStore, SignatureWorkerStore, SQLBas
                     )
                     SELECT event_id FROM auth_chain
                 """
-            txn.execute(sql, (list(event_ids),))
+                txn.execute(sql, (chunk,))
 
-            results = set(event_id for event_id, in txn)
-
-            if include_given:
-                results.update(event_ids)
+                results.update(event_id for event_id, in txn)
 
             return list(results)
 
         # Database doesn't necessarily support recursive CTE, so we fall
         # back to do doing it manually.
-        if include_given:
-            results = set(event_ids)
-        else:
-            results = set()
 
         base_sql = "SELECT auth_id FROM event_auth WHERE "