matrix-synapse 1.145.0rc2__cp310-abi3-manylinux_2_28_x86_64.whl

synapse/util/caches/response_cache.py

@@ -0,0 +1,323 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 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
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Awaitable,
+    Callable,
+    Generic,
+    Iterable,
+    TypeVar,
+)
+
+import attr
+
+from twisted.internet import defer
+
+from synapse.logging.context import make_deferred_yieldable, run_in_background
+from synapse.logging.opentracing import (
+    active_span,
+    start_active_span,
+    start_active_span_follows_from,
+)
+from synapse.util.async_helpers import AbstractObservableDeferred, ObservableDeferred
+from synapse.util.caches import EvictionReason, register_cache
+from synapse.util.clock import Clock
+from synapse.util.duration import Duration
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    import opentracing
+
+# the type of the key in the cache
+KV = TypeVar("KV")
+
+# the type of the result from the operation
+RV = TypeVar("RV")
+
+
+@attr.s(auto_attribs=True)
+class ResponseCacheContext(Generic[KV]):
+    """Information about a missed ResponseCache hit
+
+    This object can be passed into the callback for additional feedback
+    """
+
+    cache_key: KV
+    """The cache key that caused the cache miss
+
+    This should be considered read-only.
+
+    TODO: in attrs 20.1, make it frozen with an on_setattr.
+    """
+
+    should_cache: bool = True
+    """Whether the result should be cached once the request completes.
+
+    This can be modified by the callback if it decides its result should not be cached.
+    """
+
+
+@attr.s(auto_attribs=True)
+class ResponseCacheEntry:
+    result: AbstractObservableDeferred
+    """The (possibly incomplete) result of the operation.
+
+    Note that we continue to store an ObservableDeferred even after the operation
+    completes (rather than switching to an immediate value), since that makes it
+    easier to cache Failure results.
+    """
+
+    opentracing_span_context: "opentracing.SpanContext | None"
+    """The opentracing span which generated/is generating the result"""
+
+
+class ResponseCache(Generic[KV]):
+    """
+    This caches a deferred response. Until the deferred completes it will be
+    returned from the cache. This means that if the client retries the request
+    while the response is still being computed, that original response will be
+    used rather than trying to compute a new response.
+    """
+
+    def __init__(
+        self,
+        *,
+        clock: Clock,
+        name: str,
+        server_name: str,
+        timeout_ms: float = 0,
+        enable_logging: bool = True,
+    ):
+        """
+        Args:
+            clock
+            name
+            server_name: The homeserver name that this cache is associated
+                with (used to label the metric) (`hs.hostname`).
+            timeout_ms
+            enable_logging
+        """
+        self._result_cache: dict[KV, ResponseCacheEntry] = {}
+
+        self.clock = clock
+        self.timeout = Duration(milliseconds=timeout_ms)
+
+        self._name = name
+        self._metrics = register_cache(
+            cache_type="response_cache",
+            cache_name=name,
+            cache=self,
+            server_name=server_name,
+            resizable=False,
+        )
+        self._enable_logging = enable_logging
+
+    def size(self) -> int:
+        return len(self._result_cache)
+
+    def __len__(self) -> int:
+        return self.size()
+
+    def keys(self) -> Iterable[KV]:
+        """Get the keys currently in the result cache
+
+        Returns both incomplete entries, and (if the timeout on this cache is non-zero),
+        complete entries which are still in the cache.
+
+        Note that the returned iterator is not safe in the face of concurrent execution:
+        behaviour is undefined if `wrap` is called during iteration.
+        """
+        return self._result_cache.keys()
+
+    def _get(self, key: KV) -> ResponseCacheEntry | None:
+        """Look up the given key.
+
+        Args:
+            key: key to get in the cache
+
+        Returns:
+            The entry for this key, if any; else None.
+        """
+        entry = self._result_cache.get(key)
+        if entry is not None:
+            self._metrics.inc_hits()
+            return entry
+        else:
+            self._metrics.inc_misses()
+            return None
+
+    def _set(
+        self,
+        context: ResponseCacheContext[KV],
+        deferred: "defer.Deferred[RV]",
+        opentracing_span_context: "opentracing.SpanContext | None",
+    ) -> ResponseCacheEntry:
+        """Set the entry for the given key to the given deferred.
+
+        *deferred* should run its callbacks in the sentinel logcontext (ie,
+        you should wrap normal synapse deferreds with
+        synapse.logging.context.run_in_background).
+
+        Args:
+            context: Information about the cache miss
+            deferred: The deferred which resolves to the result.
+            opentracing_span_context: An opentracing span wrapping the calculation
+
+        Returns:
+            The cache entry object.
+        """
+        result = ObservableDeferred(deferred, consumeErrors=True)
+        key = context.cache_key
+        entry = ResponseCacheEntry(result, opentracing_span_context)
+        self._result_cache[key] = entry
+
+        def on_complete(r: RV) -> RV:
+            # if this cache has a non-zero timeout, and the callback has not cleared
+            # the should_cache bit, we leave it in the cache for now and schedule
+            # its removal later.
+            if self.timeout and context.should_cache:
+                self.clock.call_later(
+                    self.timeout,
+                    self._entry_timeout,
+                    key,
+                    # We don't need to track these calls since they don't hold any strong
+                    # references which would keep the `HomeServer` in memory after shutdown.
+                    # We don't want to track these because they can get cancelled really
+                    # quickly and thrash the tracking mechanism, ie. during repeated calls
+                    # to /sync.
+                    call_later_cancel_on_shutdown=False,
+                )
+            else:
+                # otherwise, remove the result immediately.
+                self.unset(key)
+            return r
+
+        # make sure we do this *after* adding the entry to result_cache,
+        # in case the result is already complete (in which case flipping the order would
+        # leave us with a stuck entry in the cache).
+        result.addBoth(on_complete)
+        return entry
+
+    def unset(self, key: KV) -> None:
+        """Remove the cached value for this key from the cache, if any.
+
+        Args:
+            key: key used to remove the cached value
+        """
+        self._metrics.inc_evictions(EvictionReason.invalidation)
+        self._result_cache.pop(key, None)
+
+    def _entry_timeout(self, key: KV) -> None:
+        """For the call_later to remove from the cache"""
+        self._metrics.inc_evictions(EvictionReason.time)
+        self._result_cache.pop(key, None)
+
+    async def wrap(
+        self,
+        key: KV,
+        callback: Callable[..., Awaitable[RV]],
+        *args: Any,
+        cache_context: bool = False,
+        **kwargs: Any,
+    ) -> RV:
+        """Wrap together a *get* and *set* call, taking care of logcontexts
+
+        First looks up the key in the cache, and if it is present makes it
+        follow the synapse logcontext rules and returns it.
+
+        Otherwise, makes a call to *callback(*args, **kwargs)*, which should
+        follow the synapse logcontext rules, and adds the result to the cache.
+
+        Example usage:
+
+            async def handle_request(request):
+                # etc
+                return result
+
+            result = await response_cache.wrap(
+                key,
+                handle_request,
+                request,
+            )
+
+        Args:
+            key: key to get/set in the cache
+
+            callback: function to call if the key is not found in
+                the cache
+
+            *args: positional parameters to pass to the callback, if it is used
+
+            cache_context: if set, the callback will be given a `cache_context` kw arg,
+                which will be a ResponseCacheContext object.
+
+            **kwargs: named parameters to pass to the callback, if it is used
+
+        Returns:
+            The result of the callback (from the cache, or otherwise)
+        """
+        entry = self._get(key)
+        if not entry:
+            if self._enable_logging:
+                logger.debug(
+                    "[%s]: no cached result for [%s], calculating new one",
+                    self._name,
+                    key,
+                )
+            context = ResponseCacheContext(cache_key=key)
+            if cache_context:
+                kwargs["cache_context"] = context
+
+            span_context: opentracing.SpanContext | None = None
+
+            async def cb() -> RV:
+                # NB it is important that we do not `await` before setting span_context!
+                nonlocal span_context
+                with start_active_span(f"ResponseCache[{self._name}].calculate"):
+                    span = active_span()
+                    if span:
+                        span_context = span.context
+                    return await callback(*args, **kwargs)
+
+            d = run_in_background(cb)
+            entry = self._set(context, d, span_context)
+            return await make_deferred_yieldable(entry.result.observe())
+
+        result = entry.result.observe()
+        if self._enable_logging:
+            if result.called:
+                logger.info(
+                    "[%s]: using completed cached result for [%s]", self._name, key
+                )
+            else:
+                logger.info(
+                    "[%s]: using incomplete cached result for [%s]", self._name, key
+                )
+
+        span_context = entry.opentracing_span_context
+        with start_active_span_follows_from(
+            f"ResponseCache[{self._name}].wait",
+            contexts=(span_context,) if span_context else (),
+        ):
+            return await make_deferred_yieldable(result)

synapse/util/caches/stream_change_cache.py

@@ -0,0 +1,370 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 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 math
+from typing import Collection, Mapping
+
+import attr
+from sortedcontainers import SortedDict
+
+from synapse.util import caches
+
+logger = logging.getLogger(__name__)
+
+# for now, assume all entities in the cache are strings
+EntityType = str
+
+
+@attr.s(auto_attribs=True, frozen=True, slots=True)
+class AllEntitiesChangedResult:
+    """Return type of `get_all_entities_changed`.
+
+    Callers must check that there was a cache hit, via `result.hit`, before
+    using the entities in `result.entities`.
+
+    This specifically does *not* implement helpers such as `__bool__` to ensure
+    that callers do the correct checks.
+    """
+
+    _entities: list[EntityType] | None
+
+    @property
+    def hit(self) -> bool:
+        return self._entities is not None
+
+    @property
+    def entities(self) -> list[EntityType]:
+        assert self._entities is not None
+        return self._entities
+
+
+class StreamChangeCache:
+    """
+    Keeps track of the stream positions of the latest change in a set of entities.
+
+    The entity will is typically a room ID or user ID, but can be any string.
+
+    Can be queried for whether a specific entity has changed after a stream position
+    or for a list of changed entities after a stream position. See the individual
+    methods for more information.
+
+    Only tracks to a maximum cache size, any position earlier than the earliest
+    known stream position must be treated as unknown.
+    """
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        server_name: str,
+        current_stream_pos: int,
+        max_size: int = 10000,
+        prefilled_cache: Mapping[EntityType, int] | None = None,
+    ) -> None:
+        """
+        Args:
+            name
+            server_name: The homeserver name that this cache is associated with
+                (used to label the metric) (`hs.hostname`).
+            current_stream_pos
+            max_size
+            prefilled_cache
+        """
+
+        self._original_max_size: int = max_size
+        self._max_size = math.floor(max_size)
+
+        # map from stream id to the set of entities which changed at that stream id.
+        self._cache: SortedDict[int, set[EntityType]] = SortedDict()
+        # map from entity to the stream ID of the latest change for that entity.
+        #
+        # Must be kept in sync with _cache.
+        self._entity_to_key: dict[EntityType, int] = {}
+
+        # the earliest stream_pos for which we can reliably answer
+        # get_all_entities_changed. In other words, one less than the earliest
+        # stream_pos for which we know _cache is valid.
+        #
+        self._earliest_known_stream_pos = current_stream_pos
+
+        self.name = name
+        self.metrics = caches.register_cache(
+            cache_type="cache",
+            cache_name=self.name,
+            server_name=server_name,
+            cache=self._cache,
+            resize_callback=self.set_cache_factor,
+        )
+
+        if prefilled_cache:
+            for entity, stream_pos in prefilled_cache.items():
+                self.entity_has_changed(entity, stream_pos)
+
+    def set_cache_factor(self, factor: float) -> bool:
+        """
+        Set the cache factor for this individual cache.
+
+        This will trigger a resize if it changes, which may require evicting
+        items from the cache.
+
+        Returns:
+            Whether the cache changed size or not.
+        """
+        new_size = math.floor(self._original_max_size * factor)
+        if new_size != self._max_size:
+            self._max_size = new_size
+            self._evict()
+            return True
+        return False
+
+    def has_entity_changed(self, entity: EntityType, stream_pos: int) -> bool:
+        """
+        Returns True if the entity may have been updated after stream_pos.
+
+        Args:
+            entity: The entity to check for changes.
+            stream_pos: The stream position to check for changes after.
+
+        Return:
+            True if the entity may have been updated, this happens if:
+                * The given stream position is at or earlier than the earliest
+                  known stream position.
+                * The given stream position is earlier than the latest change for
+                  the entity.
+
+            False otherwise:
+                * The entity is unknown.
+                * The given stream position is at or later than the latest change
+                  for the entity.
+        """
+        assert isinstance(stream_pos, int)
+
+        # _cache is not valid before the earliest known stream position, so
+        # return that the entity has changed.
+        if stream_pos < self._earliest_known_stream_pos:
+            self.metrics.inc_misses()
+            return True
+
+        # If the entity is unknown, it hasn't changed.
+        latest_entity_change_pos = self._entity_to_key.get(entity, None)
+        if latest_entity_change_pos is None:
+            self.metrics.inc_hits()
+            return False
+
+        # This is a known entity, return true if the stream position is earlier
+        # than the last change.
+        if stream_pos < latest_entity_change_pos:
+            self.metrics.inc_misses()
+            return True
+
+        # Otherwise, the stream position is after the latest change: return false.
+        self.metrics.inc_hits()
+        return False
+
+    def get_entities_changed(
+        self, entities: Collection[EntityType], stream_pos: int, _perf_factor: int = 1
+    ) -> set[EntityType] | frozenset[EntityType]:
+        """
+        Returns the subset of the given entities that have had changes after the given position.
+
+        Entities unknown to the cache will be returned.
+
+        If the position is too old it will just return the given list.
+
+        Args:
+            entities: Entities to check for changes.
+            stream_pos: The stream position to check for changes after.
+            _perf_factor: Used by unit tests to choose when to use each
+                optimisation.
+
+        Return:
+            A subset of entities which have changed after the given stream position.
+
+            This will be all entities if the given stream position is at or earlier
+            than the earliest known stream position.
+        """
+        if not self._cache or stream_pos < self._earliest_known_stream_pos:
+            self.metrics.inc_misses()
+            return set(entities)
+
+        # If there have been tonnes of changes compared with the number of
+        # entities, it is faster to check each entities stream ordering
+        # one-by-one.
+        max_stream_pos, _ = self._cache.peekitem()
+        if max_stream_pos - stream_pos > _perf_factor * len(entities):
+            self.metrics.inc_hits()
+            return {
+                entity
+                for entity in entities
+                if self._entity_to_key.get(entity, -1) > stream_pos
+            }
+
+        cache_result = self.get_all_entities_changed(stream_pos)
+        if cache_result.hit:
+            # We now do an intersection, trying to do so in the most efficient
+            # way possible (some of these sets are *large*). First check in the
+            # given iterable is already a set that we can reuse, otherwise we
+            # create a set of the *smallest* of the two iterables and call
+            # `intersection(..)` on it (this can be twice as fast as the reverse).
+            if isinstance(entities, (set, frozenset)):
+                result = entities.intersection(cache_result.entities)
+            elif len(cache_result.entities) < len(entities):
+                result = set(cache_result.entities).intersection(entities)
+            else:
+                result = set(entities).intersection(cache_result.entities)
+            self.metrics.inc_hits()
+        else:
+            result = set(entities)
+            self.metrics.inc_misses()
+
+        return result
+
+    def has_any_entity_changed(self, stream_pos: int) -> bool:
+        """
+        Returns true if any entity has changed after the given stream position.
+
+        Args:
+            stream_pos: The stream position to check for changes after.
+
+        Return:
+            True if any entity has changed after the given stream position or
+            if the given stream position is at or earlier than the earliest
+            known stream position.
+
+            False otherwise.
+        """
+        assert isinstance(stream_pos, int)
+
+        # _cache is not valid before the earliest known stream position, so
+        # return that an entity has changed.
+        if stream_pos < self._earliest_known_stream_pos:
+            self.metrics.inc_misses()
+            return True
+
+        # If the cache is empty, nothing can have changed.
+        if not self._cache:
+            self.metrics.inc_misses()
+            return False
+
+        self.metrics.inc_hits()
+        return stream_pos < self._cache.peekitem()[0]
+
+    def get_all_entities_changed(self, stream_pos: int) -> AllEntitiesChangedResult:
+        """
+        Returns all entities that have had changes after the given position.
+
+        If the stream change cache does not go far enough back, i.e. the
+        position is too old, it will return None.
+
+        Returns the entities in the order that they were changed.
+
+        Args:
+            stream_pos: The stream position to check for changes after.
+
+        Return:
+            A class indicating if we have the requested data cached, and if so
+            includes the entities in the order they were changed.
+        """
+        assert isinstance(stream_pos, int)
+
+        # _cache is not valid before the earliest known stream position, so
+        # return None to mark that it is unknown if an entity has changed.
+        if stream_pos < self._earliest_known_stream_pos:
+            return AllEntitiesChangedResult(None)
+
+        changed_entities: list[EntityType] = []
+
+        for k in self._cache.islice(start=self._cache.bisect_right(stream_pos)):
+            changed_entities.extend(self._cache[k])
+        return AllEntitiesChangedResult(changed_entities)
+
+    def entity_has_changed(self, entity: EntityType, stream_pos: int) -> None:
+        """
+        Informs the cache that the entity has been changed at the given position.
+
+        Args:
+            entity: The entity to mark as changed.
+            stream_pos: The stream position to update the entity to.
+        """
+        assert isinstance(stream_pos, int)
+
+        # For a change before _cache is valid (e.g. at or before the earliest known
+        # stream position) there's nothing to do.
+        if stream_pos <= self._earliest_known_stream_pos:
+            return
+
+        old_pos = self._entity_to_key.get(entity, None)
+        if old_pos is not None:
+            if old_pos >= stream_pos:
+                # nothing to do
+                return
+            e = self._cache[old_pos]
+            e.remove(entity)
+            if not e:
+                # cache at this point is now empty
+                del self._cache[old_pos]
+
+        e1 = self._cache.get(stream_pos)
+        if e1 is None:
+            e1 = self._cache[stream_pos] = set()
+        e1.add(entity)
+        self._entity_to_key[entity] = stream_pos
+        self._evict()
+
+    def all_entities_changed(self, stream_pos: int) -> None:
+        """
+        Mark all entities as changed. This is useful when the cache is invalidated and
+        there may be some potential change for all of the entities.
+        """
+        self._cache.clear()
+        self._entity_to_key.clear()
+        self._earliest_known_stream_pos = stream_pos
+
+    def _evict(self) -> None:
+        """
+        Ensure the cache has not exceeded the maximum size.
+
+        Evicts entries until it is at the maximum size.
+        """
+        # if the cache is too big, remove entries
+        while len(self._cache) > self._max_size:
+            k, r = self._cache.popitem(0)
+            self._earliest_known_stream_pos = max(k, self._earliest_known_stream_pos)
+            for entity in r:
+                self._entity_to_key.pop(entity, None)
+
+    def get_max_pos_of_last_change(self, entity: EntityType) -> int | None:
+        """Returns an upper bound of the stream id of the last change to an
+        entity.
+
+        Args:
+            entity: The entity to check.
+
+        Return:
+            The stream position of the latest change for the given entity, if
+            known
+        """
+        return self._entity_to_key.get(entity)
+
+    def get_earliest_known_position(self) -> int:
+        """Returns the earliest position in the cache."""
+        return self._earliest_known_stream_pos