matrix-synapse 1.143.0__cp310-abi3-manylinux_2_28_aarch64.whl

synapse/storage/util/partial_state_events_tracker.py

@@ -0,0 +1,194 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2022 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 logging
+from collections import defaultdict
+from typing import Collection
+
+from twisted.internet import defer
+from twisted.internet.defer import Deferred
+
+from synapse.logging.context import PreserveLoggingContext, make_deferred_yieldable
+from synapse.logging.opentracing import trace_with_opname
+from synapse.storage.databases.main.events_worker import EventsWorkerStore
+from synapse.storage.databases.main.room import RoomWorkerStore
+from synapse.util import unwrapFirstError
+from synapse.util.cancellation import cancellable
+
+logger = logging.getLogger(__name__)
+
+
+class PartialStateEventsTracker:
+    """Keeps track of which events have partial state, after a partial-state join"""
+
+    def __init__(self, store: EventsWorkerStore):
+        self._store = store
+        # a map from event id to a set of Deferreds which are waiting for that event to be
+        # un-partial-stated.
+        self._observers: dict[str, set[Deferred[None]]] = defaultdict(set)
+
+    def notify_un_partial_stated(self, event_id: str) -> None:
+        """Notify that we now have full state for a given event
+
+        Called by the state-resynchronization loop whenever we resynchronize the state
+        for a particular event. Unblocks any callers to await_full_state() for that
+        event.
+
+        Args:
+            event_id: the event that now has full state.
+        """
+        observers = self._observers.pop(event_id, None)
+        if not observers:
+            return
+        logger.info(
+            "Notifying %i things waiting for un-partial-stating of event %s",
+            len(observers),
+            event_id,
+        )
+        with PreserveLoggingContext():
+            for o in observers:
+                o.callback(None)
+
+    @trace_with_opname("PartialStateEventsTracker.await_full_state")
+    @cancellable
+    async def await_full_state(self, event_ids: Collection[str]) -> None:
+        """Wait for all the given events to have full state.
+
+        Args:
+            event_ids: the list of event ids that we want full state for
+        """
+        # first try the happy path: if there are no partial-state events, we can return
+        # quickly
+        partial_state_event_ids = [
+            ev
+            for ev, p in (await self._store.get_partial_state_events(event_ids)).items()
+            if p
+        ]
+
+        if not partial_state_event_ids:
+            return
+
+        logger.info(
+            "Awaiting un-partial-stating of events %s",
+            partial_state_event_ids,
+            stack_info=True,
+        )
+
+        # create an observer for each lazy-joined event
+        observers: dict[str, Deferred[None]] = {
+            event_id: Deferred() for event_id in partial_state_event_ids
+        }
+        for event_id, observer in observers.items():
+            self._observers[event_id].add(observer)
+
+        try:
+            # some of them may have been un-lazy-joined between us checking the db and
+            # registering the observer, in which case we'd wait forever for the
+            # notification. Call back the observers now.
+            for event_id, partial in (
+                await self._store.get_partial_state_events(observers.keys())
+            ).items():
+                # there may have been a call to notify_un_partial_stated during the
+                # db query, so the observers may already have been called.
+                if not partial and not observers[event_id].called:
+                    observers[event_id].callback(None)
+
+            await make_deferred_yieldable(
+                defer.gatherResults(
+                    observers.values(),
+                    consumeErrors=True,
+                )
+            ).addErrback(unwrapFirstError)
+            logger.info("Events %s all un-partial-stated", observers.keys())
+        finally:
+            # remove any observers we created. This should happen when the notification
+            # is received, but that might not happen for two reasons:
+            #   (a) we're bailing out early on an exception (including us being
+            #       cancelled during the await)
+            #   (b) the event got de-lazy-joined before we set up the observer.
+            for event_id, observer in observers.items():
+                observer_set = self._observers.get(event_id)
+                if observer_set:
+                    observer_set.discard(observer)
+                    if not observer_set:
+                        del self._observers[event_id]
+
+
+class PartialCurrentStateTracker:
+    """Keeps track of which rooms have partial state, after partial-state joins"""
+
+    def __init__(self, store: RoomWorkerStore):
+        self._store = store
+
+        # a map from room id to a set of Deferreds which are waiting for that room to be
+        # un-partial-stated.
+        self._observers: dict[str, set[Deferred[None]]] = defaultdict(set)
+
+    def notify_un_partial_stated(self, room_id: str) -> None:
+        """Notify that we now have full current state for a given room
+
+        Unblocks any callers to await_full_state() for that room.
+
+        Args:
+            room_id: the room that now has full current state.
+        """
+        observers = self._observers.pop(room_id, None)
+        if not observers:
+            return
+        logger.info(
+            "Notifying %i things waiting for un-partial-stating of room %s",
+            len(observers),
+            room_id,
+        )
+        with PreserveLoggingContext():
+            for o in observers:
+                o.callback(None)
+
+    @trace_with_opname("PartialCurrentStateTracker.await_full_state")
+    @cancellable
+    async def await_full_state(self, room_id: str) -> None:
+        # We add the deferred immediately so that the DB call to check for
+        # partial state doesn't race when we unpartial the room.
+        d: Deferred[None] = Deferred()
+        self._observers.setdefault(room_id, set()).add(d)
+
+        try:
+            # Check if the room has partial current state or not.
+            has_partial_state = await self._store.is_partial_state_room(room_id)
+            if not has_partial_state:
+                return
+
+            logger.info(
+                "Awaiting un-partial-stating of room %s",
+                room_id,
+                stack_info=True,
+            )
+
+            await make_deferred_yieldable(d)
+
+            logger.info("Room has un-partial-stated")
+        finally:
+            # Remove the added observer, and remove the room entry if its empty.
+            ds = self._observers.get(room_id)
+            if ds is not None:
+                ds.discard(d)
+                if not ds:
+                    self._observers.pop(room_id, None)

synapse/storage/util/sequence.py

@@ -0,0 +1,315 @@
+#
+# 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 abc
+import logging
+import threading
+from typing import TYPE_CHECKING, Callable
+
+from synapse.storage.engines import (
+    BaseDatabaseEngine,
+    IncorrectDatabaseSetup,
+    PostgresEngine,
+)
+from synapse.storage.types import Connection, Cursor
+
+if TYPE_CHECKING:
+    from synapse.storage.database import LoggingDatabaseConnection
+
+logger = logging.getLogger(__name__)
+
+
+_INCONSISTENT_STREAM_ERROR = """
+Postgres sequence '%(seq)s' is inconsistent with associated stream position
+of '%(stream_name)s' in the 'stream_positions' table.
+
+This is likely a programming error and should be reported at
+https://github.com/matrix-org/synapse.
+
+A temporary workaround to fix this error is to shut down Synapse (including
+any and all workers) and run the following SQL:
+
+    DELETE FROM stream_positions WHERE stream_name = '%(stream_name)s';
+
+This will need to be done every time the server is restarted.
+"""
+
+
+class SequenceGenerator(metaclass=abc.ABCMeta):
+    """A class which generates a unique sequence of integers"""
+
+    @abc.abstractmethod
+    def get_next_id_txn(self, txn: Cursor) -> int:
+        """Gets the next ID in the sequence"""
+        ...
+
+    @abc.abstractmethod
+    def get_next_mult_txn(self, txn: Cursor, n: int) -> list[int]:
+        """Get the next `n` IDs in the sequence"""
+        ...
+
+    @abc.abstractmethod
+    def check_consistency(
+        self,
+        db_conn: "LoggingDatabaseConnection",
+        table: str,
+        id_column: str,
+        stream_name: str | None = None,
+        positive: bool = True,
+    ) -> None:
+        """Should be called during start up to test that the current value of
+        the sequence is greater than or equal to the maximum ID in the table.
+
+        This is to handle various cases where the sequence value can get out of
+        sync with the table, e.g. if Synapse gets rolled back to a previous
+        version and the rolled forwards again.
+
+        If a stream name is given then this will check that any value in the
+        `stream_positions` table is less than or equal to the current sequence
+        value. If it isn't then it's likely that streams have been crossed
+        somewhere (e.g. two ID generators have the same stream name).
+        """
+        ...
+
+    @abc.abstractmethod
+    def get_max_allocated(self, txn: Cursor) -> int:
+        """Get the maximum ID that we have allocated"""
+
+
+class PostgresSequenceGenerator(SequenceGenerator):
+    """An implementation of SequenceGenerator which uses a postgres sequence"""
+
+    def __init__(self, sequence_name: str):
+        self._sequence_name = sequence_name
+
+    def get_next_id_txn(self, txn: Cursor) -> int:
+        txn.execute("SELECT nextval(?)", (self._sequence_name,))
+        fetch_res = txn.fetchone()
+        assert fetch_res is not None
+        return fetch_res[0]
+
+    def get_next_mult_txn(self, txn: Cursor, n: int) -> list[int]:
+        txn.execute(
+            "SELECT nextval(?) FROM generate_series(1, ?)", (self._sequence_name, n)
+        )
+        return [i for (i,) in txn]
+
+    def check_consistency(
+        self,
+        db_conn: "LoggingDatabaseConnection",
+        table: str,
+        id_column: str,
+        stream_name: str | None = None,
+        positive: bool = True,
+    ) -> None:
+        """See SequenceGenerator.check_consistency for docstring."""
+
+        txn = db_conn.cursor(txn_name="sequence.check_consistency")
+
+        # First we get the current max ID from the table.
+        table_sql = "SELECT GREATEST(%(agg)s(%(id)s), 0) FROM %(table)s" % {
+            "id": id_column,
+            "table": table,
+            "agg": "MAX" if positive else "-MIN",
+        }
+
+        txn.execute(table_sql)
+        row = txn.fetchone()
+        if not row:
+            # Table is empty, so nothing to do.
+            txn.close()
+            return
+
+        # Now we fetch the current value from the sequence and compare with the
+        # above.
+        max_stream_id = row[0]
+        txn.execute(
+            "SELECT last_value, is_called FROM %(seq)s" % {"seq": self._sequence_name}
+        )
+        fetch_res = txn.fetchone()
+        assert fetch_res is not None
+        last_value, is_called = fetch_res
+
+        # If we have an associated stream check the stream_positions table.
+        max_in_stream_positions = None
+        if stream_name:
+            txn.execute(
+                "SELECT MAX(stream_id) FROM stream_positions WHERE stream_name = ?",
+                (stream_name,),
+            )
+            row = txn.fetchone()
+            if row:
+                max_in_stream_positions = row[0]
+
+        # If `is_called` is False then `last_value` is actually the value that
+        # will be generated next, so we decrement to get the true "last value".
+        if not is_called:
+            last_value -= 1
+
+        if max_stream_id > last_value:
+            # The sequence is lagging behind the tables. This is probably due to
+            # rolling back to a version before the sequence was used and then
+            # forwards again. We resolve this by setting the sequence to the
+            # right value.
+            logger.warning(
+                "Postgres sequence %s is behind table %s: %d < %d. Updating sequence.",
+                self._sequence_name,
+                table,
+                last_value,
+                max_stream_id,
+            )
+
+            sql = f"""
+                SELECT setval('{self._sequence_name}', GREATEST(
+                    (SELECT last_value FROM {self._sequence_name}),
+                    ({table_sql})
+                ));
+            """
+            txn.execute(sql)
+
+        txn.close()
+
+        # If we have values in the stream positions table then they have to be
+        # less than or equal to `last_value`
+        if max_in_stream_positions and max_in_stream_positions > last_value:
+            raise IncorrectDatabaseSetup(
+                _INCONSISTENT_STREAM_ERROR
+                % {"seq": self._sequence_name, "stream_name": stream_name}
+            )
+
+    def get_max_allocated(self, txn: Cursor) -> int:
+        # We just read from the sequence what the last value we fetched was.
+        txn.execute(f"SELECT last_value, is_called FROM {self._sequence_name}")
+        row = txn.fetchone()
+        assert row is not None
+
+        last_value, is_called = row
+        if not is_called:
+            last_value -= 1
+        return last_value
+
+
+GetFirstCallbackType = Callable[[Cursor], int]
+
+
+class LocalSequenceGenerator(SequenceGenerator):
+    """An implementation of SequenceGenerator which uses local locking
+
+    This only works reliably if there are no other worker processes generating IDs at
+    the same time.
+    """
+
+    def __init__(self, get_first_callback: GetFirstCallbackType):
+        """
+        Args:
+            get_first_callback: a callback which is called on the first call to
+                 get_next_id_txn; should return the current maximum id
+        """
+        # the callback. this is cleared after it is called, so that it can be GCed.
+        self._callback: GetFirstCallbackType | None = get_first_callback
+
+        # The current max value, or None if we haven't looked in the DB yet.
+        self._current_max_id: int | None = None
+        self._lock = threading.Lock()
+
+    def get_next_id_txn(self, txn: Cursor) -> int:
+        # We do application locking here since if we're using sqlite then
+        # we are a single process synapse.
+        with self._lock:
+            if self._current_max_id is None:
+                assert self._callback is not None
+                self._current_max_id = self._callback(txn)
+                self._callback = None
+
+            self._current_max_id += 1
+            return self._current_max_id
+
+    def get_next_mult_txn(self, txn: Cursor, n: int) -> list[int]:
+        with self._lock:
+            if self._current_max_id is None:
+                assert self._callback is not None
+                self._current_max_id = self._callback(txn)
+                self._callback = None
+
+            first_id = self._current_max_id + 1
+            self._current_max_id += n
+            return [first_id + i for i in range(n)]
+
+    def check_consistency(
+        self,
+        db_conn: Connection,
+        table: str,
+        id_column: str,
+        stream_name: str | None = None,
+        positive: bool = True,
+    ) -> None:
+        # There is nothing to do for in memory sequences
+        pass
+
+    def get_max_allocated(self, txn: Cursor) -> int:
+        with self._lock:
+            if self._current_max_id is None:
+                assert self._callback is not None
+                self._current_max_id = self._callback(txn)
+                self._callback = None
+
+            return self._current_max_id
+
+
+def build_sequence_generator(
+    db_conn: "LoggingDatabaseConnection",
+    database_engine: BaseDatabaseEngine,
+    get_first_callback: GetFirstCallbackType,
+    sequence_name: str,
+    table: str | None,
+    id_column: str | None,
+    stream_name: str | None = None,
+    positive: bool = True,
+) -> SequenceGenerator:
+    """Get the best impl of SequenceGenerator available
+
+    This uses PostgresSequenceGenerator on postgres, and a locally-locked impl on
+    sqlite.
+
+    Args:
+        database_engine: the database engine we are connected to
+        get_first_callback: a callback which gets the next sequence ID. Used if
+            we're on sqlite.
+        sequence_name: the name of a postgres sequence to use.
+        table, id_column, stream_name, positive: If set then `check_consistency`
+            is called on the created sequence. See docstring for
+            `check_consistency` details.
+    """
+    if isinstance(database_engine, PostgresEngine):
+        seq: SequenceGenerator = PostgresSequenceGenerator(sequence_name)
+    else:
+        seq = LocalSequenceGenerator(get_first_callback)
+
+    if table:
+        assert id_column
+        seq.check_consistency(
+            db_conn=db_conn,
+            table=table,
+            id_column=id_column,
+            stream_name=stream_name,
+            positive=positive,
+        )
+
+    return seq

synapse/streams/__init__.py

@@ -0,0 +1,43 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2014-2016 OpenMarket Ltd
+# 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]
+#
+#
+from abc import ABC, abstractmethod
+from typing import Generic, TypeVar
+
+from synapse.types import StrCollection, UserID
+
+# The key, this is either a stream token or int.
+K = TypeVar("K")
+# The return type.
+R = TypeVar("R")
+
+
+class EventSource(ABC, Generic[K, R]):
+    @abstractmethod
+    async def get_new_events(
+        self,
+        user: UserID,
+        from_key: K,
+        limit: int,
+        room_ids: StrCollection,
+        is_guest: bool,
+        explicit_room_id: str | None = None,
+    ) -> tuple[list[R], K]:
+        raise NotImplementedError()

synapse/streams/config.py

@@ -0,0 +1,91 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2014-2016 OpenMarket Ltd
+# 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 logging
+
+import attr
+
+from synapse.api.constants import Direction
+from synapse.api.errors import SynapseError
+from synapse.http.servlet import parse_enum, parse_integer, parse_string
+from synapse.http.site import SynapseRequest
+from synapse.storage.databases.main import DataStore
+from synapse.types import StreamToken
+
+logger = logging.getLogger(__name__)
+
+
+MAX_LIMIT = 1000
+
+
+@attr.s(slots=True, auto_attribs=True)
+class PaginationConfig:
+    """A configuration object which stores pagination parameters."""
+
+    from_token: StreamToken | None
+    to_token: StreamToken | None
+    direction: Direction
+    limit: int
+
+    @classmethod
+    async def from_request(
+        cls,
+        store: "DataStore",
+        request: SynapseRequest,
+        default_limit: int,
+        default_dir: Direction = Direction.FORWARDS,
+    ) -> "PaginationConfig":
+        direction = parse_enum(request, "dir", Direction, default=default_dir)
+
+        from_tok_str = parse_string(request, "from")
+        to_tok_str = parse_string(request, "to")
+
+        try:
+            from_tok = None
+            if from_tok_str == "END":
+                from_tok = None  # For backwards compat.
+            elif from_tok_str:
+                from_tok = await StreamToken.from_string(store, from_tok_str)
+        except Exception:
+            raise SynapseError(400, "'from' parameter is invalid")
+
+        try:
+            to_tok = None
+            if to_tok_str:
+                to_tok = await StreamToken.from_string(store, to_tok_str)
+        except Exception:
+            raise SynapseError(400, "'to' parameter is invalid")
+
+        limit = parse_integer(request, "limit", default=default_limit)
+        limit = min(limit, MAX_LIMIT)
+
+        try:
+            return PaginationConfig(from_tok, to_tok, direction, limit)
+        except Exception:
+            logger.exception("Failed to create pagination config")
+            raise SynapseError(400, "Invalid request.")
+
+    def __repr__(self) -> str:
+        return "PaginationConfig(from_tok=%r, to_tok=%r, direction=%r, limit=%r)" % (
+            self.from_token,
+            self.to_token,
+            self.direction,
+            self.limit,
+        )