Initial stab at implementing a batched get_missing_pdus request

1 files changed, 58 insertions, 5 deletions

diff --git a/synapse/storage/event_federation.py b/synapse/storage/event_federation.py

index 3fbc090224..22bf7ad832 100644
--- a/synapse/storage/event_federation.py
+++ b/synapse/storage/event_federation.py
@@ -32,15 +32,15 @@ class EventFederationStore(SQLBaseStore):
     and backfilling from another server respectively.
     """
 
-    def get_auth_chain(self, event_ids):
+    def get_auth_chain(self, event_ids, have_ids=set()):
         return self.runInteraction(
             "get_auth_chain",
             self._get_auth_chain_txn,
-            event_ids
+            event_ids, have_ids
         )
 
-    def _get_auth_chain_txn(self, txn, event_ids):
-        results = self._get_auth_chain_ids_txn(txn, event_ids)
+    def _get_auth_chain_txn(self, txn, event_ids, have_ids):
+        results = self._get_auth_chain_ids_txn(txn, event_ids, have_ids)
 
         return self._get_events_txn(txn, results)
 
@@ -51,8 +51,9 @@ class EventFederationStore(SQLBaseStore):
             event_ids
         )
 
-    def _get_auth_chain_ids_txn(self, txn, event_ids):
+    def _get_auth_chain_ids_txn(self, txn, event_ids, have_ids):
         results = set()
+        have_ids = set(have_ids)
 
         base_sql = (
             "SELECT auth_id FROM event_auth WHERE event_id = ?"
@@ -64,6 +65,10 @@ class EventFederationStore(SQLBaseStore):
             for f in front:
                 txn.execute(base_sql, (f,))
                 new_front.update([r[0] for r in txn.fetchall()])
+
+            new_front -= results
+            new_front -= have_ids
+
             front = new_front
             results.update(front)
 
@@ -378,3 +383,51 @@ class EventFederationStore(SQLBaseStore):
             event_results += new_front
 
         return self._get_events_txn(txn, event_results)
+
+    def get_missing_events(self, room_id, earliest_events, latest_events,
+                           limit, min_depth):
+        return self.runInteraction(
+            "get_missing_events",
+            self._get_missing_events,
+            room_id, earliest_events, latest_events, limit, min_depth
+        )
+
+    def _get_missing_events(self, txn, room_id, earliest_events, latest_events,
+                            limit, min_depth):
+
+        earliest_events = set(earliest_events)
+        front = set(latest_events) - earliest_events
+
+        event_results = set()
+
+        query = (
+            "SELECT prev_event_id FROM event_edges "
+            "WHERE room_id = ? AND event_id = ? AND is_state = 0 "
+            "LIMIT ?"
+        )
+
+        while front and len(event_results) < limit:
+            new_front = set()
+            for event_id in front:
+                txn.execute(
+                    query,
+                    (room_id, event_id, limit - len(event_results))
+                )
+
+                for e_id, in txn.fetchall():
+                    new_front.add(e_id)
+
+            new_front -= earliest_events
+            new_front -= event_results
+
+            front = new_front
+            event_results |= new_front
+
+        events = self._get_events_txn(txn, event_results)
+
+        events = sorted(
+            [ev for ev in events if ev.depth >= min_depth],
+            key=lambda e: e.depth,
+        )
+
+        return events[:limit]