diff --git a/lisa/tests/base.py b/lisa/tests/base.py index a9d659dd9d7bfa80ba8b4e55fd4cb66e64b7f529..b372ce8f9ce59e5734d7b57a2f18f757e55df5cd 100644 --- a/lisa/tests/base.py +++ b/lisa/tests/base.py @@ -597,10 +597,7 @@ class RTATestBundle(TestBundle, metaclass=RTATestBundleMeta): # Guard before the cache, so we don't accidentally start depending on the # LRU cache for functionnal correctness. @non_recursive_property - # Use LRU cache instead of memoized, to avoid caching the trace forever, in - # case the thread is manipulating a large number of TestBundles without - # deleting them. - @functools.lru_cache(maxsize=30, typed=True) + @memoized def trace(self): """ :returns: a :class:`lisa.trace.TraceView` diff --git a/lisa/utils.py b/lisa/utils.py index c97f8449e8a200e6117875804f8cfcf3eac03b12..336fb5af2a9ce83f16cbe702a739a608dae3047f 100644 --- a/lisa/utils.py +++ b/lisa/utils.py @@ -35,6 +35,7 @@ import numbers import difflib import threading import itertools +import weakref from weakref import WeakKeyDictionary import lisa @@ -140,18 +141,47 @@ class HideExekallID: """ pass -def memoized(callable_): +def memoized(f): """ - Decorator to memoize the result of a callable, - based on :func:`functools.lru_cache` + Decorator to memoize the result of a callable, based on + :func:`functools.lru_cache` + + .. note:: The first parameter of the callable is cached with a weak + reference. This suits well the method use-case, since we don't want the + memoization of methods to prevent garbage collection of the instances + they are bound to. """ - # It is important to have one separate call to lru_cache for every call to - # memoized, otherwise all uses of the decorator will end up using the same - # wrapper and all hells will break loose. - # maxsize should be a power of two for better speed, see: - # https://docs.python.org/3/library/functools.html#functools.lru_cache - return functools.lru_cache(maxsize=1024, typed=True)(callable_) + def apply_lru(f): + # maxsize should be a power of two for better speed, see: + # https://docs.python.org/3/library/functools.html#functools.lru_cache + return functools.lru_cache(maxsize=1024, typed=True)(f) + + # We need at least one positional parameter for the WeakKeyDictionary + if inspect.signature(f).parameters: + cache_map = WeakKeyDictionary() + + @functools.wraps(f) + def wrapper(first, *args, **kwargs): + try: + partial = cache_map[first] + except KeyError: + # Only keep a weak reference here for the "partial" closure + ref = weakref.ref(first) + + # This partial function does not take "first" as parameter, so + # that the lru_cache will not keep a reference on it + @apply_lru + def partial(*args, **kwargs): + return f(ref(), *args, **kwargs) + + cache_map[first] = partial + + return partial(*args, **kwargs) + + return wrapper + else: + return apply_lru(f) def resolve_dotted_name(name): """Only resolve names where __qualname__ == __name__, i.e the callable is a