synapse/tests/storage/test_id_generators.py

#
# This file is licensed under the Affero General Public License (AGPL) version 3.
#
# Copyright 2020 The Matrix.org Foundation C.I.C.
# Copyright (C) 2023 New Vector, Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# See the GNU Affero General Public License for more details:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
#
# Originally licensed under the Apache License, Version 2.0:
# <http://www.apache.org/licenses/LICENSE-2.0>.
#
# [This file includes modifications made by New Vector Limited]
#
#

import sqlite3
from unittest import mock

from immutabledict import immutabledict

from twisted.internet.defer import CancelledError, Deferred, ensureDeferred
from twisted.internet.testing import MemoryReactor

from synapse.logging.context import LoggingContext, make_deferred_yieldable
from synapse.server import HomeServer
from synapse.storage.database import (
    DatabasePool,
    LoggingDatabaseConnection,
    LoggingTransaction,
)
from synapse.storage.engines import Sqlite3Engine
from synapse.storage.types import Cursor
from synapse.storage.util.id_generators import (
    MultiWriterIdGenerator,
    advance_multiwriter_sharded_token_after_partial_read,
    make_multiwriter_sharded_token_bounds_sql,
    stream_current_position_gauge,
)
from synapse.storage.util.sequence import (
    LocalSequenceGenerator,
    PostgresSequenceGenerator,
    SequenceGenerator,
)
from synapse.types import MultiWriterStreamToken
from synapse.util.clock import Clock

from tests.unittest import HomeserverTestCase, TestCase
from tests.utils import USE_POSTGRES_FOR_TESTS


class MultiWriterIdGeneratorBase(HomeserverTestCase):
    positive: bool = True
    tables: list[str] = ["foobar"]

    def prepare(self, reactor: MemoryReactor, clock: Clock, hs: HomeServer) -> None:
        self.store = hs.get_datastores().main
        self.db_pool: DatabasePool = self.store.db_pool
        self.instances: dict[str, MultiWriterIdGenerator] = {}

        self.get_success(self.db_pool.runInteraction("_setup_db", self._setup_db))

        if USE_POSTGRES_FOR_TESTS:
            self.seq_gen: SequenceGenerator = PostgresSequenceGenerator("foobar_seq")
        else:
            self.seq_gen = LocalSequenceGenerator(lambda _: 0)

    def _setup_db(self, txn: LoggingTransaction) -> None:
        if USE_POSTGRES_FOR_TESTS:
            txn.execute("CREATE SEQUENCE foobar_seq")

        for table in self.tables:
            txn.execute(
                """
                CREATE TABLE %s (
                    stream_id BIGINT NOT NULL,
                    instance_name TEXT NOT NULL,
                    data TEXT
                );
                """
                % (table,)
            )

    def _create_id_generator(
        self,
        instance_name: str = "master",
        writers: list[str] | None = None,
    ) -> MultiWriterIdGenerator:
        def _create(conn: LoggingDatabaseConnection) -> MultiWriterIdGenerator:
            id_gen = MultiWriterIdGenerator(
                db_conn=conn,
                db=self.db_pool,
                notifier=self.hs.get_replication_notifier(),
                stream_name="test_stream",
                server_name=self.hs.hostname,
                instance_name=instance_name,
                tables=[(table, "instance_name", "stream_id") for table in self.tables],
                sequence_name="foobar_seq",
                writers=writers or ["master"],
                positive=self.positive,
            )
            # Constructing the generator prunes stale `stream_positions` rows
            # (writers no longer in the config); commit so that persists for the
            # next generator we create.
            #
            # Note we need to commit manually here as the generator is created
            # in a `runWithConnection` call, which doesn't automatically
            # commit/rollback.
            conn.commit()
            return id_gen

        self.instances[instance_name] = self.get_success_or_raise(
            self.db_pool.runWithConnection(_create)
        )
        return self.instances[instance_name]

    def _get_reported_metric_position(self) -> int:
        """The position the metric reports for the test stream."""

        return self.get_prometheus_metric_current_value(
            stream_current_position_gauge,
            stream_name="test_stream",
        )

    def _replicate(self, instance_name: str) -> None:
        """Similate a replication event for the given instance."""

        writer = self.instances[instance_name]
        token = writer.get_current_token_for_writer(instance_name)
        for generator in self.instances.values():
            if writer != generator:
                generator.advance(instance_name, token)

    def _replicate_all(self) -> None:
        """Similate a replication event for all instances."""

        for instance_name in self.instances:
            self._replicate(instance_name)

    def _insert_row(
        self, instance_name: str, stream_id: int, table: str | None = None
    ) -> None:
        """Insert one row as the given instance with given stream_id."""

        if table is None:
            table = self.tables[0]

        factor = 1 if self.positive else -1

        def _insert(txn: LoggingTransaction) -> None:
            txn.execute(
                "INSERT INTO %s VALUES (?, ?)" % (table,),
                (
                    stream_id,
                    instance_name,
                ),
            )
            txn.execute(
                """
                INSERT INTO stream_positions VALUES ('test_stream', ?, ?)
                ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = ?
                """,
                (instance_name, stream_id * factor, stream_id * factor),
            )

        self.get_success(self.db_pool.runInteraction("_insert_row", _insert))

    def _insert_rows(
        self,
        instance_name: str,
        number: int,
        table: str | None = None,
        update_stream_table: bool = True,
    ) -> None:
        """Insert N rows as the given instance, inserting with stream IDs pulled
        from the postgres sequence.
        """

        if table is None:
            table = self.tables[0]

        factor = 1 if self.positive else -1

        def _insert(txn: LoggingTransaction) -> None:
            for _ in range(number):
                next_val = self.seq_gen.get_next_id_txn(txn)
                txn.execute(
                    "INSERT INTO %s (stream_id, instance_name) VALUES (?, ?)"
                    % (table,),
                    (next_val, instance_name),
                )

                if update_stream_table:
                    txn.execute(
                        """
                        INSERT INTO stream_positions VALUES ('test_stream', ?,  ?)
                        ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = ?
                        """,
                        (instance_name, next_val * factor, next_val * factor),
                    )

        self.get_success(self.db_pool.runInteraction("_insert_rows", _insert))


class MultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase):
    def test_empty(self) -> None:
        """Test an ID generator against an empty database gives sensible
        current positions.
        """

        id_gen = self._create_id_generator()

        # The table is empty so we expect the map for positions to have a dummy
        # minimum value.
        self.assertEqual(id_gen.get_positions(), {"master": 1})

    def test_single_instance(self) -> None:
        """Test that reads and writes from a single process are handled
        correctly.
        """

        # Prefill table with 7 rows written by 'master'
        self._insert_rows("master", 7)

        id_gen = self._create_id_generator()

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        # Try allocating a new ID gen and check that we only see position
        # advanced after we leave the context manager.

        async def _get_next_async() -> None:
            async with id_gen.get_next() as stream_id:
                self.assertEqual(stream_id, 8)

                self.assertEqual(id_gen.get_positions(), {"master": 7})
                self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        self.get_success(_get_next_async())

        self.assertEqual(id_gen.get_positions(), {"master": 8})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 8)

    def test_current_position_metric(self) -> None:
        """The reported position follows `get_current_token`."""

        self._insert_rows("master", 7)

        id_gen = self._create_id_generator()

        self.assertEqual(self._get_reported_metric_position(), 7)

        async def _get_next_async() -> None:
            async with id_gen.get_next():
                pass

        self.get_success(_get_next_async())

        self.assertEqual(id_gen.get_current_token(), 8)
        self.assertEqual(self._get_reported_metric_position(), 8)

    def test_cancelled_enter_does_not_wedge_position(self) -> None:
        """Reproduces presence getting stuck.

        If the `get_next()` async context manager is cancelled while
        `__aenter__` is allocating a stream ID, the DB interaction that runs the
        sequence has already recorded the ID as unfinished, but `__aexit__`
        is never called (Python only invokes `__aexit__` if `__aenter__`
        returned). The abandoned ID is therefore leaked into the unfinished IDs
        forever, which permanently pins the persisted stream position: new rows
        keep getting higher IDs, but `get_current_token()` can never advance past
        `leaked_id - 1` until the process restarts.

        This mirrors a `/sync` request being cancelled part-way through
        persisting a presence update. `/sync` became `@cancellable` in #19499,
        and on a monolith the presence write in `PresenceStore.update_presence`
        is awaited inside that cancellable request scope.
        """
        # Prefill table with 7 rows written by 'master'; position starts at 7.
        self._insert_rows("master", 7)

        id_gen = self._create_id_generator()
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        # We model the cancellation at the seam it actually happens in
        # production: `__aenter__` awaits `runInteraction("_load_next_mult_id")`,
        # whose transaction runs in a thread pool and so *always* completes -
        # allocating stream ID 8 and recording it as unfinished - but the
        # awaiting coroutine is handed a `CancelledError` because the enclosing
        # `/sync` request was cancelled. We reproduce that by letting the real
        # interaction run (applying its side effects) and then failing the
        # awaited deferred with `CancelledError`.
        cancel_enter: "Deferred[None]" = Deferred()
        original_run_interaction = id_gen._db.runInteraction

        async def blocking_run_interaction(desc, func, *args, **kwargs):  # type: ignore[no-untyped-def]
            result = await original_run_interaction(desc, func, *args, **kwargs)
            if desc == "_load_next_mult_id":
                # Stream ID 8 is now allocated and recorded as unfinished.
                # Deliver the cancellation here, exactly as a cancelled `/sync`
                # would land it on this `await`.
                await make_deferred_yieldable(cancel_enter)
            return result

        async def presence_like_write() -> None:
            # Mirrors `PresenceStore.update_presence`: allocate an ID and
            # "persist" under the context manager.
            with LoggingContext(name="sync", server_name=self.hs.hostname):
                async with id_gen.get_next():
                    pass

        with mock.patch.object(
            id_gen._db, "runInteraction", new=blocking_run_interaction
        ):
            write = ensureDeferred(presence_like_write())

            # The write is now blocked inside `__aenter__`, i.e. after stream ID
            # 8 has been allocated and recorded as unfinished.
            self.assertNoResult(write)

            # The client goes away and the `/sync` request is cancelled.
            cancel_enter.errback(CancelledError())

            # The cancellation must surface as a `CancelledError`.
            self.get_failure(write, CancelledError)

        # The cancelled write never persisted a row for ID 8, so the generator
        # must not let that abandoned ID wedge the position. A subsequent
        # *successful* write should be able to advance the persisted token.
        async def _successful_write() -> None:
            async with id_gen.get_next():
                pass

        self.get_success(_successful_write())

        # On the buggy code the token is still stuck at 7 (ID 8 is leaked in
        # the unfinished IDs, blocking everything behind it). Once the leak is
        # fixed, the token advances to 9: ID 8 was allocated (and abandoned) by
        # the cancelled write, so the successful write above takes ID 9.
        self.assertEqual(
            id_gen.get_current_token_for_writer("master"),
            9,
            "presence stream position is wedged by the cancelled allocation",
        )

    def test_out_of_order_finish(self) -> None:
        """Test that IDs persisted out of order are correctly handled"""

        # Prefill table with 7 rows written by 'master'
        self._insert_rows("master", 7)

        id_gen = self._create_id_generator()

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        ctx1 = id_gen.get_next()
        ctx2 = id_gen.get_next()
        ctx3 = id_gen.get_next()
        ctx4 = id_gen.get_next()

        s1 = self.get_success(ctx1.__aenter__())
        s2 = self.get_success(ctx2.__aenter__())
        s3 = self.get_success(ctx3.__aenter__())
        s4 = self.get_success(ctx4.__aenter__())

        self.assertEqual(s1, 8)
        self.assertEqual(s2, 9)
        self.assertEqual(s3, 10)
        self.assertEqual(s4, 11)

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        self.get_success(ctx2.__aexit__(None, None, None))

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        self.get_success(ctx1.__aexit__(None, None, None))

        self.assertEqual(id_gen.get_positions(), {"master": 9})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 9)

        self.get_success(ctx4.__aexit__(None, None, None))

        self.assertEqual(id_gen.get_positions(), {"master": 9})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 9)

        self.get_success(ctx3.__aexit__(None, None, None))

        self.assertEqual(id_gen.get_positions(), {"master": 11})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 11)

    def test_get_next_txn(self) -> None:
        """Test that the `get_next_txn` function works correctly."""

        # Prefill table with 7 rows written by 'master'
        self._insert_rows("master", 7)

        id_gen = self._create_id_generator()

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        # Try allocating a new ID gen and check that we only see position
        # advanced after we leave the context manager.

        def _get_next_txn(txn: LoggingTransaction) -> None:
            stream_id = id_gen.get_next_txn(txn)
            self.assertEqual(stream_id, 8)

            self.assertEqual(id_gen.get_positions(), {"master": 7})
            self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        self.get_success(self.db_pool.runInteraction("test", _get_next_txn))

        self.assertEqual(id_gen.get_positions(), {"master": 8})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 8)

    def test_restart_during_out_of_order_persistence(self) -> None:
        """Test that restarting a process while another process is writing out
        of order updates are handled correctly.
        """

        # Prefill table with 7 rows written by 'master'
        self._insert_rows("master", 7)

        id_gen = self._create_id_generator()

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        # Persist two rows at once
        ctx1 = id_gen.get_next()
        ctx2 = id_gen.get_next()

        s1 = self.get_success(ctx1.__aenter__())
        s2 = self.get_success(ctx2.__aenter__())

        self.assertEqual(s1, 8)
        self.assertEqual(s2, 9)

        self.assertEqual(id_gen.get_positions(), {"master": 7})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), 7)

        # We finish persisting the second row before restart
        self.get_success(ctx2.__aexit__(None, None, None))

        # We simulate a restart of another worker by just creating a new ID gen.
        id_gen_worker = self._create_id_generator("worker")

        # Restarted worker should not see the second persisted row
        self.assertEqual(id_gen_worker.get_positions(), {"master": 7})
        self.assertEqual(id_gen_worker.get_current_token_for_writer("master"), 7)

        # Now if we persist the first row then both instances should jump ahead
        # correctly.
        self.get_success(ctx1.__aexit__(None, None, None))

        self.assertEqual(id_gen.get_positions(), {"master": 9})
        id_gen_worker.advance("master", 9)
        self.assertEqual(id_gen_worker.get_positions(), {"master": 9})


class WorkerMultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase):
    if not USE_POSTGRES_FOR_TESTS:
        skip = "Requires Postgres"

    def _insert_row_with_id(self, instance_name: str, stream_id: int) -> None:
        """Insert one row as the given instance with given stream_id, updating
        the postgres sequence position to match.
        """

        def _insert(txn: LoggingTransaction) -> None:
            txn.execute(
                "INSERT INTO foobar (stream_id, instance_name) VALUES (?, ?)",
                (
                    stream_id,
                    instance_name,
                ),
            )

            txn.execute("SELECT setval('foobar_seq', ?)", (stream_id,))

            txn.execute(
                """
                INSERT INTO stream_positions VALUES ('test_stream', ?, ?)
                ON CONFLICT (stream_name, instance_name) DO UPDATE SET stream_id = ?
                """,
                (instance_name, stream_id, stream_id),
            )

        self.get_success(self.db_pool.runInteraction("_insert_row_with_id", _insert))

    def test_get_persisted_upto_position(self) -> None:
        """Test that `get_persisted_upto_position` correctly tracks updates to
        positions.
        """

        # The following tests are a bit cheeky in that we notify about new
        # positions via `advance` without *actually* advancing the postgres
        # sequence.

        self._insert_row_with_id("first", 3)
        self._insert_row_with_id("second", 5)

        id_gen = self._create_id_generator("worker", writers=["first", "second"])

        self.assertEqual(id_gen.get_positions(), {"first": 3, "second": 5})

        # Min is 3 and there is a gap between 5, so we expect it to be 3.
        self.assertEqual(id_gen.get_persisted_upto_position(), 3)

        # We advance "first" straight to 6. Min is now 5 but there is no gap so
        # we expect it to be 6
        id_gen.advance("first", 6)
        self.assertEqual(id_gen.get_persisted_upto_position(), 6)

        # No gap, so we expect 7.
        id_gen.advance("second", 7)
        self.assertEqual(id_gen.get_persisted_upto_position(), 7)

        # We haven't seen 8 yet, so we expect 7 still.
        id_gen.advance("second", 9)
        self.assertEqual(id_gen.get_persisted_upto_position(), 7)

        # Now that we've seen 7, 8 and 9 we can got straight to 9.
        id_gen.advance("first", 8)
        self.assertEqual(id_gen.get_persisted_upto_position(), 9)

        # Jump forward with gaps. The minimum is 11, even though we haven't seen
        # 10 we know that everything before 11 must be persisted.
        id_gen.advance("first", 11)
        id_gen.advance("second", 15)
        self.assertEqual(id_gen.get_persisted_upto_position(), 11)

    def test_get_persisted_upto_position_get_next(self) -> None:
        """Test that `get_persisted_upto_position` correctly tracks updates to
        positions when `get_next` is called.
        """

        self._insert_row_with_id("first", 3)
        self._insert_row_with_id("second", 5)

        id_gen = self._create_id_generator("first", writers=["first", "second"])

        # When the writer is created, it assumes its own position is the current head of
        # the sequence
        self.assertEqual(id_gen.get_positions(), {"first": 5, "second": 5})

        self.assertEqual(id_gen.get_persisted_upto_position(), 5)

        async def _get_next_async() -> None:
            async with id_gen.get_next() as stream_id:
                self.assertEqual(stream_id, 6)
                self.assertEqual(id_gen.get_persisted_upto_position(), 5)

        self.get_success(_get_next_async())

        self.assertEqual(id_gen.get_persisted_upto_position(), 6)

        # We assume that so long as `get_next` does correctly advance the
        # `persisted_upto_position` in this case, then it will be correct in the
        # other cases that are tested above (since they'll hit the same code).

    def test_multi_instance(self) -> None:
        """Test that reads and writes from multiple processes are handled
        correctly.
        """
        self._insert_rows("first", 3)
        first_id_gen = self._create_id_generator("first", writers=["first", "second"])

        self._insert_rows("second", 4)
        second_id_gen = self._create_id_generator("second", writers=["first", "second"])

        self._replicate_all()

        self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 7})
        self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)

        self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)

        # Try allocating a new ID gen and check that we only see position
        # advanced after we leave the context manager.

        async def _get_next_async() -> None:
            async with first_id_gen.get_next() as stream_id:
                self.assertEqual(stream_id, 8)

                self.assertEqual(
                    first_id_gen.get_positions(), {"first": 3, "second": 7}
                )
                self.assertEqual(
                    second_id_gen.get_positions(), {"first": 3, "second": 7}
                )
                self.assertEqual(first_id_gen.get_persisted_upto_position(), 7)

        self.get_success(_get_next_async())

        self.assertEqual(first_id_gen.get_positions(), {"first": 8, "second": 7})

        # However the ID gen on the second instance won't have seen the update
        self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})

        # ... but calling `get_next` on the second instance should give a unique
        # stream ID

        async def _get_next_async2() -> None:
            async with second_id_gen.get_next() as stream_id:
                self.assertEqual(stream_id, 9)

                self.assertEqual(
                    second_id_gen.get_positions(), {"first": 3, "second": 7}
                )

        self.get_success(_get_next_async2())

        self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 9})

        # If the second ID gen gets told about the first, it correctly updates
        second_id_gen.advance("first", 8)
        self.assertEqual(second_id_gen.get_positions(), {"first": 8, "second": 9})

    def test_multi_instance_empty_row(self) -> None:
        """Test that reads and writes from multiple processes are handled
        correctly, when one of the writers starts without any rows.
        """
        # Insert some rows for two out of three of the ID gens.
        self._insert_rows("first", 3)
        first_id_gen = self._create_id_generator(
            "first", writers=["first", "second", "third"]
        )

        self._insert_rows("second", 4)
        second_id_gen = self._create_id_generator(
            "second", writers=["first", "second", "third"]
        )
        third_id_gen = self._create_id_generator(
            "third", writers=["first", "second", "third"]
        )

        self._replicate_all()

        self.assertEqual(
            first_id_gen.get_positions(), {"first": 3, "second": 7, "third": 7}
        )
        self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)
        self.assertEqual(first_id_gen.get_current_token_for_writer("third"), 7)

        self.assertEqual(
            second_id_gen.get_positions(), {"first": 3, "second": 7, "third": 7}
        )
        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
        self.assertEqual(second_id_gen.get_current_token_for_writer("third"), 7)

        # Try allocating a new ID gen and check that we only see position
        # advanced after we leave the context manager.

        async def _get_next_async() -> None:
            async with third_id_gen.get_next() as stream_id:
                self.assertEqual(stream_id, 8)

                self.assertEqual(
                    third_id_gen.get_positions(), {"first": 3, "second": 7, "third": 7}
                )
                self.assertEqual(third_id_gen.get_persisted_upto_position(), 7)

        self.get_success(_get_next_async())

        self.assertEqual(
            third_id_gen.get_positions(), {"first": 3, "second": 7, "third": 8}
        )

    def test_writer_config_change(self) -> None:
        """Test that changing the writer config correctly works."""

        self._insert_row_with_id("first", 3)
        self._insert_row_with_id("second", 5)

        # Initial config has two writers
        id_gen = self._create_id_generator("worker", writers=["first", "second"])
        self.assertEqual(id_gen.get_persisted_upto_position(), 3)
        self.assertEqual(id_gen.get_current_token_for_writer("first"), 3)
        self.assertEqual(id_gen.get_current_token_for_writer("second"), 5)

        # New config removes one of the configs. Note that if the writer is
        # removed from config we assume that it has been shut down and has
        # finished persisting, hence why the persisted upto position is 5.
        id_gen_2 = self._create_id_generator("second", writers=["second"])
        self.assertEqual(id_gen_2.get_persisted_upto_position(), 5)
        self.assertEqual(id_gen_2.get_current_token_for_writer("second"), 5)

        # This config points to a single, previously unused writer.
        id_gen_3 = self._create_id_generator("third", writers=["third"])
        self.assertEqual(id_gen_3.get_persisted_upto_position(), 5)

        # For new writers we assume their initial position to be the current
        # persisted up to position. This stops Synapse from doing a full table
        # scan when a new writer comes along.
        self.assertEqual(id_gen_3.get_current_token_for_writer("third"), 5)

        id_gen_4 = self._create_id_generator("fourth", writers=["third"])
        self.assertEqual(id_gen_4.get_current_token_for_writer("third"), 5)

        # Check that we get a sane next stream ID with this new config.

        async def _get_next_async() -> None:
            async with id_gen_3.get_next() as stream_id:
                self.assertEqual(stream_id, 6)

        self.get_success(_get_next_async())
        self.assertEqual(id_gen_3.get_persisted_upto_position(), 6)

        # If we add back the old "first" then we shouldn't see the persisted up
        # to position revert back to 3.
        id_gen_5 = self._create_id_generator("five", writers=["first", "third"])
        self.assertEqual(id_gen_5.get_persisted_upto_position(), 6)
        self.assertEqual(id_gen_5.get_current_token_for_writer("first"), 6)
        self.assertEqual(id_gen_5.get_current_token_for_writer("third"), 6)

    def test_sequence_consistency(self) -> None:
        """Test that we correct the sequence if the table and sequence diverges."""

        # Prefill with some rows
        self._insert_row_with_id("master", 3)

        # Now we add a row *without* updating the stream ID
        def _insert(txn: Cursor) -> None:
            txn.execute("INSERT INTO foobar VALUES (26, 'master')")

        self.get_success(self.db_pool.runInteraction("_insert", _insert))

        # Creating the ID gen should now fix the inconsistency
        id_gen = self._create_id_generator()

        async def _get_next_async() -> None:
            async with id_gen.get_next() as stream_id:
                self.assertEqual(stream_id, 27)

        self.get_success(_get_next_async())

    def test_minimal_local_token(self) -> None:
        self._insert_rows("first", 3)
        first_id_gen = self._create_id_generator("first", writers=["first", "second"])

        self._insert_rows("second", 4)
        second_id_gen = self._create_id_generator("second", writers=["first", "second"])

        self._replicate_all()

        self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 7})
        self.assertEqual(first_id_gen.get_minimal_local_current_token(), 3)

        self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
        self.assertEqual(second_id_gen.get_minimal_local_current_token(), 7)

    def test_current_token_gap(self) -> None:
        """Test that getting the current token for a writer returns the maximal
        token when there are no writes.
        """
        self._insert_rows("first", 3)
        first_id_gen = self._create_id_generator(
            "first", writers=["first", "second", "third"]
        )

        self._insert_rows("second", 4)
        second_id_gen = self._create_id_generator(
            "second", writers=["first", "second", "third"]
        )

        self._replicate_all()

        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
        self.assertEqual(second_id_gen.get_current_token(), 7)

        # Check that the first ID gen advancing causes the second ID gen to
        # advance (as the second ID gen has nothing in flight).

        async def _get_next_async() -> None:
            async with first_id_gen.get_next_mult(2):
                pass

        self.get_success(_get_next_async())
        second_id_gen.advance("first", 9)

        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 9)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 9)
        self.assertEqual(second_id_gen.get_current_token(), 7)

        # Check that the first ID gen advancing doesn't advance the second ID
        # gen when the second ID gen has stuff in flight.
        self.get_success(_get_next_async())

        ctxmgr = second_id_gen.get_next()
        self.get_success(ctxmgr.__aenter__())

        second_id_gen.advance("first", 11)

        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 11)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 9)
        self.assertEqual(second_id_gen.get_current_token(), 7)

        self.get_success(ctxmgr.__aexit__(None, None, None))

        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 11)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 12)
        self.assertEqual(second_id_gen.get_current_token(), 7)


class BackwardsMultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase):
    """Tests MultiWriterIdGenerator that produce *negative* stream IDs."""

    if not USE_POSTGRES_FOR_TESTS:
        skip = "Requires Postgres"

    positive = False

    def test_single_instance(self) -> None:
        """Test that reads and writes from a single process are handled
        correctly.
        """
        id_gen = self._create_id_generator()

        async def _get_next_async() -> None:
            async with id_gen.get_next() as stream_id:
                self._insert_row("master", stream_id)

        self.get_success(_get_next_async())

        self.assertEqual(id_gen.get_positions(), {"master": -1})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), -1)
        self.assertEqual(id_gen.get_persisted_upto_position(), -1)

        async def _get_next_async2() -> None:
            async with id_gen.get_next_mult(3) as stream_ids:
                for stream_id in stream_ids:
                    self._insert_row("master", stream_id)

        self.get_success(_get_next_async2())

        self.assertEqual(id_gen.get_positions(), {"master": -4})
        self.assertEqual(id_gen.get_current_token_for_writer("master"), -4)
        self.assertEqual(id_gen.get_persisted_upto_position(), -4)

        # Test loading from DB by creating a second ID gen
        second_id_gen = self._create_id_generator()

        self.assertEqual(second_id_gen.get_positions(), {"master": -4})
        self.assertEqual(second_id_gen.get_current_token_for_writer("master"), -4)
        self.assertEqual(second_id_gen.get_persisted_upto_position(), -4)

    def test_multiple_instance(self) -> None:
        """Tests that having multiple instances that get advanced over
        federation works corretly.
        """
        id_gen_1 = self._create_id_generator("first", writers=["first", "second"])
        id_gen_2 = self._create_id_generator("second", writers=["first", "second"])

        async def _get_next_async() -> None:
            async with id_gen_1.get_next() as stream_id:
                self._insert_row("first", stream_id)
            self._replicate("first")

        self.get_success(_get_next_async())

        self.assertEqual(id_gen_1.get_positions(), {"first": -1, "second": -1})
        self.assertEqual(id_gen_2.get_positions(), {"first": -1, "second": -1})
        self.assertEqual(id_gen_1.get_persisted_upto_position(), -1)
        self.assertEqual(id_gen_2.get_persisted_upto_position(), -1)

        async def _get_next_async2() -> None:
            async with id_gen_2.get_next() as stream_id:
                self._insert_row("second", stream_id)
            self._replicate("second")

        self.get_success(_get_next_async2())

        self.assertEqual(id_gen_1.get_positions(), {"first": -1, "second": -2})
        self.assertEqual(id_gen_2.get_positions(), {"first": -1, "second": -2})
        self.assertEqual(id_gen_1.get_persisted_upto_position(), -2)
        self.assertEqual(id_gen_2.get_persisted_upto_position(), -2)

    def test_restart_with_own_rows_past_stream_positions(self) -> None:
        """On restart, rows of our own past `stream_positions` move our
        position to the furthest of them, not the nearest.

        The loader once sorted the rows by their stored, negative stream ID
        before negating them, so it visited ours furthest first and left our
        position at the nearest.
        """

        def _insert(txn: LoggingTransaction) -> None:
            # The table and `stream_positions` hold the negative IDs, the
            # sequence the positive ones.
            for stream_id, instance_name in (
                (-1, "second"),
                (-2, "first"),
                (-3, "first"),
            ):
                txn.execute(
                    "INSERT INTO foobar (stream_id, instance_name) VALUES (?, ?)",
                    (stream_id, instance_name),
                )
            txn.execute("SELECT setval('foobar_seq', 3)")
            # `second` is at 1, so the loader reads every row from 1 on. Our
            # rows at 2 and 3 never reached `stream_positions`.
            txn.execute(
                "INSERT INTO stream_positions VALUES ('test_stream', 'second', -1)"
            )

        self.get_success(self.db_pool.runInteraction("_insert", _insert))

        id_gen = self._create_id_generator("first", writers=["first", "second"])

        self.assertEqual(id_gen.get_positions()["first"], -3)
        self.assertEqual(id_gen.get_persisted_upto_position(), -3)


class MultiTableMultiWriterIdGeneratorTestCase(MultiWriterIdGeneratorBase):
    if not USE_POSTGRES_FOR_TESTS:
        skip = "Requires Postgres"

    tables = ["foobar1", "foobar2"]

    def test_load_existing_stream(self) -> None:
        """Test creating ID gens with multiple tables that have rows from after
        the position in `stream_positions` table.
        """
        self._insert_rows("first", 3, table="foobar1")
        first_id_gen = self._create_id_generator("first", writers=["first", "second"])

        self._insert_rows("second", 3, table="foobar2")
        self._insert_rows("second", 1, table="foobar2", update_stream_table=False)
        second_id_gen = self._create_id_generator("second", writers=["first", "second"])

        self._replicate_all()

        self.assertEqual(first_id_gen.get_positions(), {"first": 3, "second": 7})
        self.assertEqual(first_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(first_id_gen.get_current_token_for_writer("second"), 7)
        self.assertEqual(first_id_gen.get_persisted_upto_position(), 7)

        self.assertEqual(second_id_gen.get_positions(), {"first": 3, "second": 7})
        self.assertEqual(second_id_gen.get_current_token_for_writer("first"), 7)
        self.assertEqual(second_id_gen.get_current_token_for_writer("second"), 7)
        self.assertEqual(second_id_gen.get_persisted_upto_position(), 7)


class ShardedTokenHelpersPureTestCase(TestCase):
    """
    Non-database tests for the helpers for reading multi-writer streams.
    """

    def test_bounds_sql_documented_example(self) -> None:
        """
        Tests that the example in the docstring is what we actually generate.
        """
        clause, values = make_multiwriter_sharded_token_bounds_sql(
            Sqlite3Engine({}),
            stream_id_column="se.stream_id",
            instance_name_column="se.instance_name",
            from_token_exclusive=MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 8})
            ),
            to_token_inclusive=MultiWriterStreamToken(
                stream=10, instance_map=immutabledict({"worker2": 14})
            ),
        )
        self.assertEqualNormalisingWhitespace(
            clause,
            """
            (
                ? < se.stream_id
                AND se.stream_id <= ?
                AND NOT (se.instance_name IS ? AND se.stream_id <= ?)
                AND (
                    se.stream_id <= ?
                    OR (se.instance_name IS ? AND se.stream_id <= ?)
                )
            )
            """,
        )
        self.assertEqual(list(values), [5, 14, "worker1", 8, 10, "worker2", 14])

    def test_token_after_partial_read_does_not_go_backwards(self) -> None:
        """
        Tests that advancing the token after a partial read doesn't let it go
        backwards.

        This is relevant because Synapse workers don't always advance their current
        position at the same time.
        """
        # As a scenario: the client has already read up to a baseline position of 60,
        # but this reader worker has only caught up to 10.
        # It can nonetheless see that worker2 has reached 70.
        from_token = MultiWriterStreamToken(stream=60)
        to_token = MultiWriterStreamToken(
            stream=10, instance_map=immutabledict({"worker2": 70})
        )

        resume_token = advance_multiwriter_sharded_token_after_partial_read(
            from_token_exclusive=from_token,
            to_token_inclusive=to_token,
            last_read_stream_id=64,
        )

        # worker2 advances to what we read; everyone else stays where they were.
        self.assertEqual(
            resume_token,
            MultiWriterStreamToken(
                stream=60, instance_map=immutabledict({"worker2": 64})
            ),
        )
        self.assertTrue(
            from_token.is_before_or_eq(resume_token),
            f"Expected {from_token} <= {resume_token}",
        )


class ShardedTokenHelpersDatabaseTestCase(TestCase):
    """Tests for the helpers that read a range of a multi-writer stream.

    These don't need a homeserver: they only exercise the SQL that
    `make_multiwriter_sharded_token_bounds_sql` builds, against a throwaway
    SQLite table.

    NOTE: `ShardedTokenHelpersDatabaseNullTestCase` inherits all these tests
        with a separate set of tables.
    """

    ROWS: list[tuple[int, str | None]] = [
        # (stream_id, instance_name)
        (6, "worker1"),
        (7, "worker1"),
        (8, "worker2"),
        (9, "worker1"),
        (10, "worker3"),
        (11, "worker1"),
        (12, "worker2"),
        (13, "worker3"),
        (14, "worker2"),
    ]
    """
    (stream_id, instance_name) rows to insert into the example stream table.
    """

    SHARDED_TOKEN_RANGES = [
        (
            MultiWriterStreamToken(stream=5),
            MultiWriterStreamToken(stream=14),
        ),
        (
            MultiWriterStreamToken(stream=5),
            MultiWriterStreamToken(
                stream=10, instance_map=immutabledict({"worker2": 14})
            ),
        ),
        (
            MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 9})
            ),
            MultiWriterStreamToken(
                stream=10, instance_map=immutabledict({"worker2": 14})
            ),
        ),
        (
            MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 7})
            ),
            MultiWriterStreamToken(
                stream=8, instance_map=immutabledict({"worker1": 11, "worker3": 13})
            ),
        ),
        (
            MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 9, "worker2": 8})
            ),
            MultiWriterStreamToken(
                stream=10,
                instance_map=immutabledict(
                    {"worker1": 11, "worker2": 14, "worker3": 13}
                ),
            ),
        ),
    ]
    """(from, to) token pairs to exercise the bounds against `ROWS`."""

    def setUp(self) -> None:
        # Just used for detecting SQL language support,
        # doesn't actually connect...
        self.database_engine = Sqlite3Engine({})
        self.conn = sqlite3.connect(":memory:")
        self.addCleanup(self.conn.close)
        self.conn.execute(
            "CREATE TABLE streamtable (stream_id INTEGER PRIMARY KEY, instance_name TEXT)"
        )
        self.conn.executemany("INSERT INTO streamtable VALUES (?, ?)", self.ROWS)

    def _select(
        self,
        from_token: MultiWriterStreamToken,
        to_token: MultiWriterStreamToken,
        limit: int | None = None,
    ) -> list[tuple[int, str]]:
        """
        Helper that selects the rows within the given bounds, in stream order,
        using `make_multiwriter_sharded_token_bounds_sql`.

        Bounds: from_token < ... <= to_token
        """
        clause, values = make_multiwriter_sharded_token_bounds_sql(
            self.database_engine,
            stream_id_column="streamtable.stream_id",
            instance_name_column="streamtable.instance_name",
            from_token_exclusive=from_token,
            to_token_inclusive=to_token,
        )

        limit_clause = ""
        if limit is not None:
            limit_clause = f"LIMIT {limit}"

        return list(
            self.conn.execute(
                f"""
            SELECT stream_id, instance_name
            FROM streamtable
            WHERE {clause}
            ORDER BY stream_id ASC
            {limit_clause}
            """,
                values,
            )
        )

    def _expected(
        self,
        from_token: MultiWriterStreamToken,
        to_token: MultiWriterStreamToken,
    ) -> list[tuple[int, str | None]]:
        """
        Helper that returns the rows that ought to be within the given bounds.
        It uses `MultiWriterStreamToken.is_stream_position_in_range` as the 'ground truth'.

        Bounds: from_token < ... <= to_token
        """

        return [
            (stream_id, instance_name)
            for stream_id, instance_name in self.ROWS
            if MultiWriterStreamToken.is_stream_position_in_range(
                from_token, to_token, instance_name, stream_id
            )
        ]

    def test_bounds_sql_selects_expected_rows(self) -> None:
        """
        Tests that the `make_multiwriter_sharded_token_bounds_sql` selects the correct rows,
        using `MultiWriterStreamToken.is_stream_position_in_range` as the 'ground truth'.
        """
        for from_token, to_token in self.SHARDED_TOKEN_RANGES:
            with self.subTest(from_token=str(from_token), to_token=str(to_token)):
                self.assertEqual(
                    self._select(from_token, to_token),
                    self._expected(from_token, to_token),
                )

    def test_token_after_partial_read(self) -> None:
        """
        Tests that when advancing the token after a partial read, the subsequent read returns all the remaining rows
        but no more (especially not duplicates).
        """
        for from_token, to_token in self.SHARDED_TOKEN_RANGES:
            all_rows = self._expected(from_token, to_token)
            for limit in range(1, len(all_rows) + 1):
                with self.subTest(
                    from_token=str(from_token), to_token=str(to_token), limit=limit
                ):
                    read = self._select(from_token, to_token, limit=limit)
                    resume_token = advance_multiwriter_sharded_token_after_partial_read(
                        from_token_exclusive=from_token,
                        to_token_inclusive=to_token,
                        last_read_stream_id=read[-1][0],
                    )

                    # The resumption point must be somewhere within the range we
                    # were asked to read.
                    self.assertTrue(
                        from_token.is_before_or_eq(resume_token),
                        f"Expected {from_token} <= {resume_token}",
                    )
                    self.assertTrue(
                        resume_token.is_before_or_eq(to_token),
                        f"Expected {resume_token} <= {to_token}",
                    )

                    self.assertEqual(
                        read,
                        self._expected(from_token, resume_token),
                        f"The rows we read between {from_token} < ... <= {resume_token} seem wrong.",
                    )

                    self.assertEqual(
                        read + self._select(resume_token, to_token),
                        all_rows,
                        f"We did a limited read of {limit} rows within {from_token} < ... <= {to_token}, "
                        f"giving us an advanced token {resume_token}. "
                        f"We then read {resume_token} < ... <= {to_token} and expected to get "
                        "all the rows within range (in order, no duplicates), but didn't.",
                    )


class ShardedTokenHelpersDatabaseNullTestCase(ShardedTokenHelpersDatabaseTestCase):
    """
    Variant of `ShardedTokenHelpersDatabaseTestCase` where there are rows with
    `instance_name` being `NULL`, corresponding to rows that existed before the stream was sharded.
    """

    ROWS = [
        # (stream_id, instance_name)
        (4, None),
        (5, None),
        (6, None),
        (7, None),
        (8, "worker2"),
        (9, "worker1"),
        (10, "worker3"),
        (11, "worker1"),
        (12, "worker2"),
        (13, "worker3"),
        (14, "worker2"),
    ]
    """
    (stream_id, instance_name) rows to insert into the example stream table.
    """

    SHARDED_TOKEN_RANGES = [
        (
            MultiWriterStreamToken(stream=5),
            MultiWriterStreamToken(stream=14),
        ),
        (
            MultiWriterStreamToken(stream=5),
            MultiWriterStreamToken(
                stream=10, instance_map=immutabledict({"worker2": 14})
            ),
        ),
        (
            MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 9})
            ),
            MultiWriterStreamToken(
                stream=10, instance_map=immutabledict({"worker2": 14})
            ),
        ),
        (
            MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 7})
            ),
            MultiWriterStreamToken(
                stream=8, instance_map=immutabledict({"worker1": 11, "worker3": 13})
            ),
        ),
        (
            MultiWriterStreamToken(
                stream=5, instance_map=immutabledict({"worker1": 9, "worker2": 8})
            ),
            MultiWriterStreamToken(
                stream=10,
                instance_map=immutabledict(
                    {"worker1": 11, "worker2": 14, "worker3": 13}
                ),
            ),
        ),
    ]
    """(from, to) token pairs to exercise the bounds against `ROWS`."""