diff --git a/synapse/storage/events.py b/synapse/storage/events.py
index a37b7bad5a..d946024c9b 100644
--- a/synapse/storage/events.py
+++ b/synapse/storage/events.py
@@ -441,10 +441,10 @@ class EventsStore(SQLBaseStore):
         Assumes that we are only persisting events for one room at a time.
 
         Returns:
-            2-tuple (to_delete, to_insert) where both are state dicts, i.e.
-            (type, state_key) -> event_id. `to_delete` are the entries to
+            3-tuple (to_delete, to_insert, new_state) where both are state dicts,
+            i.e. (type, state_key) -> event_id. `to_delete` are the entries to
             first be deleted from current_state_events, `to_insert` are entries
-            to insert.
+            to insert. `new_state` is the full set of state.
             May return None if there are no changes to be applied.
         """
         # Now we need to work out the different state sets for
diff --git a/synapse/storage/state.py b/synapse/storage/state.py
index 1e1ce87e0e..5d6f7dfa28 100644
--- a/synapse/storage/state.py
+++ b/synapse/storage/state.py
@@ -227,6 +227,9 @@ class StateStore(SQLBaseStore):
                     ],
                 )
 
+            # Prefill the state group cache with this group.
+            # It's fine to use the sequence like this as the state group map
+            # is immutable.
             txn.call_after(
                 self._state_group_cache.update,
                 self._state_group_cache.sequence,