From a23a3136f92e18ca02ece53c3e6fcd613b23d55a Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 30 Sep 2020 11:16:33 +0100
Subject: [PATCH 01/13] Add an immutable mapping from instance name to ID

---
 .../delta/58/19instance_map.sql.postgres      | 25 +++++++++
 synapse/storage/databases/main/stream.py      | 53 +++++++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres

diff --git a/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres b/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
new file mode 100644
index 000000000000..bdbbab573835
--- /dev/null
+++ b/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
@@ -0,0 +1,25 @@
+/* Copyright 2020 The Matrix.org Foundation C.I.C
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+
+-- A unique and immutable mapping between instance name and an integer ID. This
+-- let's us refer to instances via a small ID in e.g. stream tokens, without
+-- having to encode the full name.
+CREATE TABLE instance_map (
+    instance_id SERIAL PRIMARY KEY,
+    instance_name TEXT NOT NULL
+);
+
+CREATE UNIQUE INDEX instance_map_idx ON instance_map(instance_name);
diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 37249f1e3f99..e6449d68b854 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -54,6 +54,7 @@
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine
 from synapse.types import Collection, PersistedEventPosition, RoomStreamToken
+from synapse.util.caches.descriptors import cached
 from synapse.util.caches.stream_change_cache import StreamChangeCache
 
 if TYPE_CHECKING:
@@ -1095,6 +1096,58 @@ async def paginate_room_events(
 
         return (events, token)
 
+    @cached()
+    async def get_id_for_instance(self, instance_name: str) -> int:
+        """Get a unique, immutable ID that corresponds to the instance.
+        """
+
+        def _get_id_for_instance_txn(txn):
+            instance_id = self.db_pool.simple_select_one_onecol_txn(
+                txn,
+                table="instance_map",
+                keyvalues={"instance_name": instance_name},
+                retcol="instance_id",
+                allow_none=True,
+            )
+            if instance_id is not None:
+                return instance_id
+
+            # If we don't have an entry upsert one.
+            #
+            # We could do this before the first check, and rely on the cache for
+            # efficiency, but each UPSERT causes the next ID to increment which
+            # can quickly bloat the size of the generated IDs for new instances.
+            self.db_pool.simple_upsert_txn(
+                txn,
+                table="instance_map",
+                keyvalues={"instance_name": instance_name},
+                values={},
+            )
+
+            return self.db_pool.simple_select_one_onecol_txn(
+                txn,
+                table="instance_map",
+                keyvalues={"instance_name": instance_name},
+                retcol="instance_id",
+            )
+
+        return await self.db_pool.runInteraction(
+            "get_id_for_instance", _get_id_for_instance_txn
+        )
+
+    @cached()
+    async def get_name_from_instance_id(self, instance_id: int) -> str:
+        """Get the instance name from an ID previously returned by
+        `get_id_for_instance`.
+        """
+
+        return await self.db_pool.simple_select_one_onecol(
+            table="instance_map",
+            keyvalues={"instance_id": instance_id},
+            retcol="instance_name",
+            desc="get_name_from_instance_id",
+        )
+
 
 class StreamStore(StreamWorkerStore):
     def get_room_max_stream_ordering(self) -> int:

From b2172da9bba51628bced7979a516acb58b79f241 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 1 Oct 2020 16:14:00 +0100
Subject: [PATCH 02/13] Create new token format for encoding instance positions

---
 synapse/types.py | 65 ++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 60 insertions(+), 5 deletions(-)

diff --git a/synapse/types.py b/synapse/types.py
index bd271f9f1611..550a968a03be 100644
--- a/synapse/types.py
+++ b/synapse/types.py
@@ -22,6 +22,7 @@
     TYPE_CHECKING,
     Any,
     Dict,
+    Iterable,
     Mapping,
     MutableMapping,
     Optional,
@@ -43,7 +44,7 @@
 if sys.version_info[:3] >= (3, 6, 0):
     from typing import Collection
 else:
-    from typing import Container, Iterable, Sized
+    from typing import Container, Sized
 
     T_co = TypeVar("T_co", covariant=True)
 
@@ -375,7 +376,7 @@ def f2(m):
     return username.decode("ascii")
 
 
-@attr.s(frozen=True, slots=True)
+@attr.s(frozen=True, slots=True, cmp=False)
 class RoomStreamToken:
     """Tokens are positions between events. The token "s1" comes after event 1.
 
@@ -395,8 +396,30 @@ class RoomStreamToken:
 
     Live tokens start with an "s" followed by the "stream_ordering" id of the
     event it comes after. Historic tokens start with a "t" followed by the
-    "topological_ordering" id of the event it comes after, followed by "-",
+    "topological_ordering" id of the event it comes after, followed by "~",
     followed by the "stream_ordering" id of the event it comes after.
+
+    There is also a third mode for live tokens where the token starts with "m",
+    which is sometimes used when using sharded event persisters. In this case
+    the events stream is considered to be a set of streams (one for each writer)
+    and the token encodes the vector clock of positions of each writer in their
+    respective streams.
+
+    The format of the token in such case is an initial integer min position,
+    followed by the mapping of instance ID to position seperate by '.' and '~':
+
+        m{min_pos}.{writer1}~{pos1}.{writer2}~{pos2}. ...
+
+    The `min_pos` corresponds to the minimum position all writers have persisted
+    up to, and then only writers that are ahead of that position need to be
+    encoded. An example token is:
+
+        m56.2~58.3~59
+
+    Which corresponds to a set of three (or more writers) where instances 2 and
+    3 (these are instance IDs that can be looked up in the DB to fetch the more
+    commonly used instance names) are at positions 58 and 59 respectively, and
+    all other instances are at position 56.
     """
 
     topological = attr.ib(
@@ -405,6 +428,8 @@ class RoomStreamToken:
     )
     stream = attr.ib(type=int, validator=attr.validators.instance_of(int))
 
+    instance_map = attr.ib(type=Dict[str, int], factory=dict)
+
     @classmethod
     async def parse(cls, store: "DataStore", string: str) -> "RoomStreamToken":
         try:
@@ -413,6 +438,20 @@ async def parse(cls, store: "DataStore", string: str) -> "RoomStreamToken":
             if string[0] == "t":
                 parts = string[1:].split("-", 1)
                 return cls(topological=int(parts[0]), stream=int(parts[1]))
+            if string[0] == "m":
+                parts = string[1:].split(".")
+                stream = int(parts[0])
+
+                instance_map = {}
+                for part in parts[1:]:
+                    key, value = part.split("~")
+                    instance_id = int(key)
+                    pos = int(value)
+
+                    instance_name = await store.get_name_from_instance_id(instance_id)
+                    instance_map[instance_name] = pos
+
+                return cls(topological=None, stream=stream, instance_map=instance_map,)
         except Exception:
             pass
         raise SynapseError(400, "Invalid token %r" % (string,))
@@ -436,7 +475,15 @@ def copy_and_advance(self, other: "RoomStreamToken") -> "RoomStreamToken":
 
         max_stream = max(self.stream, other.stream)
 
-        return RoomStreamToken(None, max_stream)
+        instance_map = {
+            instance: max(
+                self.instance_map.get(instance, self.stream),
+                other.instance_map.get(instance, other.stream),
+            )
+            for instance in set(self.instance_map).union(other.instance_map)
+        }
+
+        return RoomStreamToken(None, max_stream, instance_map)
 
     def as_tuple(self) -> Tuple[Optional[int], int]:
         return (self.topological, self.stream)
@@ -444,6 +491,14 @@ def as_tuple(self) -> Tuple[Optional[int], int]:
     async def to_string(self, store: "DataStore") -> str:
         if self.topological is not None:
             return "t%d-%d" % (self.topological, self.stream)
+        elif self.instance_map:
+            entries = []
+            for name, pos in self.instance_map.items():
+                instance_id = await store.get_id_for_instance(name)
+                entries.append("{}~{}".format(instance_id, pos))
+
+            encoded_map = ".".join(entries)
+            return "m{}.{}".format(self.stream, encoded_map)
         else:
             return "s%d" % (self.stream,)
 
@@ -535,7 +590,7 @@ class PersistedEventPosition:
     stream = attr.ib(type=int)
 
     def persisted_after(self, token: RoomStreamToken) -> bool:
-        return token.stream < self.stream
+        return token.get_stream_pos_for_instance(self.instance_name) < self.stream
 
     def to_room_stream_token(self) -> RoomStreamToken:
         """Converts the position to a room stream token such that events

From d7da8ca8a8202145b39676cbcb9a5b90b0f3c874 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 1 Oct 2020 16:29:52 +0100
Subject: [PATCH 03/13] Wire up new token when fetching events streams

---
 synapse/storage/databases/main/stream.py | 200 +++++++++++++++++++----
 synapse/types.py                         |  15 ++
 2 files changed, 179 insertions(+), 36 deletions(-)

diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index e6449d68b854..6eff289a0d42 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -53,6 +53,7 @@
 )
 from synapse.storage.databases.main.events_worker import EventsWorkerStore
 from synapse.storage.engines import BaseDatabaseEngine, PostgresEngine
+from synapse.storage.util.id_generators import MultiWriterIdGenerator
 from synapse.types import Collection, PersistedEventPosition, RoomStreamToken
 from synapse.util.caches.descriptors import cached
 from synapse.util.caches.stream_change_cache import StreamChangeCache
@@ -209,6 +210,49 @@ def _make_generic_sql_bound(
     )
 
 
+def _filter_results(
+    direction: str,
+    from_token: Optional[RoomStreamToken],
+    to_token: Optional[RoomStreamToken],
+    instance_name: str,
+    stream_ordering: int,
+) -> bool:
+    """Filter results from fetching events in the DB against the given tokens.
+
+    This is necessary to handle the case where the tokens include positions
+    maps, which we handle by fetching more than necessary from the DB and then
+    filtering (rather than attempting to construct a complicated SQL query).
+    """
+
+    # We will have already filtered by the topological tokens, so we don't
+    # bother checking topological token bounds again.
+    if from_token and from_token.topological:
+        from_token = None
+
+    if to_token and to_token.topological:
+        to_token = None
+
+    lower_bound = None
+    if direction == "f" and from_token:
+        lower_bound = from_token.get_stream_pos_for_instance(instance_name)
+    elif direction == "b" and to_token:
+        lower_bound = to_token.get_stream_pos_for_instance(instance_name)
+
+    if lower_bound and stream_ordering <= lower_bound:
+        return False
+
+    upper_bound = None
+    if direction == "b" and from_token:
+        upper_bound = from_token.get_stream_pos_for_instance(instance_name)
+    elif direction == "f" and to_token:
+        upper_bound = to_token.get_stream_pos_for_instance(instance_name)
+
+    if upper_bound and upper_bound < stream_ordering:
+        return False
+
+    return True
+
+
 def filter_to_clause(event_filter: Optional[Filter]) -> Tuple[str, List[str]]:
     # NB: This may create SQL clauses that don't optimise well (and we don't
     # have indices on all possible clauses). E.g. it may create
@@ -306,7 +350,26 @@ def get_room_min_stream_ordering(self) -> int:
         raise NotImplementedError()
 
     def get_room_max_token(self) -> RoomStreamToken:
-        return RoomStreamToken(None, self.get_room_max_stream_ordering())
+        min_pos = self._stream_id_gen.get_current_token()
+
+        positions = {}
+        if isinstance(self._stream_id_gen, MultiWriterIdGenerator):
+            # The `min_pos` is the minimum position that we know all instances
+            # have finished persisting to, so we only care about instances whose
+            # positions are ahead of that. (Instance positions can be behind the
+            # min position as there are times we can work out that the minimum
+            # position is ahead of the naive minimum across all current
+            # positions. See MultiWriterIdGenerator for details)
+            positions = {
+                i: p
+                for i, p in self._stream_id_gen.get_positions().items()
+                if p > min_pos
+            }
+
+            if set(positions.values()) == {min_pos}:
+                positions = {}
+
+        return RoomStreamToken(None, min_pos, positions)
 
     async def get_room_events_stream_for_rooms(
         self,
@@ -405,25 +468,39 @@ async def get_room_events_stream_for_room(
         if from_key == to_key:
             return [], from_key
 
-        from_id = from_key.stream
-        to_id = to_key.stream
-
-        has_changed = self._events_stream_cache.has_entity_changed(room_id, from_id)
+        has_changed = self._events_stream_cache.has_entity_changed(
+            room_id, from_key.stream
+        )
 
         if not has_changed:
             return [], from_key
 
         def f(txn):
-            sql = (
-                "SELECT event_id, stream_ordering FROM events WHERE"
-                " room_id = ?"
-                " AND not outlier"
-                " AND stream_ordering > ? AND stream_ordering <= ?"
-                " ORDER BY stream_ordering %s LIMIT ?"
-            ) % (order,)
-            txn.execute(sql, (room_id, from_id, to_id, limit))
-
-            rows = [_EventDictReturn(row[0], None, row[1]) for row in txn]
+            # To handle tokens with a non-empty instance_map we fetch more
+            # results than necessary and the filter down
+            min_from_id = from_key.stream
+            max_to_id = to_key.get_max_stream_pos()
+
+            sql = """
+                SELECT event_id, instance_name, stream_ordering
+                FROM events
+                WHERE
+                    room_id = ?
+                    AND not outlier
+                    AND stream_ordering > ? AND stream_ordering <= ?
+                ORDER BY stream_ordering %s LIMIT ?
+            """ % (
+                order,
+            )
+            txn.execute(sql, (room_id, min_from_id, max_to_id, 2 * limit))
+
+            rows = [
+                _EventDictReturn(event_id, None, stream_ordering)
+                for event_id, instance_name, stream_ordering in txn
+                if _filter_results(
+                    "f", from_key, to_key, instance_name, stream_ordering
+                )
+            ][:limit]
             return rows
 
         rows = await self.db_pool.runInteraction("get_room_events_stream_for_room", f)
@@ -432,7 +509,7 @@ def f(txn):
             [r.event_id for r in rows], get_prev_content=True
         )
 
-        self._set_before_and_after(ret, rows, topo_order=from_id is None)
+        self._set_before_and_after(ret, rows, topo_order=from_key.stream is None)
 
         if order.lower() == "desc":
             ret.reverse()
@@ -449,31 +526,39 @@ def f(txn):
     async def get_membership_changes_for_user(
         self, user_id: str, from_key: RoomStreamToken, to_key: RoomStreamToken
     ) -> List[EventBase]:
-        from_id = from_key.stream
-        to_id = to_key.stream
-
         if from_key == to_key:
             return []
 
-        if from_id:
+        if from_key:
             has_changed = self._membership_stream_cache.has_entity_changed(
-                user_id, int(from_id)
+                user_id, int(from_key.stream)
             )
             if not has_changed:
                 return []
 
         def f(txn):
-            sql = (
-                "SELECT m.event_id, stream_ordering FROM events AS e,"
-                " room_memberships AS m"
-                " WHERE e.event_id = m.event_id"
-                " AND m.user_id = ?"
-                " AND e.stream_ordering > ? AND e.stream_ordering <= ?"
-                " ORDER BY e.stream_ordering ASC"
-            )
-            txn.execute(sql, (user_id, from_id, to_id))
+            # To handle tokens with a non-empty instance_map we fetch more
+            # results than necessary and the filter down
+            min_from_id = from_key.stream
+            max_to_id = to_key.get_max_stream_pos()
+
+            sql = """
+                SELECT m.event_id, instance_name, stream_ordering
+                FROM events AS e, room_memberships AS m
+                WHERE e.event_id = m.event_id
+                    AND m.user_id = ?
+                    AND e.stream_ordering > ? AND e.stream_ordering <= ?
+                ORDER BY e.stream_ordering ASC
+            """
+            txn.execute(sql, (user_id, min_from_id, max_to_id,))
 
-            rows = [_EventDictReturn(row[0], None, row[1]) for row in txn]
+            rows = [
+                _EventDictReturn(event_id, None, stream_ordering)
+                for event_id, instance_name, stream_ordering in txn
+                if _filter_results(
+                    "f", from_key, to_key, instance_name, stream_ordering
+                )
+            ]
 
             return rows
 
@@ -980,11 +1065,44 @@ def _paginate_room_events_txn(
         else:
             order = "ASC"
 
+        # The bounds for the stream tokens are complicated by the fact the fact
+        # that we need to handle the instance_map part of the tokens. We do this
+        # by fetching all events between the min stream token and the maximum
+        # stream token (as return by `RoomStreamToken.get_max_stream_pos`) and
+        # then filtering the results.
+        if from_token.topological is not None:
+            from_bound = from_token.as_tuple()
+        elif direction == "b":
+            from_bound = (
+                None,
+                from_token.get_max_stream_pos(),
+            )
+        else:
+            from_bound = (
+                None,
+                from_token.stream,
+            )
+
+        to_bound = None
+        if to_token:
+            if to_token.topological is not None:
+                to_bound = to_token.as_tuple()
+            elif direction == "b":
+                to_bound = (
+                    None,
+                    to_token.stream,
+                )
+            else:
+                to_bound = (
+                    None,
+                    to_token.get_max_stream_pos(),
+                )
+
         bounds = generate_pagination_where_clause(
             direction=direction,
             column_names=("topological_ordering", "stream_ordering"),
-            from_token=from_token.as_tuple(),
-            to_token=to_token.as_tuple() if to_token else None,
+            from_token=from_bound,
+            to_token=to_bound,
             engine=self.database_engine,
         )
 
@@ -994,7 +1112,8 @@ def _paginate_room_events_txn(
             bounds += " AND " + filter_clause
             args.extend(filter_args)
 
-        args.append(int(limit))
+        # We fetch more events as we'll filter the result set
+        args.append(int(limit) * 2)
 
         select_keywords = "SELECT"
         join_clause = ""
@@ -1016,7 +1135,9 @@ def _paginate_room_events_txn(
                 select_keywords += "DISTINCT"
 
         sql = """
-            %(select_keywords)s event_id, topological_ordering, stream_ordering
+            %(select_keywords)s
+                event_id, instance_name,
+                topological_ordering, stream_ordering
             FROM events
             %(join_clause)s
             WHERE outlier = ? AND room_id = ? AND %(bounds)s
@@ -1031,7 +1152,14 @@ def _paginate_room_events_txn(
 
         txn.execute(sql, args)
 
-        rows = [_EventDictReturn(row[0], row[1], row[2]) for row in txn]
+        # Filter the result set.
+        rows = [
+            _EventDictReturn(event_id, topological_ordering, stream_ordering)
+            for event_id, instance_name, topological_ordering, stream_ordering in txn
+            if _filter_results(
+                direction, from_token, to_token, instance_name, stream_ordering
+            )
+        ][:limit]
 
         if rows:
             topo = rows[-1].topological_ordering
diff --git a/synapse/types.py b/synapse/types.py
index 550a968a03be..b2f18fe55eaa 100644
--- a/synapse/types.py
+++ b/synapse/types.py
@@ -488,6 +488,21 @@ def copy_and_advance(self, other: "RoomStreamToken") -> "RoomStreamToken":
     def as_tuple(self) -> Tuple[Optional[int], int]:
         return (self.topological, self.stream)
 
+    def get_stream_pos_for_instance(self, instance_name: str) -> int:
+        """Get the stream position for the instance
+        """
+        return self.instance_map.get(instance_name, self.stream)
+
+    def get_max_stream_pos(self) -> int:
+        """Get the maximum stream position referenced in this token.
+
+        The corresponding "min" position is, by definition just `self.stream`.
+
+        This is used to handle tokens that have non-empty `instance_map`, and so
+        reference stream positions after the `self.stream` position.
+        """
+        return max(self.instance_map.values(), default=self.stream)
+
     async def to_string(self, store: "DataStore") -> str:
         if self.topological is not None:
             return "t%d-%d" % (self.topological, self.stream)

From fdf6c05a979929c106baa24c7fcb6b7555e9758c Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Thu, 1 Oct 2020 16:32:03 +0100
Subject: [PATCH 04/13] Newsfile

---
 changelog.d/8439.misc | 1 +
 1 file changed, 1 insertion(+)
 create mode 100644 changelog.d/8439.misc

diff --git a/changelog.d/8439.misc b/changelog.d/8439.misc
new file mode 100644
index 000000000000..237cb3b31135
--- /dev/null
+++ b/changelog.d/8439.misc
@@ -0,0 +1 @@
+Allow events to be sent to clients sooner when using sharded event persisters.

From eed3cf588cff3523d704496ea3c785a904005346 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Fri, 2 Oct 2020 10:49:20 +0100
Subject: [PATCH 05/13] Typos

Co-authored-by: Andrew Morgan <1342360+anoadragon453@users.noreply.github.com>
---
 .../main/schema/delta/58/19instance_map.sql.postgres   |  2 +-
 synapse/storage/databases/main/stream.py               | 10 +++++-----
 synapse/types.py                                       |  2 +-
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres b/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
index bdbbab573835..c5788598cefa 100644
--- a/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
+++ b/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
@@ -15,7 +15,7 @@
 
 
 -- A unique and immutable mapping between instance name and an integer ID. This
--- let's us refer to instances via a small ID in e.g. stream tokens, without
+-- lets us refer to instances via a small ID in e.g. stream tokens, without
 -- having to encode the full name.
 CREATE TABLE instance_map (
     instance_id SERIAL PRIMARY KEY,
diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 6eff289a0d42..d280607f0da6 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -219,7 +219,7 @@ def _filter_results(
 ) -> bool:
     """Filter results from fetching events in the DB against the given tokens.
 
-    This is necessary to handle the case where the tokens include positions
+    This is necessary to handle the case where the tokens include position
     maps, which we handle by fetching more than necessary from the DB and then
     filtering (rather than attempting to construct a complicated SQL query).
     """
@@ -477,7 +477,7 @@ async def get_room_events_stream_for_room(
 
         def f(txn):
             # To handle tokens with a non-empty instance_map we fetch more
-            # results than necessary and the filter down
+            # results than necessary and then filter down
             min_from_id = from_key.stream
             max_to_id = to_key.get_max_stream_pos()
 
@@ -538,7 +538,7 @@ async def get_membership_changes_for_user(
 
         def f(txn):
             # To handle tokens with a non-empty instance_map we fetch more
-            # results than necessary and the filter down
+            # results than necessary and then filter down
             min_from_id = from_key.stream
             max_to_id = to_key.get_max_stream_pos()
 
@@ -1065,10 +1065,10 @@ def _paginate_room_events_txn(
         else:
             order = "ASC"
 
-        # The bounds for the stream tokens are complicated by the fact the fact
+        # The bounds for the stream tokens are complicated by the fact
         # that we need to handle the instance_map part of the tokens. We do this
         # by fetching all events between the min stream token and the maximum
-        # stream token (as return by `RoomStreamToken.get_max_stream_pos`) and
+        # stream token (as returned by `RoomStreamToken.get_max_stream_pos`) and
         # then filtering the results.
         if from_token.topological is not None:
             from_bound = from_token.as_tuple()
diff --git a/synapse/types.py b/synapse/types.py
index b2f18fe55eaa..dd8160f355cb 100644
--- a/synapse/types.py
+++ b/synapse/types.py
@@ -406,7 +406,7 @@ class RoomStreamToken:
     respective streams.
 
     The format of the token in such case is an initial integer min position,
-    followed by the mapping of instance ID to position seperate by '.' and '~':
+    followed by the mapping of instance ID to position separated by '.' and '~':
 
         m{min_pos}.{writer1}~{pos1}.{writer2}~{pos2}. ...
 

From f709da9e05d0a6eddd3fbd3af5dea863e9288d5e Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 6 Oct 2020 09:59:02 +0100
Subject: [PATCH 06/13] Apply suggestions from code review

Co-authored-by: Richard van der Hoff <1389908+richvdh@users.noreply.github.com>
---
 .../main/schema/delta/58/19instance_map.sql.postgres          | 4 ++--
 synapse/storage/databases/main/stream.py                      | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres b/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
index c5788598cefa..841186b8267e 100644
--- a/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
+++ b/synapse/storage/databases/main/schema/delta/58/19instance_map.sql.postgres
@@ -17,9 +17,9 @@
 -- A unique and immutable mapping between instance name and an integer ID. This
 -- lets us refer to instances via a small ID in e.g. stream tokens, without
 -- having to encode the full name.
-CREATE TABLE instance_map (
+CREATE TABLE IF NOT EXISTS instance_map (
     instance_id SERIAL PRIMARY KEY,
     instance_name TEXT NOT NULL
 );
 
-CREATE UNIQUE INDEX instance_map_idx ON instance_map(instance_name);
+CREATE UNIQUE INDEX IF NOT EXISTS instance_map_idx ON instance_map(instance_name);
diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 36afd17f4f34..8006568b0aba 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -1226,7 +1226,7 @@ async def paginate_room_events(
 
     @cached()
     async def get_id_for_instance(self, instance_name: str) -> int:
-        """Get a unique, immutable ID that corresponds to the instance.
+        """Get a unique, immutable ID that corresponds to the given Synapse worker instance.
         """
 
         def _get_id_for_instance_txn(txn):

From e91f0e9d13c3d15f6f48d51c31f739e6cd05c1cf Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 6 Oct 2020 10:29:29 +0100
Subject: [PATCH 07/13] Fix up RoomStreamToken

---
 synapse/storage/databases/main/stream.py |  8 +++--
 synapse/types.py                         | 37 +++++++++++++++++++++---
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 8006568b0aba..257c5584d0c6 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -1071,7 +1071,9 @@ def _paginate_room_events_txn(
         # stream token (as returned by `RoomStreamToken.get_max_stream_pos`) and
         # then filtering the results.
         if from_token.topological is not None:
-            from_bound = from_token.as_tuple()
+            from_bound = (
+                from_token.as_historical_tuple()
+            )  # type: Tuple[Optional[int], int]
         elif direction == "b":
             from_bound = (
                 None,
@@ -1083,10 +1085,10 @@ def _paginate_room_events_txn(
                 from_token.stream,
             )
 
-        to_bound = None
+        to_bound = None  # type: Optional[Tuple[Optional[int], int]]
         if to_token:
             if to_token.topological is not None:
-                to_bound = to_token.as_tuple()
+                to_bound = to_token.as_historical_tuple()
             elif direction == "b":
                 to_bound = (
                     None,
diff --git a/synapse/types.py b/synapse/types.py
index dd8160f355cb..dd36477e9d1c 100644
--- a/synapse/types.py
+++ b/synapse/types.py
@@ -396,7 +396,7 @@ class RoomStreamToken:
 
     Live tokens start with an "s" followed by the "stream_ordering" id of the
     event it comes after. Historic tokens start with a "t" followed by the
-    "topological_ordering" id of the event it comes after, followed by "~",
+    "topological_ordering" id of the event it comes after, followed by "-",
     followed by the "stream_ordering" id of the event it comes after.
 
     There is also a third mode for live tokens where the token starts with "m",
@@ -420,6 +420,9 @@ class RoomStreamToken:
     3 (these are instance IDs that can be looked up in the DB to fetch the more
     commonly used instance names) are at positions 58 and 59 respectively, and
     all other instances are at position 56.
+
+    Note: The `RoomStreamToken` cannot have both a topological part and an
+    instance map.
     """
 
     topological = attr.ib(
@@ -428,7 +431,19 @@ class RoomStreamToken:
     )
     stream = attr.ib(type=int, validator=attr.validators.instance_of(int))
 
-    instance_map = attr.ib(type=Dict[str, int], factory=dict)
+    instance_map = attr.ib(
+        type=Dict[str, int], factory=dict, validator=attr.validators.instance_of(dict)
+    )
+
+    @instance_map.validator
+    def _validate_topological_or_live(self, attribute, value):
+        """Validates that both `topological` and `instance_map` aren't set.
+        """
+
+        if value and self.topological:
+            raise ValueError(
+                "Cannot set both 'topological' and 'instance_map' on 'RoomStreamToken'."
+            )
 
     @classmethod
     async def parse(cls, store: "DataStore", string: str) -> "RoomStreamToken":
@@ -485,12 +500,26 @@ def copy_and_advance(self, other: "RoomStreamToken") -> "RoomStreamToken":
 
         return RoomStreamToken(None, max_stream, instance_map)
 
-    def as_tuple(self) -> Tuple[Optional[int], int]:
+    def as_historical_tuple(self) -> Tuple[int, int]:
+        """Returns a tuple of `(topological, stream)` for historical tokens.
+
+        Raises if not an historical token (i.e. doesn't have a topological part).
+        """
+        if self.topological is None:
+            raise Exception(
+                "Cannot call `RoomStreamToken.as_historical_tuple` on live token"
+            )
+
         return (self.topological, self.stream)
 
     def get_stream_pos_for_instance(self, instance_name: str) -> int:
-        """Get the stream position for the instance
+        """Get the stream position that the writer was at at this token.
+
+        This only makes sense for "live" tokens that may have a vector clock
+        component.
         """
+        # If we don't have an entry for the instance we can assume that it was
+        # at `self.stream`.
         return self.instance_map.get(instance_name, self.stream)
 
     def get_max_stream_pos(self) -> int:

From 604e33fbb3492a91019de6e2c0c612c6a17aa5fe Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 6 Oct 2020 11:44:10 +0100
Subject: [PATCH 08/13] Fix up _filter_results

---
 synapse/storage/databases/main/stream.py | 94 ++++++++++++++----------
 1 file changed, 56 insertions(+), 38 deletions(-)

diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 257c5584d0c6..cd708e1c991c 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -211,44 +211,50 @@ def _make_generic_sql_bound(
 
 
 def _filter_results(
-    direction: str,
-    from_token: Optional[RoomStreamToken],
-    to_token: Optional[RoomStreamToken],
+    lower_token: Optional[RoomStreamToken],
+    upper_token: Optional[RoomStreamToken],
     instance_name: str,
+    topological_ordering: int,
     stream_ordering: int,
 ) -> bool:
-    """Filter results from fetching events in the DB against the given tokens.
-
-    This is necessary to handle the case where the tokens include position
-    maps, which we handle by fetching more than necessary from the DB and then
-    filtering (rather than attempting to construct a complicated SQL query).
+    """Returns True if the event persisted by the given instance at the given
+    topological/stream_ordering falls between the two tokens (taking a None
+    token to mean unbounded).
+
+    Used to filter results from fetching events in the DB against the given
+    tokens. This is necessary to handle the case where the tokens include
+    position maps, which we handle by fetching more than necessary from the DB
+    and then filtering (rather than attempting to construct a complicated SQL
+    query).
     """
 
-    # We will have already filtered by the topological tokens, so we don't
-    # bother checking topological token bounds again.
-    if from_token and from_token.topological:
-        from_token = None
-
-    if to_token and to_token.topological:
-        to_token = None
-
-    lower_bound = None
-    if direction == "f" and from_token:
-        lower_bound = from_token.get_stream_pos_for_instance(instance_name)
-    elif direction == "b" and to_token:
-        lower_bound = to_token.get_stream_pos_for_instance(instance_name)
-
-    if lower_bound and stream_ordering <= lower_bound:
-        return False
+    event_historical_tuple = (
+        topological_ordering,
+        stream_ordering,
+    )
 
-    upper_bound = None
-    if direction == "b" and from_token:
-        upper_bound = from_token.get_stream_pos_for_instance(instance_name)
-    elif direction == "f" and to_token:
-        upper_bound = to_token.get_stream_pos_for_instance(instance_name)
+    if lower_token:
+        if lower_token.topological:
+            # If these are historical tokens we compare the `(topological, stream)`
+            # tuples.
+            if event_historical_tuple <= lower_token.as_historical_tuple():
+                return False
 
-    if upper_bound and upper_bound < stream_ordering:
-        return False
+        else:
+            # If these are live tokens we compare the stream ordering against the
+            # writers stream position.
+            if stream_ordering <= lower_token.get_stream_pos_for_instance(
+                instance_name
+            ):
+                return False
+
+    if upper_token:
+        if upper_token.topological:
+            if upper_token.as_historical_tuple() < event_historical_tuple:
+                return False
+        else:
+            if upper_token.get_stream_pos_for_instance(instance_name) < stream_ordering:
+                return False
 
     return True
 
@@ -482,7 +488,7 @@ def f(txn):
             max_to_id = to_key.get_max_stream_pos()
 
             sql = """
-                SELECT event_id, instance_name, stream_ordering
+                SELECT event_id, instance_name, topological_ordering, stream_ordering
                 FROM events
                 WHERE
                     room_id = ?
@@ -496,9 +502,13 @@ def f(txn):
 
             rows = [
                 _EventDictReturn(event_id, None, stream_ordering)
-                for event_id, instance_name, stream_ordering in txn
+                for event_id, instance_name, topological_ordering, stream_ordering in txn
                 if _filter_results(
-                    "f", from_key, to_key, instance_name, stream_ordering
+                    from_key,
+                    to_key,
+                    instance_name,
+                    topological_ordering,
+                    stream_ordering,
                 )
             ][:limit]
             return rows
@@ -543,7 +553,7 @@ def f(txn):
             max_to_id = to_key.get_max_stream_pos()
 
             sql = """
-                SELECT m.event_id, instance_name, stream_ordering
+                SELECT m.event_id, instance_name, topological_ordering, stream_ordering
                 FROM events AS e, room_memberships AS m
                 WHERE e.event_id = m.event_id
                     AND m.user_id = ?
@@ -554,9 +564,13 @@ def f(txn):
 
             rows = [
                 _EventDictReturn(event_id, None, stream_ordering)
-                for event_id, instance_name, stream_ordering in txn
+                for event_id, instance_name, topological_ordering, stream_ordering in txn
                 if _filter_results(
-                    "f", from_key, to_key, instance_name, stream_ordering
+                    from_key,
+                    to_key,
+                    instance_name,
+                    topological_ordering,
+                    stream_ordering,
                 )
             ]
 
@@ -1159,7 +1173,11 @@ def _paginate_room_events_txn(
             _EventDictReturn(event_id, topological_ordering, stream_ordering)
             for event_id, instance_name, topological_ordering, stream_ordering in txn
             if _filter_results(
-                direction, from_token, to_token, instance_name, stream_ordering
+                lower_token=to_token if direction == "b" else from_token,
+                upper_token=from_token if direction == "b" else to_token,
+                instance_name=instance_name,
+                topological_ordering=topological_ordering,
+                stream_ordering=stream_ordering,
             )
         ][:limit]
 

From 6425f430d0ab7fae374c22fbbd735db9d2add874 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 6 Oct 2020 11:52:23 +0100
Subject: [PATCH 09/13] Fixup get_room_max_token

---
 synapse/storage/databases/main/stream.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index cd708e1c991c..6a950f1ba33c 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -356,6 +356,14 @@ def get_room_min_stream_ordering(self) -> int:
         raise NotImplementedError()
 
     def get_room_max_token(self) -> RoomStreamToken:
+        """Get a `RoomStreamToken` that marks the current maximum persisted
+        position of the events stream. Useful to get a token that represents
+        "now".
+
+        The token returned is a "live" token that may have an instance_map
+        component.
+        """
+
         min_pos = self._stream_id_gen.get_current_token()
 
         positions = {}
@@ -372,9 +380,6 @@ def get_room_max_token(self) -> RoomStreamToken:
                 if p > min_pos
             }
 
-            if set(positions.values()) == {min_pos}:
-                positions = {}
-
         return RoomStreamToken(None, min_pos, positions)
 
     async def get_room_events_stream_for_rooms(

From 93c809114541c50d2528f9d57f029ee288bc0098 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 6 Oct 2020 11:57:07 +0100
Subject: [PATCH 10/13] Fix always False condition to just be False

---
 synapse/storage/databases/main/stream.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index 6a950f1ba33c..cdfff1fc4464 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -524,7 +524,7 @@ def f(txn):
             [r.event_id for r in rows], get_prev_content=True
         )
 
-        self._set_before_and_after(ret, rows, topo_order=from_key.stream is None)
+        self._set_before_and_after(ret, rows, topo_order=False)
 
         if order.lower() == "desc":
             ret.reverse()

From a321252e33f574be7dde9246e26a9c61eb834c8b Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Tue, 6 Oct 2020 11:58:53 +0100
Subject: [PATCH 11/13] Switch '.' and '~' in new token format

---
 synapse/types.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/synapse/types.py b/synapse/types.py
index dd36477e9d1c..b98cbdc15795 100644
--- a/synapse/types.py
+++ b/synapse/types.py
@@ -408,13 +408,13 @@ class RoomStreamToken:
     The format of the token in such case is an initial integer min position,
     followed by the mapping of instance ID to position separated by '.' and '~':
 
-        m{min_pos}.{writer1}~{pos1}.{writer2}~{pos2}. ...
+        m{min_pos}~{writer1}.{pos1}~{writer2}.{pos2}. ...
 
     The `min_pos` corresponds to the minimum position all writers have persisted
     up to, and then only writers that are ahead of that position need to be
     encoded. An example token is:
 
-        m56.2~58.3~59
+        m56~2.58~3.59
 
     Which corresponds to a set of three (or more writers) where instances 2 and
     3 (these are instance IDs that can be looked up in the DB to fetch the more
@@ -454,12 +454,12 @@ async def parse(cls, store: "DataStore", string: str) -> "RoomStreamToken":
                 parts = string[1:].split("-", 1)
                 return cls(topological=int(parts[0]), stream=int(parts[1]))
             if string[0] == "m":
-                parts = string[1:].split(".")
+                parts = string[1:].split("~")
                 stream = int(parts[0])
 
                 instance_map = {}
                 for part in parts[1:]:
-                    key, value = part.split("~")
+                    key, value = part.split(".")
                     instance_id = int(key)
                     pos = int(value)
 
@@ -539,10 +539,10 @@ async def to_string(self, store: "DataStore") -> str:
             entries = []
             for name, pos in self.instance_map.items():
                 instance_id = await store.get_id_for_instance(name)
-                entries.append("{}~{}".format(instance_id, pos))
+                entries.append("{}.{}".format(instance_id, pos))
 
-            encoded_map = ".".join(entries)
-            return "m{}.{}".format(self.stream, encoded_map)
+            encoded_map = "~".join(entries)
+            return "m{}~{}".format(self.stream, encoded_map)
         else:
             return "s%d" % (self.stream,)
 

From cb7c48cec86711dadfccd2f61f35ecb307a28e77 Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 7 Oct 2020 13:56:41 +0100
Subject: [PATCH 12/13] Improve RoomStreamToken

---
 synapse/types.py | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/synapse/types.py b/synapse/types.py
index b98cbdc15795..5bde67cc078a 100644
--- a/synapse/types.py
+++ b/synapse/types.py
@@ -432,15 +432,20 @@ class RoomStreamToken:
     stream = attr.ib(type=int, validator=attr.validators.instance_of(int))
 
     instance_map = attr.ib(
-        type=Dict[str, int], factory=dict, validator=attr.validators.instance_of(dict)
+        type=Dict[str, int],
+        factory=dict,
+        validator=attr.validators.deep_mapping(
+            key_validator=attr.validators.instance_of(str),
+            value_validator=attr.validators.instance_of(int),
+            mapping_validator=attr.validators.instance_of(dict),
+        ),
     )
 
-    @instance_map.validator
-    def _validate_topological_or_live(self, attribute, value):
+    def __attrs_post_init__(self):
         """Validates that both `topological` and `instance_map` aren't set.
         """
 
-        if value and self.topological:
+        if self.instance_map and self.topological:
             raise ValueError(
                 "Cannot set both 'topological' and 'instance_map' on 'RoomStreamToken'."
             )
@@ -513,11 +518,13 @@ def as_historical_tuple(self) -> Tuple[int, int]:
         return (self.topological, self.stream)
 
     def get_stream_pos_for_instance(self, instance_name: str) -> int:
-        """Get the stream position that the writer was at at this token.
+        """Get the stream position that the given writer was at at this token.
 
         This only makes sense for "live" tokens that may have a vector clock
-        component.
+        component, and so asserts that this is a "live" token.
         """
+        assert self.topological is None
+
         # If we don't have an entry for the instance we can assume that it was
         # at `self.stream`.
         return self.instance_map.get(instance_name, self.stream)

From 05c2058fc0ca44a1ef4786ff6b657cd8f37b448d Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 7 Oct 2020 14:45:10 +0100
Subject: [PATCH 13/13] Fix unit tests

---
 synapse/storage/databases/main/stream.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/synapse/storage/databases/main/stream.py b/synapse/storage/databases/main/stream.py
index f713148f5259..e3b9ff5ca6b3 100644
--- a/synapse/storage/databases/main/stream.py
+++ b/synapse/storage/databases/main/stream.py
@@ -234,7 +234,7 @@ def _filter_results(
     )
 
     if lower_token:
-        if lower_token.topological:
+        if lower_token.topological is not None:
             # If these are historical tokens we compare the `(topological, stream)`
             # tuples.
             if event_historical_tuple <= lower_token.as_historical_tuple():
@@ -249,7 +249,7 @@ def _filter_results(
                 return False
 
     if upper_token:
-        if upper_token.topological:
+        if upper_token.topological is not None:
             if upper_token.as_historical_tuple() < event_historical_tuple:
                 return False
         else: