1 files changed, 22 insertions, 22 deletions
diff --git a/synapse/storage/util/id_generators.py b/synapse/storage/util/id_generators.py
index 9d461d5e96..e40eb8a8c4 100644
--- a/synapse/storage/util/id_generators.py
+++ b/synapse/storage/util/id_generators.py
@@ -30,15 +30,13 @@ class IdGenerator(object):
 
     @defer.inlineCallbacks
     def get_next(self):
-        with self._lock:
-            if not self._next_id:
-                res = yield self.store._execute_and_decode(
-                    "IdGenerator_%s" % (self.table,),
-                    "SELECT MAX(%s) as mx FROM %s" % (self.column, self.table,)
-                )
-
-                self._next_id = (res and res[0] and res[0]["mx"]) or 1
+        if self._next_id is None:
+            yield self.store.runInteraction(
+                "IdGenerator_%s" % (self.table,),
+                self.get_next_txn,
+            )
 
+        with self._lock:
             i = self._next_id
             self._next_id += 1
             defer.returnValue(i)
@@ -86,10 +84,10 @@ class StreamIdGenerator(object):
             with stream_id_gen.get_next_txn(txn) as stream_id:
                 # ... persist event ...
         """
-        with self._lock:
-            if not self._current_max:
-                self._compute_current_max(txn)
+        if not self._current_max:
+            self._get_or_compute_current_max(txn)
 
+        with self._lock:
             self._current_max += 1
             next_id = self._current_max
 
@@ -110,22 +108,24 @@ class StreamIdGenerator(object):
         """Returns the maximum stream id such that all stream ids less than or
         equal to it have been successfully persisted.
         """
+        if not self._current_max:
+            yield store.runInteraction(
+                "_compute_current_max",
+                self._get_or_compute_current_max,
+            )
+
         with self._lock:
             if self._unfinished_ids:
                 defer.returnValue(self._unfinished_ids[0] - 1)
 
-            if not self._current_max:
-                yield store.runInteraction(
-                    "_compute_current_max",
-                    self._compute_current_max,
-                )
-
             defer.returnValue(self._current_max)
 
-    def _compute_current_max(self, txn):
-        txn.execute("SELECT MAX(stream_ordering) FROM events")
-        val, = txn.fetchone()
+    def _get_or_compute_current_max(self, txn):
+        with self._lock:
+            txn.execute("SELECT MAX(stream_ordering) FROM events")
+            rows = txn.fetchall()
+            val, = rows[0]
 
-        self._current_max = int(val) if val else 1
+            self._current_max = int(val) if val else 1
 
-        return self._current_max
+            return self._current_max