123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515 |
- # Copyright 2015, 2016 OpenMarket Ltd
- # Copyright 2018 New Vector Ltd
- # Copyright 2020 The Matrix.org Foundation C.I.C.
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- import abc
- import enum
- import threading
- from typing import (
- Callable,
- Collection,
- Dict,
- Generic,
- MutableMapping,
- Optional,
- Set,
- Sized,
- Tuple,
- TypeVar,
- Union,
- cast,
- )
- from prometheus_client import Gauge
- from twisted.internet import defer
- from twisted.python.failure import Failure
- from synapse.util.async_helpers import ObservableDeferred
- from synapse.util.caches.lrucache import LruCache
- from synapse.util.caches.treecache import TreeCache, iterate_tree_cache_entry
- cache_pending_metric = Gauge(
- "synapse_util_caches_cache_pending",
- "Number of lookups currently pending for this cache",
- ["name"],
- )
- T = TypeVar("T")
- KT = TypeVar("KT")
- VT = TypeVar("VT")
- class _Sentinel(enum.Enum):
- # defining a sentinel in this way allows mypy to correctly handle the
- # type of a dictionary lookup.
- sentinel = object()
- class DeferredCache(Generic[KT, VT]):
- """Wraps an LruCache, adding support for Deferred results.
- It expects that each entry added with set() will be a Deferred; likewise get()
- will return a Deferred.
- """
- __slots__ = (
- "cache",
- "thread",
- "_pending_deferred_cache",
- )
- def __init__(
- self,
- name: str,
- max_entries: int = 1000,
- tree: bool = False,
- iterable: bool = False,
- apply_cache_factor_from_config: bool = True,
- prune_unread_entries: bool = True,
- ):
- """
- Args:
- name: The name of the cache
- max_entries: Maximum amount of entries that the cache will hold
- tree: Use a TreeCache instead of a dict as the underlying cache type
- iterable: If True, count each item in the cached object as an entry,
- rather than each cached object
- apply_cache_factor_from_config: Whether cache factors specified in the
- config file affect `max_entries`
- 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.
- """
- cache_type = TreeCache if tree else dict
- # _pending_deferred_cache maps from the key value to a `CacheEntry` object.
- self._pending_deferred_cache: Union[
- TreeCache, "MutableMapping[KT, CacheEntry[KT, VT]]"
- ] = cache_type()
- def metrics_cb() -> None:
- cache_pending_metric.labels(name).set(len(self._pending_deferred_cache))
- # cache is used for completed results and maps to the result itself, rather than
- # a Deferred.
- self.cache: LruCache[KT, VT] = LruCache(
- max_size=max_entries,
- cache_name=name,
- cache_type=cache_type,
- size_callback=(
- (lambda d: len(cast(Sized, d)) or 1)
- # Argument 1 to "len" has incompatible type "VT"; expected "Sized"
- # We trust that `VT` is `Sized` when `iterable` is `True`
- if iterable
- else None
- ),
- metrics_collection_callback=metrics_cb,
- apply_cache_factor_from_config=apply_cache_factor_from_config,
- prune_unread_entries=prune_unread_entries,
- )
- self.thread: Optional[threading.Thread] = None
- @property
- def max_entries(self) -> int:
- return self.cache.max_size
- 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,
- callback: Optional[Callable[[], None]] = None,
- update_metrics: bool = True,
- ) -> defer.Deferred:
- """Looks the key up in the caches.
- For symmetry with set(), this method does *not* follow the synapse logcontext
- rules: the logcontext will not be cleared on return, and the Deferred will run
- its callbacks in the sentinel context. In other words: wrap the result with
- make_deferred_yieldable() before `await`ing it.
- Args:
- key:
- callback: Gets called when the entry in the cache is invalidated
- update_metrics: whether to update the cache hit rate metrics
- Returns:
- A Deferred which completes with the result. Note that this may later fail
- if there is an ongoing set() operation which later completes with a failure.
- Raises:
- KeyError if the key is not found in the cache
- """
- val = self._pending_deferred_cache.get(key, _Sentinel.sentinel)
- if val is not _Sentinel.sentinel:
- val.add_invalidation_callback(key, callback)
- if update_metrics:
- m = self.cache.metrics
- assert m # we always have a name, so should always have metrics
- m.inc_hits()
- return val.deferred(key)
- callbacks = (callback,) if callback else ()
- val2 = self.cache.get(
- key, _Sentinel.sentinel, callbacks=callbacks, update_metrics=update_metrics
- )
- if val2 is _Sentinel.sentinel:
- raise KeyError()
- else:
- return defer.succeed(val2)
- def get_bulk(
- self,
- keys: Collection[KT],
- callback: Optional[Callable[[], None]] = None,
- ) -> Tuple[Dict[KT, VT], Optional["defer.Deferred[Dict[KT, VT]]"], Collection[KT]]:
- """Bulk lookup of items in the cache.
- Returns:
- A 3-tuple of:
- 1. a dict of key/value of items already cached;
- 2. a deferred that resolves to a dict of key/value of items
- we're already fetching; and
- 3. a collection of keys that don't appear in the previous two.
- """
- # The cached results
- cached = {}
- # List of pending deferreds
- pending = []
- # Dict that gets filled out when the pending deferreds complete
- pending_results = {}
- # List of keys that aren't in either cache
- missing = []
- callbacks = (callback,) if callback else ()
- for key in keys:
- # Check if its in the main cache.
- immediate_value = self.cache.get(
- key,
- _Sentinel.sentinel,
- callbacks=callbacks,
- )
- if immediate_value is not _Sentinel.sentinel:
- cached[key] = immediate_value
- continue
- # Check if its in the pending cache
- pending_value = self._pending_deferred_cache.get(key, _Sentinel.sentinel)
- if pending_value is not _Sentinel.sentinel:
- pending_value.add_invalidation_callback(key, callback)
- def completed_cb(value: VT, key: KT) -> VT:
- pending_results[key] = value
- return value
- # Add a callback to fill out `pending_results` when that completes
- d = pending_value.deferred(key).addCallback(completed_cb, key)
- pending.append(d)
- continue
- # Not in either cache
- missing.append(key)
- # If we've got pending deferreds, squash them into a single one that
- # returns `pending_results`.
- pending_deferred = None
- if pending:
- pending_deferred = defer.gatherResults(
- pending, consumeErrors=True
- ).addCallback(lambda _: pending_results)
- return (cached, pending_deferred, missing)
- def get_immediate(
- self, key: KT, default: T, update_metrics: bool = True
- ) -> Union[VT, T]:
- """If we have a *completed* cached value, return it."""
- return self.cache.get(key, default, update_metrics=update_metrics)
- def set(
- self,
- key: KT,
- value: "defer.Deferred[VT]",
- callback: Optional[Callable[[], None]] = None,
- ) -> defer.Deferred:
- """Adds a new entry to the cache (or updates an existing one).
- The given `value` *must* be a Deferred.
- First any existing entry for the same key is invalidated. Then a new entry
- is added to the cache for the given key.
- Until the `value` completes, calls to `get()` for the key will also result in an
- incomplete Deferred, which will ultimately complete with the same result as
- `value`.
- If `value` completes successfully, subsequent calls to `get()` will then return
- a completed deferred with the same result. If it *fails*, the cache is
- invalidated and subequent calls to `get()` will raise a KeyError.
- If another call to `set()` happens before `value` completes, then (a) any
- invalidation callbacks registered in the interim will be called, (b) any
- `get()`s in the interim will continue to complete with the result from the
- *original* `value`, (c) any future calls to `get()` will complete with the
- result from the *new* `value`.
- It is expected that `value` does *not* follow the synapse logcontext rules - ie,
- if it is incomplete, it runs its callbacks in the sentinel context.
- Args:
- key: Key to be set
- value: a deferred which will complete with a result to add to the cache
- callback: An optional callback to be called when the entry is invalidated
- """
- self.check_thread()
- self._pending_deferred_cache.pop(key, None)
- # XXX: why don't we invalidate the entry in `self.cache` yet?
- # otherwise, we'll add an entry to the _pending_deferred_cache for now,
- # and add callbacks to add it to the cache properly later.
- entry = CacheEntrySingle[KT, VT](value)
- entry.add_invalidation_callback(key, callback)
- self._pending_deferred_cache[key] = entry
- deferred = entry.deferred(key).addCallbacks(
- self._completed_callback,
- self._error_callback,
- callbackArgs=(entry, key),
- errbackArgs=(entry, key),
- )
- # we return a new Deferred which will be called before any subsequent observers.
- return deferred
- def start_bulk_input(
- self,
- keys: Collection[KT],
- callback: Optional[Callable[[], None]] = None,
- ) -> "CacheMultipleEntries[KT, VT]":
- """Bulk set API for use when fetching multiple keys at once from the DB.
- Called *before* starting the fetch from the DB, and the caller *must*
- call either `complete_bulk(..)` or `error_bulk(..)` on the return value.
- """
- entry = CacheMultipleEntries[KT, VT]()
- entry.add_global_invalidation_callback(callback)
- for key in keys:
- self._pending_deferred_cache[key] = entry
- return entry
- def _completed_callback(
- self, value: VT, entry: "CacheEntry[KT, VT]", key: KT
- ) -> VT:
- """Called when a deferred is completed."""
- # We check if the current entry matches the entry associated with the
- # deferred. If they don't match then it got invalidated.
- current_entry = self._pending_deferred_cache.pop(key, None)
- if current_entry is not entry:
- if current_entry:
- self._pending_deferred_cache[key] = current_entry
- return value
- self.cache.set(key, value, entry.get_invalidation_callbacks(key))
- return value
- def _error_callback(
- self,
- failure: Failure,
- entry: "CacheEntry[KT, VT]",
- key: KT,
- ) -> Failure:
- """Called when a deferred errors."""
- # We check if the current entry matches the entry associated with the
- # deferred. If they don't match then it got invalidated.
- current_entry = self._pending_deferred_cache.pop(key, None)
- if current_entry is not entry:
- if current_entry:
- self._pending_deferred_cache[key] = current_entry
- return failure
- for cb in entry.get_invalidation_callbacks(key):
- cb()
- return failure
- def prefill(
- self, key: KT, value: VT, callback: Optional[Callable[[], None]] = None
- ) -> None:
- callbacks = (callback,) if callback else ()
- self.cache.set(key, value, callbacks=callbacks)
- self._pending_deferred_cache.pop(key, None)
- def invalidate(self, key: KT) -> None:
- """Delete a key, or tree of entries
- If the cache is backed by a regular dict, then "key" must be of
- the right type for this cache
- If the cache is backed by a TreeCache, then "key" must be a tuple, but
- may be of lower cardinality than the TreeCache - in which case the whole
- subtree is deleted.
- """
- self.check_thread()
- self.cache.del_multi(key)
- # if we have a pending lookup for this key, remove it from the
- # _pending_deferred_cache, which will (a) stop it being returned for
- # future queries and (b) stop it being persisted as a proper entry
- # in self.cache.
- entry = self._pending_deferred_cache.pop(key, None)
- if entry:
- # _pending_deferred_cache.pop should either return a CacheEntry, or, in the
- # case of a TreeCache, a dict of keys to cache entries. Either way calling
- # iterate_tree_cache_entry on it will do the right thing.
- for iter_entry in iterate_tree_cache_entry(entry):
- for cb in iter_entry.get_invalidation_callbacks(key):
- cb()
- def invalidate_all(self) -> None:
- self.check_thread()
- self.cache.clear()
- for key, entry in self._pending_deferred_cache.items():
- for cb in entry.get_invalidation_callbacks(key):
- cb()
- self._pending_deferred_cache.clear()
- class CacheEntry(Generic[KT, VT], metaclass=abc.ABCMeta):
- """Abstract class for entries in `DeferredCache[KT, VT]`"""
- @abc.abstractmethod
- def deferred(self, key: KT) -> "defer.Deferred[VT]":
- """Get a deferred that a caller can wait on to get the value at the
- given key"""
- ...
- @abc.abstractmethod
- def add_invalidation_callback(
- self, key: KT, callback: Optional[Callable[[], None]]
- ) -> None:
- """Add an invalidation callback"""
- ...
- @abc.abstractmethod
- def get_invalidation_callbacks(self, key: KT) -> Collection[Callable[[], None]]:
- """Get all invalidation callbacks"""
- ...
- class CacheEntrySingle(CacheEntry[KT, VT]):
- """An implementation of `CacheEntry` wrapping a deferred that results in a
- single cache entry.
- """
- __slots__ = ["_deferred", "_callbacks"]
- def __init__(self, deferred: "defer.Deferred[VT]") -> None:
- self._deferred = ObservableDeferred(deferred, consumeErrors=True)
- self._callbacks: Set[Callable[[], None]] = set()
- def deferred(self, key: KT) -> "defer.Deferred[VT]":
- return self._deferred.observe()
- def add_invalidation_callback(
- self, key: KT, callback: Optional[Callable[[], None]]
- ) -> None:
- if callback is None:
- return
- self._callbacks.add(callback)
- def get_invalidation_callbacks(self, key: KT) -> Collection[Callable[[], None]]:
- return self._callbacks
- class CacheMultipleEntries(CacheEntry[KT, VT]):
- """Cache entry that is used for bulk lookups and insertions."""
- __slots__ = ["_deferred", "_callbacks", "_global_callbacks"]
- def __init__(self) -> None:
- self._deferred: Optional[ObservableDeferred[Dict[KT, VT]]] = None
- self._callbacks: Dict[KT, Set[Callable[[], None]]] = {}
- self._global_callbacks: Set[Callable[[], None]] = set()
- def deferred(self, key: KT) -> "defer.Deferred[VT]":
- if not self._deferred:
- self._deferred = ObservableDeferred(defer.Deferred(), consumeErrors=True)
- return self._deferred.observe().addCallback(lambda res: res[key])
- def add_invalidation_callback(
- self, key: KT, callback: Optional[Callable[[], None]]
- ) -> None:
- if callback is None:
- return
- self._callbacks.setdefault(key, set()).add(callback)
- def get_invalidation_callbacks(self, key: KT) -> Collection[Callable[[], None]]:
- return self._callbacks.get(key, set()) | self._global_callbacks
- def add_global_invalidation_callback(
- self, callback: Optional[Callable[[], None]]
- ) -> None:
- """Add a callback for when any keys get invalidated."""
- if callback is None:
- return
- self._global_callbacks.add(callback)
- def complete_bulk(
- self,
- cache: DeferredCache[KT, VT],
- result: Dict[KT, VT],
- ) -> None:
- """Called when there is a result"""
- for key, value in result.items():
- cache._completed_callback(value, self, key)
- if self._deferred:
- self._deferred.callback(result)
- def error_bulk(
- self, cache: DeferredCache[KT, VT], keys: Collection[KT], failure: Failure
- ) -> None:
- """Called when bulk lookup failed."""
- for key in keys:
- cache._error_callback(failure, self, key)
- if self._deferred:
- self._deferred.errback(failure)
|