matrix-synapse 1.143.0__cp310-abi3-manylinux_2_28_aarch64.whl

synapse/util/caches/dictionary_cache.py

@@ -0,0 +1,346 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2015, 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 enum
+import logging
+import threading
+from typing import (
+    Generic,
+    Iterable,
+    Literal,
+    TypeVar,
+)
+
+import attr
+
+from synapse.util.caches.lrucache import LruCache
+from synapse.util.caches.treecache import TreeCache
+from synapse.util.clock import Clock
+
+logger = logging.getLogger(__name__)
+
+
+# The type of the cache keys.
+KT = TypeVar("KT")
+# The type of the dictionary keys.
+DKT = TypeVar("DKT")
+# The type of the dictionary values.
+DV = TypeVar("DV")
+
+
+@attr.s(slots=True, frozen=True, auto_attribs=True)
+class DictionaryEntry(Generic[DKT, DV]):
+    """Returned when getting an entry from the cache
+
+    If `full` is true then `known_absent` will be the empty set.
+
+    Attributes:
+        full: Whether the cache has the full or dict or just some keys.
+            If not full then not all requested keys will necessarily be present
+            in `value`
+        known_absent: Keys that were looked up in the dict and were not there.
+        value: The full or partial dict value
+    """
+
+    full: bool
+    known_absent: set[DKT]
+    value: dict[DKT, DV]
+
+    def __len__(self) -> int:
+        return len(self.value)
+
+
+class _FullCacheKey(enum.Enum):
+    """The key we use to cache the full dict."""
+
+    KEY = object()
+
+
+class _Sentinel(enum.Enum):
+    # defining a sentinel in this way allows mypy to correctly handle the
+    # type of a dictionary lookup.
+    sentinel = object()
+
+
+class _PerKeyValue(Generic[DV]):
+    """The cached value of a dictionary key. If `value` is the sentinel,
+    indicates that the requested key is known to *not* be in the full dict.
+    """
+
+    __slots__ = ["value"]
+
+    def __init__(self, value: DV | Literal[_Sentinel.sentinel]) -> None:
+        self.value = value
+
+    def __len__(self) -> int:
+        # We add a `__len__` implementation as we use this class in a cache
+        # where the values are variable length.
+        return 1
+
+
+class DictionaryCache(Generic[KT, DKT, DV]):
+    """Caches key -> dictionary lookups, supporting caching partial dicts, i.e.
+    fetching a subset of dictionary keys for a particular key.
+
+    This cache has two levels of key. First there is the "cache key" (of type
+    `KT`), which maps to a dict. The keys to that dict are the "dict key" (of
+    type `DKT`). The overall structure is therefore `KT->DKT->DV`. For
+    example, it might look like:
+
+       {
+           1: { 1: "a", 2: "b" },
+           2: { 1: "c" },
+       }
+
+    It is possible to look up either individual dict keys, or the *complete*
+    dict for a given cache key.
+
+    Each dict item, and the complete dict is treated as a separate LRU
+    entry for the purpose of cache expiry. For example, given:
+        dict_cache.get(1, None)  -> DictionaryEntry({1: "a", 2: "b"})
+        dict_cache.get(1, [1])  -> DictionaryEntry({1: "a"})
+        dict_cache.get(1, [2])  -> DictionaryEntry({2: "b"})
+
+    ... then the cache entry for the complete dict will expire first,
+    followed by the cache entry for the '1' dict key, and finally that
+    for the '2' dict key.
+    """
+
+    def __init__(
+        self, *, name: str, clock: Clock, server_name: str, max_entries: int = 1000
+    ):
+        """
+        Args:
+            name
+            clock: The homeserver `Clock` instance
+            server_name: The homeserver name that this cache is associated with
+                (used to label the metric) (`hs.hostname`).
+            max_entries
+        """
+
+        # We use a single LruCache to store two different types of entries:
+        #   1. Map from (key, dict_key) -> dict value (or sentinel, indicating
+        #      the key doesn't exist in the dict); and
+        #   2. Map from (key, _FullCacheKey.KEY) -> full dict.
+        #
+        # The former is used when explicit keys of the dictionary are looked up,
+        # and the latter when the full dictionary is requested.
+        #
+        # If when explicit keys are requested and not in the cache, we then look
+        # to see if we have the full dict and use that if we do. If found in the
+        # full dict each key is added into the cache.
+        #
+        # This set up allows the `LruCache` to prune the full dict entries if
+        # they haven't been used in a while, even when there have been recent
+        # queries for subsets of the dict.
+        #
+        # Typing:
+        #     * A key of `(KT, DKT)` has a value of `_PerKeyValue`
+        #     * A key of `(KT, _FullCacheKey.KEY)` has a value of `Dict[DKT, DV]`
+        self.cache: LruCache[
+            tuple[KT, DKT | Literal[_FullCacheKey.KEY]],
+            _PerKeyValue | dict[DKT, DV],
+        ] = LruCache(
+            max_size=max_entries,
+            clock=clock,
+            server_name=server_name,
+            cache_name=name,
+            cache_type=TreeCache,
+            size_callback=len,
+        )
+
+        self.name = name
+        self.sequence = 0
+        self.thread: threading.Thread | None = None
+
+    def check_thread(self) -> None:
+        expected_thread = self.thread
+        if expected_thread is None:
+            self.thread = threading.current_thread()
+        else:
+            if expected_thread is not threading.current_thread():
+                raise ValueError(
+                    "Cache objects can only be accessed from the main thread"
+                )
+
+    def get(self, key: KT, dict_keys: Iterable[DKT] | None = None) -> DictionaryEntry:
+        """Fetch an entry out of the cache
+
+        Args:
+            key
+            dict_keys: If given a set of keys then return only those keys
+                that exist in the cache. If None then returns the full dict
+                if it is in the cache.
+
+        Returns:
+            If `dict_keys` is not None then `DictionaryEntry` will contain include
+            the keys that are in the cache.
+
+            If None then will either return the full dict if in the cache, or the
+            empty dict (with `full` set to False) if it isn't.
+        """
+        if dict_keys is None:
+            # The caller wants the full set of dictionary keys for this cache key
+            return self._get_full_dict(key)
+
+        # We are being asked for a subset of keys.
+
+        # First go and check for each requested dict key in the cache, tracking
+        # which we couldn't find.
+        values = {}
+        known_absent = set()
+        missing = []
+        for dict_key in dict_keys:
+            entry = self.cache.get((key, dict_key), _Sentinel.sentinel)
+            if entry is _Sentinel.sentinel:
+                missing.append(dict_key)
+                continue
+
+            assert isinstance(entry, _PerKeyValue)
+
+            if entry.value is _Sentinel.sentinel:
+                known_absent.add(dict_key)
+            else:
+                values[dict_key] = entry.value
+
+        # If we found everything we can return immediately.
+        if not missing:
+            return DictionaryEntry(False, known_absent, values)
+
+        # We are missing some keys, so check if we happen to have the full dict in
+        # the cache.
+        #
+        # We don't update the last access time for this cache fetch, as we
+        # aren't explicitly interested in the full dict and so we don't want
+        # requests for explicit dict keys to keep the full dict in the cache.
+        entry = self.cache.get(
+            (key, _FullCacheKey.KEY),
+            _Sentinel.sentinel,
+            update_last_access=False,
+        )
+        if entry is _Sentinel.sentinel:
+            # Not in the cache, return the subset of keys we found.
+            return DictionaryEntry(False, known_absent, values)
+
+        # We have the full dict!
+        assert isinstance(entry, dict)
+
+        for dict_key in missing:
+            # We explicitly add each dict key to the cache, so that cache hit
+            # rates and LRU times for each key can be tracked separately.
+            value = entry.get(dict_key, _Sentinel.sentinel)
+            self.cache[(key, dict_key)] = _PerKeyValue(value)
+
+            if value is not _Sentinel.sentinel:
+                values[dict_key] = value
+
+        return DictionaryEntry(True, set(), values)
+
+    def _get_full_dict(
+        self,
+        key: KT,
+    ) -> DictionaryEntry:
+        """Fetch the full dict for the given key."""
+
+        # First we check if we have cached the full dict.
+        entry = self.cache.get((key, _FullCacheKey.KEY), _Sentinel.sentinel)
+        if entry is not _Sentinel.sentinel:
+            assert isinstance(entry, dict)
+            return DictionaryEntry(True, set(), entry)
+
+        return DictionaryEntry(False, set(), {})
+
+    def invalidate(self, key: KT) -> None:
+        self.check_thread()
+
+        # Increment the sequence number so that any SELECT statements that
+        # raced with the INSERT don't update the cache (SYN-369)
+        self.sequence += 1
+
+        # We want to drop all information about the dict for the given key, so
+        # we use `del_multi` to delete it all in one go.
+        #
+        # We ignore the type error here: `del_multi` accepts a truncated key
+        # (when the key type is a tuple).
+        self.cache.del_multi((key,))  # type: ignore[arg-type]
+
+    def invalidate_all(self) -> None:
+        self.check_thread()
+        self.sequence += 1
+        self.cache.clear()
+
+    def update(
+        self,
+        sequence: int,
+        key: KT,
+        value: dict[DKT, DV],
+        fetched_keys: Iterable[DKT] | None = None,
+    ) -> None:
+        """Updates the entry in the cache.
+
+        Note: This does *not* invalidate any existing entries for the `key`.
+        In particular, if we add an entry for the cached "full dict" with
+        `fetched_keys=None`, existing entries for individual dict keys are
+        not invalidated. Likewise, adding entries for individual keys does
+        not invalidate any cached value for the full dict.
+
+        In other words: if the underlying data is *changed*, the cache must
+        be explicitly invalidated via `.invalidate()`.
+
+        Args:
+            sequence
+            key
+            value: The value to update the cache with.
+            fetched_keys: All of the dictionary keys which were
+                fetched from the database.
+
+                If None, this is the complete value for key K. Otherwise, it
+                is used to infer a list of keys which we know don't exist in
+                the full dict.
+        """
+        self.check_thread()
+        if self.sequence == sequence:
+            # Only update the cache if the caches sequence number matches the
+            # number that the cache had before the SELECT was started (SYN-369)
+            if fetched_keys is None:
+                self.cache[(key, _FullCacheKey.KEY)] = value
+            else:
+                self._update_subset(key, value, fetched_keys)
+
+    def _update_subset(
+        self, key: KT, value: dict[DKT, DV], fetched_keys: Iterable[DKT]
+    ) -> None:
+        """Add the given dictionary values as explicit keys in the cache.
+
+        Args:
+            key: top-level cache key
+            value: The dictionary with all the values that we should cache
+            fetched_keys: The full set of dict keys that were looked up. Any keys
+                here not in `value` should be marked as "known absent".
+        """
+
+        for dict_key, dict_value in value.items():
+            self.cache[(key, dict_key)] = _PerKeyValue(dict_value)
+
+        for dict_key in fetched_keys:
+            if dict_key in value:
+                continue
+
+            self.cache[(key, dict_key)] = _PerKeyValue(_Sentinel.sentinel)

synapse/util/caches/expiringcache.py

@@ -0,0 +1,249 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2015, 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 collections import OrderedDict
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Iterable,
+    Literal,
+    TypeVar,
+    overload,
+)
+
+import attr
+
+from twisted.internet import defer
+
+from synapse.config import cache as cache_config
+from synapse.util.caches import EvictionReason, register_cache
+from synapse.util.clock import Clock
+
+if TYPE_CHECKING:
+    from synapse.server import HomeServer
+
+logger = logging.getLogger(__name__)
+
+
+SENTINEL: Any = object()
+
+
+T = TypeVar("T")
+KT = TypeVar("KT")
+VT = TypeVar("VT")
+
+
+class ExpiringCache(Generic[KT, VT]):
+    def __init__(
+        self,
+        *,
+        cache_name: str,
+        server_name: str,
+        hs: "HomeServer",
+        clock: Clock,
+        max_len: int = 0,
+        expiry_ms: int = 0,
+        reset_expiry_on_get: bool = False,
+        iterable: bool = False,
+    ):
+        """
+        Args:
+            cache_name: Name of this cache, used for logging.
+            server_name: The homeserver name that this cache is associated
+                with (used to label the metric) (`hs.hostname`).
+            clock
+            max_len: Max size of dict. If the dict grows larger than this
+                then the oldest items get automatically evicted. Default is 0,
+                which indicates there is no max limit.
+            expiry_ms: How long before an item is evicted from the cache
+                in milliseconds. Default is 0, indicating items never get
+                evicted based on time.
+            reset_expiry_on_get: If true, will reset the expiry time for
+                an item on access. Defaults to False.
+            iterable: If true, the size is calculated by summing the
+                sizes of all entries, rather than the number of entries.
+        """
+        self._cache_name = cache_name
+
+        self._original_max_size = max_len
+
+        self._max_size = int(max_len * cache_config.properties.default_factor_size)
+
+        self._clock = clock
+
+        self._expiry_ms = expiry_ms
+        self._reset_expiry_on_get = reset_expiry_on_get
+
+        self._cache: OrderedDict[KT, _CacheEntry[VT]] = OrderedDict()
+
+        self.iterable = iterable
+
+        self.metrics = register_cache(
+            cache_type="expiring",
+            cache_name=cache_name,
+            cache=self,
+            server_name=server_name,
+        )
+
+        if not self._expiry_ms:
+            # Don't bother starting the loop if things never expire
+            return
+
+        def f() -> "defer.Deferred[None]":
+            return hs.run_as_background_process("prune_cache", self._prune_cache)
+
+        self._clock.looping_call(f, self._expiry_ms / 2)
+
+    def __setitem__(self, key: KT, value: VT) -> None:
+        now = self._clock.time_msec()
+        self._cache[key] = _CacheEntry(now, value)
+        self.evict()
+
+    def evict(self) -> None:
+        # Evict if there are now too many items
+        while self._max_size and len(self) > self._max_size:
+            _key, value = self._cache.popitem(last=False)
+            if self.iterable:
+                # type-ignore, here and below: if self.iterable is true, then the value
+                # type VT should be Sized (i.e. have a __len__ method). We don't enforce
+                # this via the type system at present.
+                self.metrics.inc_evictions(EvictionReason.size, len(value.value))  # type: ignore[arg-type]
+            else:
+                self.metrics.inc_evictions(EvictionReason.size)
+
+    def __getitem__(self, key: KT) -> VT:
+        try:
+            entry = self._cache[key]
+            self.metrics.inc_hits()
+        except KeyError:
+            self.metrics.inc_misses()
+            raise
+
+        if self._reset_expiry_on_get:
+            entry.time = self._clock.time_msec()
+
+        return entry.value
+
+    def pop(self, key: KT, default: T = SENTINEL) -> VT | T:
+        """Removes and returns the value with the given key from the cache.
+
+        If the key isn't in the cache then `default` will be returned if
+        specified, otherwise `KeyError` will get raised.
+
+        Identical functionality to `dict.pop(..)`.
+        """
+
+        value = self._cache.pop(key, SENTINEL)
+        # The key was not found.
+        if value is SENTINEL:
+            if default is SENTINEL:
+                raise KeyError(key)
+            return default
+
+        if self.iterable:
+            self.metrics.inc_evictions(EvictionReason.invalidation, len(value.value))
+        else:
+            self.metrics.inc_evictions(EvictionReason.invalidation)
+
+        return value.value
+
+    def __contains__(self, key: KT) -> bool:
+        return key in self._cache
+
+    @overload
+    def get(self, key: KT, default: Literal[None] = None) -> VT | None: ...
+
+    @overload
+    def get(self, key: KT, default: T) -> VT | T: ...
+
+    def get(self, key: KT, default: T | None = None) -> VT | T | None:
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def setdefault(self, key: KT, value: VT) -> VT:
+        try:
+            return self[key]
+        except KeyError:
+            self[key] = value
+            return value
+
+    async def _prune_cache(self) -> None:
+        if not self._expiry_ms:
+            # zero expiry time means don't expire. This should never get called
+            # since we have this check in start too.
+            return
+        begin_length = len(self)
+
+        now = self._clock.time_msec()
+
+        keys_to_delete = set()
+
+        for key, cache_entry in self._cache.items():
+            if now - cache_entry.time > self._expiry_ms:
+                keys_to_delete.add(key)
+
+        for k in keys_to_delete:
+            value = self._cache.pop(k)
+            if self.iterable:
+                self.metrics.inc_evictions(EvictionReason.time, len(value.value))  # type: ignore[arg-type]
+            else:
+                self.metrics.inc_evictions(EvictionReason.time)
+
+        logger.debug(
+            "[%s] _prune_cache before: %d, after len: %d",
+            self._cache_name,
+            begin_length,
+            len(self),
+        )
+
+    def __len__(self) -> int:
+        if self.iterable:
+            g: Iterable[int] = (len(entry.value) for entry in self._cache.values())  # type: ignore[arg-type]
+            return sum(g)
+        else:
+            return len(self._cache)
+
+    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 = int(self._original_max_size * factor)
+        if new_size != self._max_size:
+            self._max_size = new_size
+            self.evict()
+            return True
+        return False
+
+
+@attr.s(slots=True, auto_attribs=True)
+class _CacheEntry(Generic[VT]):
+    time: int
+    value: VT