Reapply "WIP: @cached progress"

This reverts commit b4ed25e.

Author: MadLittleMods

synapse/storage/databases/main/events_worker.py

@@ -286,8 +286,7 @@ def __init__(
             AsyncLruCache(
                 max_size=hs.config.caches.event_cache_size,
                 cache_name="*getEvent*",
-                # TODO
-                # cache_manager=hs.get_cache_manager(),
+                cache_manager=hs.get_cache_manager(),
                 # `extra_index_cb` Returns a tuple as that is the key type
                 extra_index_cb=lambda _, v: (v.event.room_id,),
             )

synapse/util/caches/descriptors.py

@@ -33,6 +33,7 @@
     List,
     Mapping,
     Optional,
+    Protocol,
     Sequence,
     Tuple,
     Type,
@@ -43,13 +44,15 @@
 from weakref import WeakValueDictionary
 
 import attr
+from typing_extensions import Concatenate, ParamSpec
 
 from twisted.internet import defer
 from twisted.python.failure import Failure
 
 from synapse.logging.context import make_deferred_yieldable, preserve_fn
 from synapse.util import unwrapFirstError
 from synapse.util.async_helpers import delay_cancellation
+from synapse.util.caches import CacheManager
 from synapse.util.caches.deferred_cache import DeferredCache
 from synapse.util.caches.lrucache import LruCache
 
@@ -183,6 +186,7 @@ def foo(self, key, cache_context):
 
     Args:
         orig:
+        cache_manager: The cache manager to handle metrics
         max_entries:
         num_args: number of positional arguments (excluding ``self`` and
             ``cache_context``) to use as cache keys. Defaults to all named
@@ -196,11 +200,14 @@ def foo(self, key, cache_context):
         prune_unread_entries: If True, cache entries that haven't been read recently
             will be evicted from the cache in the background. Set to False to opt-out
             of this behaviour.
+        name: Will default to the `__name__` of the `orig` function.
     """
 
     def __init__(
         self,
         orig: Callable[..., Any],
+        *,
+        cache_manager: CacheManager,
         max_entries: int = 1000,
         num_args: Optional[int] = None,
         uncached_args: Optional[Collection[str]] = None,
@@ -217,6 +224,7 @@ def __init__(
             cache_context=cache_context,
             name=name,
         )
+        self.cache_manager = cache_manager
 
         if tree and self.num_args < 2:
             raise RuntimeError(
@@ -233,6 +241,7 @@ def __get__(
     ) -> Callable[..., "defer.Deferred[Any]"]:
         cache: DeferredCache[CacheKey, Any] = DeferredCache(
             name=self.name,
+            cache_manager=self.cache_manager,
             max_entries=self.max_entries,
             tree=self.tree,
             iterable=self.iterable,
@@ -487,10 +496,12 @@ class _CachedFunctionDescriptor:
     iterable: bool
     prune_unread_entries: bool
     name: Optional[str]
+    cache_manager: CacheManager
 
     def __call__(self, orig: F) -> CachedFunction[F]:
         d = DeferredCacheDescriptor(
             orig,
+            cache_manager=self.cache_manager,
             max_entries=self.max_entries,
             num_args=self.num_args,
             uncached_args=self.uncached_args,
@@ -503,6 +514,15 @@ def __call__(self, orig: F) -> CachedFunction[F]:
         return cast(CachedFunction[F], d)
 
 
+P = ParamSpec("P")
+R = TypeVar("R")
+
+
+class HasCacheManager(Protocol):
+    # Used to handle registering the caches
+    cache_manager: CacheManager
+
+
 def cached(
     *,
     max_entries: int = 1000,
@@ -513,17 +533,55 @@ def cached(
     iterable: bool = False,
     prune_unread_entries: bool = True,
     name: Optional[str] = None,
-) -> _CachedFunctionDescriptor:
-    return _CachedFunctionDescriptor(
-        max_entries=max_entries,
-        num_args=num_args,
-        uncached_args=uncached_args,
-        tree=tree,
-        cache_context=cache_context,
-        iterable=iterable,
-        prune_unread_entries=prune_unread_entries,
-        name=name,
-    )
+) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+    """Decorate an async method with a `Measure` context manager.
+
+    The Measure is created using `self.cache_manager`; it should only be used to decorate
+    methods in classes defining an instance-level `clock` attribute.
+
+    Usage:
+
+    @measure_func()
+    async def foo(...):
+        ...
+
+    Which is analogous to:
+
+    async def foo(...):
+        with Measure(...):
+            ...
+
+    """
+
+    def wrapper(
+        func: Callable[Concatenate[HasCacheManager, P], Awaitable[R]],
+    ) -> Callable[P, Awaitable[R]]:
+        # block_name = func.__name__ if name is None else name
+
+        @functools.wraps(func)
+        async def cached_func(
+            self: HasCacheManager, *args: P.args, **kwargs: P.kwargs
+        ) -> R:
+            return _CachedFunctionDescriptor(
+                max_entries=max_entries,
+                num_args=num_args,
+                uncached_args=uncached_args,
+                tree=tree,
+                cache_context=cache_context,
+                iterable=iterable,
+                prune_unread_entries=prune_unread_entries,
+                name=name,
+                # Grab this attribute from the instance
+                cache_manager=self.cache_manager,
+            )
+
+        # There are some shenanigans here, because we're decorating a method but
+        # explicitly making use of the `self` parameter. The key thing here is that the
+        # return type within the return type for `measure_func` itself describes how the
+        # decorated function will be called.
+        return cached_func  # type: ignore[return-value]
+
+    return wrapper  # type: ignore[return-value]
 
 
 @attr.s(auto_attribs=True, slots=True, frozen=True)