Load events for state group seperately

2 files changed, 11 insertions, 5 deletions

diff --git a/synapse/storage/_base.py b/synapse/storage/_base.py
index c9fe5a3555..0279400a83 100644
--- a/synapse/storage/_base.py
+++ b/synapse/storage/_base.py
@@ -867,9 +867,9 @@ class SQLBaseStore(object):
         return self.runInteraction("_simple_max_id", func)
 
     def _get_events(self, event_ids, check_redacted=True,
-                    get_prev_content=False):
+                    get_prev_content=False, desc="_get_events"):
         return self.runInteraction(
-            "_get_events", self._get_events_txn, event_ids,
+            desc, self._get_events_txn, event_ids,
             check_redacted=check_redacted, get_prev_content=get_prev_content,
         )
 
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index dbc0e49c1f..9ed5412999 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -43,6 +43,7 @@ class StateStore(SQLBaseStore):
       * `state_groups_state`: Maps state group to state events.
     """
 
+    @defer.inlineCallbacks
     def get_state_groups(self, event_ids):
         """ Get the state groups for the given list of event_ids
 
@@ -71,17 +72,22 @@ class StateStore(SQLBaseStore):
                     retcol="event_id",
                 )
 
-                state = self._get_events_txn(txn, state_ids)
+                # state = self._get_events_txn(txn, state_ids)
 
-                res[group] = state
+                res[group] = state_ids
 
             return res
 
-        return self.runInteraction(
+        states = yield self.runInteraction(
             "get_state_groups",
             f,
         )
 
+        for vals in states.values():
+            vals[:] = yield self._get_events(vals, desc="_get_state_groups_ev")
+
+        defer.returnValue(states)
+
     def _store_state_groups_txn(self, txn, event, context):
         if context.current_state is None:
             return