matrix-synapse 1.142.0rc3__cp314-abi3-musllinux_1_2_aarch64.whl

synapse/util/batching_queue.py

@@ -0,0 +1,202 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2021 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 typing import (
+    TYPE_CHECKING,
+    Awaitable,
+    Callable,
+    Generic,
+    Hashable,
+    TypeVar,
+)
+
+from prometheus_client import Gauge
+
+from twisted.internet import defer
+
+from synapse.logging.context import PreserveLoggingContext, make_deferred_yieldable
+from synapse.metrics import SERVER_NAME_LABEL
+from synapse.util.clock import Clock
+
+if TYPE_CHECKING:
+    from synapse.server import HomeServer
+
+logger = logging.getLogger(__name__)
+
+
+V = TypeVar("V")
+R = TypeVar("R")
+
+number_queued = Gauge(
+    "synapse_util_batching_queue_number_queued",
+    "The number of items waiting in the queue across all keys",
+    labelnames=("name", SERVER_NAME_LABEL),
+)
+
+number_in_flight = Gauge(
+    "synapse_util_batching_queue_number_pending",
+    "The number of items across all keys either being processed or waiting in a queue",
+    labelnames=("name", SERVER_NAME_LABEL),
+)
+
+number_of_keys = Gauge(
+    "synapse_util_batching_queue_number_of_keys",
+    "The number of distinct keys that have items queued",
+    labelnames=("name", SERVER_NAME_LABEL),
+)
+
+
+class BatchingQueue(Generic[V, R]):
+    """A queue that batches up work, calling the provided processing function
+    with all pending work (for a given key).
+
+    The provided processing function will only be called once at a time for each
+    key. It will be called the next reactor tick after `add_to_queue` has been
+    called, and will keep being called until the queue has been drained (for the
+    given key).
+
+    If the processing function raises an exception then the exception is proxied
+    through to the callers waiting on that batch of work.
+
+    Note that the return value of `add_to_queue` will be the return value of the
+    processing function that processed the given item. This means that the
+    returned value will likely include data for other items that were in the
+    batch.
+
+    Args:
+        name: A name for the queue, used for logging contexts and metrics.
+            This must be unique, otherwise the metrics will be wrong.
+        server_name: The homeserver name of the server (used to label metrics)
+            (this should be `hs.hostname`).
+        clock: The clock to use to schedule work.
+        process_batch_callback: The callback to to be run to process a batch of
+            work.
+    """
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        hs: "HomeServer",
+        clock: Clock,
+        process_batch_callback: Callable[[list[V]], Awaitable[R]],
+    ):
+        self._name = name
+        self.hs = hs
+        self.server_name = hs.hostname
+        self._clock = clock
+
+        # The set of keys currently being processed.
+        self._processing_keys: set[Hashable] = set()
+
+        # The currently pending batch of values by key, with a Deferred to call
+        # with the result of the corresponding `_process_batch_callback` call.
+        self._next_values: dict[Hashable, list[tuple[V, defer.Deferred]]] = {}
+
+        # The function to call with batches of values.
+        self._process_batch_callback = process_batch_callback
+
+        number_queued.labels(
+            name=self._name, **{SERVER_NAME_LABEL: self.server_name}
+        ).set_function(lambda: sum(len(q) for q in self._next_values.values()))
+
+        number_of_keys.labels(
+            name=self._name, **{SERVER_NAME_LABEL: self.server_name}
+        ).set_function(lambda: len(self._next_values))
+
+        self._number_in_flight_metric: Gauge = number_in_flight.labels(
+            name=self._name, **{SERVER_NAME_LABEL: self.server_name}
+        )
+
+    def shutdown(self) -> None:
+        """
+        Prepares the object for garbage collection by removing any handed out
+        references.
+        """
+        number_queued.remove(self._name, self.server_name)
+        number_of_keys.remove(self._name, self.server_name)
+
+    async def add_to_queue(self, value: V, key: Hashable = ()) -> R:
+        """Adds the value to the queue with the given key, returning the result
+        of the processing function for the batch that included the given value.
+
+        The optional `key` argument allows sharding the queue by some key. The
+        queues will then be processed in parallel, i.e. the process batch
+        function will be called in parallel with batched values from a single
+        key.
+        """
+
+        # First we create a defer and add it and the value to the list of
+        # pending items.
+        d: defer.Deferred[R] = defer.Deferred()
+        self._next_values.setdefault(key, []).append((value, d))
+
+        # If we're not currently processing the key fire off a background
+        # process to start processing.
+        if key not in self._processing_keys:
+            self.hs.run_as_background_process(self._name, self._process_queue, key)
+
+        with self._number_in_flight_metric.track_inprogress():
+            return await make_deferred_yieldable(d)
+
+    async def _process_queue(self, key: Hashable) -> None:
+        """A background task to repeatedly pull things off the queue for the
+        given key and call the `self._process_batch_callback` with the values.
+        """
+
+        if key in self._processing_keys:
+            return
+
+        try:
+            self._processing_keys.add(key)
+
+            while True:
+                # We purposefully wait a reactor tick to allow us to batch
+                # together requests that we're about to receive. A common
+                # pattern is to call `add_to_queue` multiple times at once, and
+                # deferring to the next reactor tick allows us to batch all of
+                # those up.
+                await self._clock.sleep(0)
+
+                next_values = self._next_values.pop(key, [])
+                if not next_values:
+                    # We've exhausted the queue.
+                    break
+
+                try:
+                    values = [value for value, _ in next_values]
+                    results = await self._process_batch_callback(values)
+
+                    with PreserveLoggingContext():
+                        for _, deferred in next_values:
+                            deferred.callback(results)
+
+                except Exception as e:
+                    with PreserveLoggingContext():
+                        for _, deferred in next_values:
+                            if deferred.called:
+                                continue
+
+                            deferred.errback(e)
+
+        finally:
+            self._processing_keys.discard(key)

synapse/util/caches/__init__.py

@@ -0,0 +1,300 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2019, 2020 The Matrix.org Foundation C.I.C.
+# 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 collections
+import logging
+import typing
+from enum import Enum, auto
+from sys import intern
+from typing import Any, Callable, Optional, Sized, TypeVar
+
+import attr
+from prometheus_client import REGISTRY
+from prometheus_client.core import Gauge
+
+from synapse.config.cache import add_resizable_cache
+from synapse.metrics import SERVER_NAME_LABEL
+from synapse.util.metrics import DynamicCollectorRegistry
+
+logger = logging.getLogger(__name__)
+
+
+# Whether to track estimated memory usage of the LruCaches.
+TRACK_MEMORY_USAGE = False
+
+# We track cache metrics in a special registry that lets us update the metrics
+# just before they are returned from the scrape endpoint.
+#
+# The `SERVER_NAME_LABEL` is included in the individual metrics added to this registry,
+# so we don't need to worry about it on the collector itself.
+CACHE_METRIC_REGISTRY = DynamicCollectorRegistry()  # type: ignore[missing-server-name-label]
+
+cache_size = Gauge(
+    "synapse_util_caches_cache_size",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+cache_hits = Gauge(
+    "synapse_util_caches_cache_hits",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+cache_evicted = Gauge(
+    "synapse_util_caches_cache_evicted_size",
+    "",
+    labelnames=["name", "reason", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+cache_total = Gauge(
+    "synapse_util_caches_cache",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+cache_max_size = Gauge(
+    "synapse_util_caches_cache_max_size",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+cache_memory_usage = Gauge(
+    "synapse_util_caches_cache_size_bytes",
+    "Estimated memory usage of the caches",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+
+response_cache_size = Gauge(
+    "synapse_util_caches_response_cache_size",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+response_cache_hits = Gauge(
+    "synapse_util_caches_response_cache_hits",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+response_cache_evicted = Gauge(
+    "synapse_util_caches_response_cache_evicted_size",
+    "",
+    labelnames=["name", "reason", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+response_cache_total = Gauge(
+    "synapse_util_caches_response_cache",
+    "",
+    labelnames=["name", SERVER_NAME_LABEL],
+    registry=CACHE_METRIC_REGISTRY,
+)
+
+
+# Register our custom cache metrics registry with the global registry
+REGISTRY.register(CACHE_METRIC_REGISTRY)
+
+
+class EvictionReason(Enum):
+    size = auto()
+    time = auto()
+    invalidation = auto()
+
+
+@attr.s(slots=True, auto_attribs=True, kw_only=True)
+class CacheMetric:
+    """
+    Used to track cache metrics
+    """
+
+    _cache: Sized
+    _cache_type: str
+    _cache_name: str
+    _collect_callback: Optional[Callable]
+    _server_name: str
+
+    hits: int = 0
+    misses: int = 0
+    eviction_size_by_reason: typing.Counter[EvictionReason] = attr.ib(
+        factory=collections.Counter
+    )
+    memory_usage: Optional[int] = None
+
+    def inc_hits(self) -> None:
+        self.hits += 1
+
+    def inc_misses(self) -> None:
+        self.misses += 1
+
+    def inc_evictions(self, reason: EvictionReason, size: int = 1) -> None:
+        self.eviction_size_by_reason[reason] += size
+
+    def inc_memory_usage(self, memory: int) -> None:
+        if self.memory_usage is None:
+            self.memory_usage = 0
+
+        self.memory_usage += memory
+
+    def dec_memory_usage(self, memory: int) -> None:
+        assert self.memory_usage is not None
+        self.memory_usage -= memory
+
+    def clear_memory_usage(self) -> None:
+        if self.memory_usage is not None:
+            self.memory_usage = 0
+
+    def describe(self) -> list[str]:
+        return []
+
+    def collect(self) -> None:
+        try:
+            labels_base = {
+                "name": self._cache_name,
+                SERVER_NAME_LABEL: self._server_name,
+            }
+            if self._cache_type == "response_cache":
+                response_cache_size.labels(**labels_base).set(len(self._cache))
+                response_cache_hits.labels(**labels_base).set(self.hits)
+                for reason in EvictionReason:
+                    response_cache_evicted.labels(
+                        **{**labels_base, "reason": reason.name}
+                    ).set(self.eviction_size_by_reason[reason])
+                response_cache_total.labels(**labels_base).set(self.hits + self.misses)
+            else:
+                cache_size.labels(**labels_base).set(len(self._cache))
+                cache_hits.labels(**labels_base).set(self.hits)
+                for reason in EvictionReason:
+                    cache_evicted.labels(**{**labels_base, "reason": reason.name}).set(
+                        self.eviction_size_by_reason[reason]
+                    )
+                cache_total.labels(**labels_base).set(self.hits + self.misses)
+                max_size = getattr(self._cache, "max_size", None)
+                if max_size:
+                    cache_max_size.labels(**labels_base).set(max_size)
+
+                if TRACK_MEMORY_USAGE:
+                    # self.memory_usage can be None if nothing has been inserted
+                    # into the cache yet.
+                    cache_memory_usage.labels(**labels_base).set(self.memory_usage or 0)
+            if self._collect_callback:
+                self._collect_callback()
+        except Exception as e:
+            logger.warning("Error calculating metrics for %s: %s", self._cache_name, e)
+            raise
+
+
+def register_cache(
+    *,
+    cache_type: str,
+    cache_name: str,
+    cache: Sized,
+    server_name: str,
+    collect_callback: Optional[Callable] = None,
+    resizable: bool = True,
+    resize_callback: Optional[Callable] = None,
+) -> CacheMetric:
+    """Register a cache object for metric collection and resizing.
+
+    Args:
+        cache_type: a string indicating the "type" of the cache. This is used
+            only for deduplication so isn't too important provided it's constant.
+        cache_name: name of the cache
+        cache: cache itself, which must implement __len__(), and may optionally implement
+             a max_size property
+        server_name: The homeserver name that this cache is associated with
+            (used to label the metric) (`hs.hostname`).
+        collect_callback: If given, a function which is called during metric
+            collection to update additional metrics.
+        resizable: Whether this cache supports being resized, in which case either
+            resize_callback must be provided, or the cache must support set_max_size().
+        resize_callback: A function which can be called to resize the cache.
+
+    Returns:
+        an object which provides inc_{hits,misses,evictions} methods
+    """
+    if resizable:
+        if not resize_callback:
+            resize_callback = cache.set_cache_factor  # type: ignore
+        add_resizable_cache(cache_name, resize_callback)
+
+    metric = CacheMetric(
+        cache=cache,
+        cache_type=cache_type,
+        cache_name=cache_name,
+        server_name=server_name,
+        collect_callback=collect_callback,
+    )
+    metric_name = "cache_%s_%s_%s" % (cache_type, cache_name, server_name)
+    CACHE_METRIC_REGISTRY.register_hook(server_name, metric_name, metric.collect)
+    return metric
+
+
+KNOWN_KEYS = {
+    key: key
+    for key in (
+        "auth_events",
+        "content",
+        "depth",
+        "event_id",
+        "hashes",
+        "origin",  # old events were created with an origin field.
+        "origin_server_ts",
+        "prev_events",
+        "room_id",
+        "sender",
+        "signatures",
+        "state_key",
+        "type",
+        "unsigned",
+        "user_id",
+    )
+}
+
+T = TypeVar("T", Optional[str], str)
+
+
+def intern_string(string: T) -> T:
+    """Takes a (potentially) unicode string and interns it if it's ascii"""
+    if string is None:
+        return None
+
+    try:
+        return intern(string)
+    except UnicodeEncodeError:
+        return string
+
+
+def intern_dict(dictionary: dict[str, Any]) -> dict[str, Any]:
+    """Takes a dictionary and interns well known keys and their values"""
+    return {
+        KNOWN_KEYS.get(key, key): _intern_known_values(key, value)
+        for key, value in dictionary.items()
+    }
+
+
+def _intern_known_values(key: str, value: Any) -> Any:
+    intern_keys = ("event_id", "room_id", "sender", "user_id", "type", "state_key")
+
+    if key in intern_keys:
+        return intern_string(value)
+
+    return value

synapse/util/caches/cached_call.py

@@ -0,0 +1,143 @@
+#
+# This file is licensed under the Affero General Public License (AGPL) version 3.
+#
+# Copyright 2021 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 enum
+from typing import Awaitable, Callable, Generic, Optional, TypeVar, Union
+
+from twisted.internet.defer import Deferred
+from twisted.python.failure import Failure
+
+from synapse.logging.context import make_deferred_yieldable, run_in_background
+
+TV = TypeVar("TV")
+
+
+class _Sentinel(enum.Enum):
+    sentinel = object()
+
+
+class CachedCall(Generic[TV]):
+    """A wrapper for asynchronous calls whose results should be shared
+
+    This is useful for wrapping asynchronous functions, where there might be multiple
+    callers, but we only want to call the underlying function once (and have the result
+    returned to all callers).
+
+    Similar results can be achieved via a lock of some form, but that typically requires
+    more boilerplate (and ends up being less efficient).
+
+    Correctly handles Synapse logcontexts (logs and resource usage for the underlying
+    function are logged against the logcontext which is active when get() is first
+    called).
+
+    Example usage:
+
+        _cached_val = CachedCall(_load_prop)
+
+        async def handle_request() -> X:
+            # We can call this multiple times, but it will result in a single call to
+            # _load_prop().
+            return await _cached_val.get()
+
+        async def _load_prop() -> X:
+            await difficult_operation()
+
+
+    The implementation is deliberately single-shot (ie, once the call is initiated,
+    there is no way to ask for it to be run). This keeps the implementation and
+    semantics simple. If you want to make a new call, simply replace the whole
+    CachedCall object.
+    """
+
+    __slots__ = ["_callable", "_deferred", "_result"]
+
+    def __init__(self, f: Callable[[], Awaitable[TV]]):
+        """
+        Args:
+            f: The underlying function. Only one call to this function will be alive
+                at once (per instance of CachedCall)
+        """
+        self._callable: Optional[Callable[[], Awaitable[TV]]] = f
+        self._deferred: Optional[Deferred] = None
+        self._result: Union[_Sentinel, TV, Failure] = _Sentinel.sentinel
+
+    async def get(self) -> TV:
+        """Kick off the call if necessary, and return the result"""
+
+        # Fire off the callable now if this is our first time
+        if not self._deferred:
+            assert self._callable is not None
+            self._deferred = run_in_background(self._callable)
+
+            # we will never need the callable again, so make sure it can be GCed
+            self._callable = None
+
+            # once the deferred completes, store the result. We cannot simply leave the
+            # result in the deferred, since `awaiting` a deferred destroys its result.
+            # (Also, if it's a Failure, GCing the deferred would log a critical error
+            # about unhandled Failures)
+            def got_result(r: Union[TV, Failure]) -> None:
+                self._result = r
+
+            self._deferred.addBoth(got_result)
+
+        # TODO: consider cancellation semantics. Currently, if the call to get()
+        #    is cancelled, the underlying call will continue (and any future calls
+        #    will get the result/exception), which I think is *probably* ok, modulo
+        #    the fact the underlying call may be logged to a cancelled logcontext,
+        #    and any eventual exception may not be reported.
+
+        # we can now await the deferred, and once it completes, return the result.
+        if isinstance(self._result, _Sentinel):
+            await make_deferred_yieldable(self._deferred)
+            assert not isinstance(self._result, _Sentinel)
+
+        if isinstance(self._result, Failure):
+            self._result.raiseException()
+            raise AssertionError("unexpected return from Failure.raiseException")
+
+        return self._result
+
+
+class RetryOnExceptionCachedCall(Generic[TV]):
+    """A wrapper around CachedCall which will retry the call if an exception is thrown
+
+    This is used in much the same way as CachedCall, but adds some extra functionality
+    so that if the underlying function throws an exception, then the next call to get()
+    will initiate another call to the underlying function. (Any calls to get() which
+    are already pending will raise the exception.)
+    """
+
+    slots = ["_cachedcall"]
+
+    def __init__(self, f: Callable[[], Awaitable[TV]]):
+        async def _wrapper() -> TV:
+            try:
+                return await f()
+            except Exception:
+                # the call raised an exception: replace the underlying CachedCall to
+                # trigger another call next time get() is called
+                self._cachedcall = CachedCall(_wrapper)
+                raise
+
+        self._cachedcall = CachedCall(_wrapper)
+
+    async def get(self) -> TV:
+        return await self._cachedcall.get()