Store state groups when event contexts are created

4 files changed, 43 insertions, 16 deletions

diff --git a/synapse/handlers/federation.py b/synapse/handlers/federation.py
index ac70730885..aab59cb47f 100644
--- a/synapse/handlers/federation.py
+++ b/synapse/handlers/federation.py
@@ -1829,7 +1829,7 @@ class FederationHandler(BaseHandler):
                 different_auth = event_auth_events - current_state
 
                 self._update_context_for_auth_events(
-                    context, auth_events, event_key,
+                    event, context, auth_events, event_key,
                 )
 
         if different_auth and not event.internal_metadata.is_outlier():
@@ -1911,7 +1911,7 @@ class FederationHandler(BaseHandler):
                 # TODO.
 
                 self._update_context_for_auth_events(
-                    context, auth_events, event_key,
+                    event, context, auth_events, event_key,
                 )
 
         try:
@@ -1920,7 +1920,7 @@ class FederationHandler(BaseHandler):
             logger.warn("Failed auth resolution for %r because %s", event, e)
             raise e
 
-    def _update_context_for_auth_events(self, context, auth_events,
+    def _update_context_for_auth_events(self, event, context, auth_events,
                                         event_key):
         """Update the state_ids in an event context after auth event resolution
 
@@ -1947,7 +1947,13 @@ class FederationHandler(BaseHandler):
         context.prev_state_ids.update({
             k: a.event_id for k, a in auth_events.iteritems()
         })
-        context.state_group = self.store.get_next_state_group()
+
+        context.state_group = yield self.store.store_state_group(
+            event.event_id, event.room_id,
+            context.prev_group,
+            context.delta_ids,
+            context.current_state_ids,
+        )
 
     @defer.inlineCallbacks
     def construct_auth_difference(self, local_auth, remote_auth):
diff --git a/synapse/state.py b/synapse/state.py
index 9e624b4937..668bbe1f16 100644
--- a/synapse/state.py
+++ b/synapse/state.py
@@ -208,7 +208,12 @@ class StateHandler(object):
                 context.current_state_ids = {}
                 context.prev_state_ids = {}
             context.prev_state_events = []
-            context.state_group = self.store.get_next_state_group()
+            context.state_group = yield self.store.store_state_group(
+                event.event_id, event.room_id,
+                context.prev_group,
+                context.delta_ids,
+                context.current_state_ids,
+            )
             defer.returnValue(context)
 
         if old_state:
@@ -216,7 +221,6 @@ class StateHandler(object):
             context.prev_state_ids = {
                 (s.type, s.state_key): s.event_id for s in old_state
             }
-            context.state_group = self.store.get_next_state_group()
 
             if event.is_state():
                 key = (event.type, event.state_key)
@@ -230,6 +234,14 @@ class StateHandler(object):
                 context.current_state_ids = context.prev_state_ids
 
             context.prev_state_events = []
+
+            context.state_group = yield self.store.store_state_group(
+                event.event_id, event.room_id,
+                context.prev_group,
+                context.delta_ids,
+                context.current_state_ids,
+            )
+
             defer.returnValue(context)
 
         logger.debug("calling resolve_state_groups from compute_event_context")
@@ -242,8 +254,6 @@ class StateHandler(object):
         context = EventContext()
         context.prev_state_ids = curr_state
         if event.is_state():
-            context.state_group = self.store.get_next_state_group()
-
             key = (event.type, event.state_key)
             if key in context.prev_state_ids:
                 replaces = context.prev_state_ids[key]
@@ -261,16 +271,31 @@ class StateHandler(object):
                 context.prev_group = entry.prev_group
                 context.delta_ids = dict(entry.delta_ids)
                 context.delta_ids[key] = event.event_id
-        else:
-            if entry.state_group is None:
-                entry.state_group = self.store.get_next_state_group()
-                entry.state_id = entry.state_group
 
+            context.state_group = yield self.store.store_state_group(
+                event.event_id, event.room_id,
+                context.prev_group,
+                context.delta_ids,
+                context.current_state_ids,
+            )
+        else:
             context.state_group = entry.state_group
             context.current_state_ids = context.prev_state_ids
             context.prev_group = entry.prev_group
             context.delta_ids = entry.delta_ids
 
+            if entry.state_group is None:
+                entry.state_group = yield self.store.store_state_group(
+                    event.event_id, event.room_id,
+                    context.prev_group,
+                    context.delta_ids,
+                    context.current_state_ids,
+                )
+
+                entry.state_id = entry.state_group
+
+            context.state_group = entry.state_group
+
         context.prev_state_events = []
         defer.returnValue(context)
 
diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index d01d46338a..f8fbd02ceb 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -124,7 +124,6 @@ class DataStore(RoomMemberStore, RoomStore,
         )
 
         self._transaction_id_gen = IdGenerator(db_conn, "sent_transactions", "id")
-        self._state_groups_id_gen = IdGenerator(db_conn, "state_groups", "id")
         self._access_tokens_id_gen = IdGenerator(db_conn, "access_tokens", "id")
         self._event_reports_id_gen = IdGenerator(db_conn, "event_reports", "id")
         self._push_rule_id_gen = IdGenerator(db_conn, "push_rules", "id")
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index ead7134291..39fa6e724b 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -766,9 +766,6 @@ class StateStore(StateGroupReadStore, BackgroundUpdateStore):
 
             return count
 
-    def get_next_state_group(self):
-        return self._state_groups_id_gen.get_next()
-
     @defer.inlineCallbacks
     def _background_deduplicate_state(self, progress, batch_size):
         """This background update will slowly deduplicate state by reencoding