Fix backfill github/erikj/chunk_pag_3 erikj/chunk_pag_3

2 files changed, 22 insertions, 30 deletions

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index 87c0615820..dc50c14039 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -714,37 +714,15 @@ class FederationHandler(BaseHandler):
         defer.returnValue(events)
 
     @defer.inlineCallbacks
-    def maybe_backfill(self, room_id, current_depth):
+    def maybe_backfill(self, room_id, extremities):
         """Checks the database to see if we should backfill before paginating,
         and if so do.
-        """
-        extremities = yield self.store.get_oldest_events_with_depth_in_room(
-            room_id
-        )
-
-        if not extremities:
-            logger.debug("Not backfilling as no extremeties found.")
-            return
-
-        # Check if we reached a point where we should start backfilling.
-        sorted_extremeties_tuple = sorted(
-            extremities.items(),
-            key=lambda e: -int(e[1])
-        )
-        max_depth = sorted_extremeties_tuple[0][1]
-
-        # We don't want to specify too many extremities as it causes the backfill
-        # request URI to be too long.
-        extremities = dict(sorted_extremeties_tuple[:5])
 
-        if current_depth > max_depth:
-            logger.debug(
-                "Not backfilling as we don't need to. %d < %d",
-                max_depth, current_depth,
-            )
-            return
-
-        # Now we need to decide which hosts to hit first.
+        Args:
+            room_id (str)
+            extremities (list[str]): List of event_ids to backfill from. These
+                should be event IDs that we don't yet have.
+        """
 
         # First we try hosts that are already in the room
         # TODO: HEURISTIC ALERT.
@@ -844,7 +822,7 @@ class FederationHandler(BaseHandler):
         tried_domains = set(likely_domains)
         tried_domains.add(self.server_name)
 
-        event_ids = list(extremities.iterkeys())
+        event_ids = list(extremities)
 
         logger.debug("calling resolve_state_groups in _maybe_backfill")
         resolve = logcontext.preserve_fn(
@@ -871,7 +849,7 @@ class FederationHandler(BaseHandler):
             } for key, state_dict in states.iteritems()
         }
 
-        for e_id, _ in sorted_extremeties_tuple:
+        for e_id in event_ids:
             likely_domains = get_domains_from_state(states[e_id])
 
             success = yield try_backfill([
diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index ca9c48cabc..9d295b536e 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -232,6 +232,20 @@ class MessageHandler(BaseHandler):
                 event_filter=event_filter,
             )
 
+            if source_config.direction == 'b' and extremities:
+                yield self.hs.get_handlers().federation_handler.maybe_backfill(
+                    room_id, extremities
+                )
+
+                events, next_key, extremities = yield self.store.paginate_room_events(
+                    room_id=room_id,
+                    from_key=source_config.from_key,
+                    to_key=source_config.to_key,
+                    direction=source_config.direction,
+                    limit=source_config.limit,
+                    event_filter=event_filter,
+                )
+
             next_token = pagin_config.from_token.copy_and_replace(
                 "room_key", next_key
             )