1 files changed, 31 insertions, 0 deletions
diff --git a/synapse/storage/util/id_generators.py b/synapse/storage/util/id_generators.py
index 89d1643f10..83eab63098 100644
--- a/synapse/storage/util/id_generators.py
+++ b/synapse/storage/util/id_generators.py
@@ -108,6 +108,37 @@ class StreamIdGenerator(object):
         defer.returnValue(manager())
 
     @defer.inlineCallbacks
+    def get_next_mult(self, store, n):
+        """
+        Usage:
+            with yield stream_id_gen.get_next(store, n) as stream_ids:
+                # ... persist events ...
+        """
+        if not self._current_max:
+            yield store.runInteraction(
+                "_compute_current_max",
+                self._get_or_compute_current_max,
+            )
+
+        with self._lock:
+            next_ids = range(self._current_max + 1, self._current_max + n + 1)
+            self._current_max += n
+
+            for next_id in next_ids:
+                self._unfinished_ids.append(next_id)
+
+        @contextlib.contextmanager
+        def manager():
+            try:
+                yield next_ids
+            finally:
+                with self._lock:
+                    for next_id in next_ids:
+                        self._unfinished_ids.remove(next_id)
+
+        defer.returnValue(manager())
+
+    @defer.inlineCallbacks
     def get_max_token(self, store):
         """Returns the maximum stream id such that all stream ids less than or
         equal to it have been successfully persisted.