Merge pull request #1178 from matrix-org/erikj/current_room_token

Fix incredibly slow back pagination query

4 files changed, 48 insertions, 12 deletions

diff --git a/synapse/handlers/message.py b/synapse/handlers/message.py
index 59eb26beaf..abfa8c65a4 100644
--- a/synapse/handlers/message.py
+++ b/synapse/handlers/message.py
@@ -82,8 +82,8 @@ class MessageHandler(BaseHandler):
             room_token = pagin_config.from_token.room_key
         else:
             pagin_config.from_token = (
-                yield self.hs.get_event_sources().get_current_token(
-                    direction='b'
+                yield self.hs.get_event_sources().get_current_token_for_room(
+                    room_id=room_id
                 )
             )
             room_token = pagin_config.from_token.room_key
diff --git a/synapse/handlers/room.py b/synapse/handlers/room.py
index a7f533f7be..59e4d1cd15 100644
--- a/synapse/handlers/room.py
+++ b/synapse/handlers/room.py
@@ -475,8 +475,11 @@ class RoomEventSource(object):
 
         defer.returnValue((events, end_key))
 
-    def get_current_key(self, direction='f'):
-        return self.store.get_room_events_max_id(direction)
+    def get_current_key(self):
+        return self.store.get_room_events_max_id()
+
+    def get_current_key_for_room(self, room_id):
+        return self.store.get_room_events_max_id(room_id)
 
     @defer.inlineCallbacks
     def get_pagination_rows(self, user, config, key):
diff --git a/synapse/storage/stream.py b/synapse/storage/stream.py
index 07ea969d4d..888b1cb35d 100644
--- a/synapse/storage/stream.py
+++ b/synapse/storage/stream.py
@@ -521,13 +521,20 @@ class StreamStore(SQLBaseStore):
         )
 
     @defer.inlineCallbacks
-    def get_room_events_max_id(self, direction='f'):
+    def get_room_events_max_id(self, room_id=None):
+        """Returns the current token for rooms stream.
+
+        By default, it returns the current global stream token. Specifying a
+        `room_id` causes it to return the current room specific topological
+        token.
+        """
         token = yield self._stream_id_gen.get_current_token()
-        if direction != 'b':
+        if room_id is None:
             defer.returnValue("s%d" % (token,))
         else:
             topo = yield self.runInteraction(
-                "_get_max_topological_txn", self._get_max_topological_txn
+                "_get_max_topological_txn", self._get_max_topological_txn,
+                room_id,
             )
             defer.returnValue("t%d-%d" % (topo, token))
 
@@ -579,11 +586,11 @@ class StreamStore(SQLBaseStore):
             lambda r: r[0][0] if r else 0
         )
 
-    def _get_max_topological_txn(self, txn):
+    def _get_max_topological_txn(self, txn, room_id):
         txn.execute(
             "SELECT MAX(topological_ordering) FROM events"
-            " WHERE outlier = ?",
-            (False,)
+            " WHERE room_id = ?",
+            (room_id,)
         )
 
         rows = txn.fetchall()
diff --git a/synapse/streams/events.py b/synapse/streams/events.py
index 6bf21d6f5e..4d44c3d4ca 100644
--- a/synapse/streams/events.py
+++ b/synapse/streams/events.py
@@ -41,13 +41,39 @@ class EventSources(object):
         self.store = hs.get_datastore()
 
     @defer.inlineCallbacks
-    def get_current_token(self, direction='f'):
+    def get_current_token(self):
         push_rules_key, _ = self.store.get_push_rules_stream_token()
         to_device_key = self.store.get_to_device_stream_token()
 
         token = StreamToken(
             room_key=(
-                yield self.sources["room"].get_current_key(direction)
+                yield self.sources["room"].get_current_key()
+            ),
+            presence_key=(
+                yield self.sources["presence"].get_current_key()
+            ),
+            typing_key=(
+                yield self.sources["typing"].get_current_key()
+            ),
+            receipt_key=(
+                yield self.sources["receipt"].get_current_key()
+            ),
+            account_data_key=(
+                yield self.sources["account_data"].get_current_key()
+            ),
+            push_rules_key=push_rules_key,
+            to_device_key=to_device_key,
+        )
+        defer.returnValue(token)
+
+    @defer.inlineCallbacks
+    def get_current_token_for_room(self, room_id):
+        push_rules_key, _ = self.store.get_push_rules_stream_token()
+        to_device_key = self.store.get_to_device_stream_token()
+
+        token = StreamToken(
+            room_key=(
+                yield self.sources["room"].get_current_key_for_room(room_id)
             ),
             presence_key=(
                 yield self.sources["presence"].get_current_key()