Update default params test

Author: shaypal5

README.rst

@@ -286,6 +286,19 @@ human readable string like ``"200MB"``.
 When ``cachier__verbose=True`` is passed to a call that returns a value
 exceeding the limit, an informative message is printed.
 
+Cache Size Limit
+~~~~~~~~~~~~~~~~
+``cache_size_limit`` constrains the total size of the cache. When the
+limit is exceeded, entries are evicted according to the chosen
+``replacement_policy``. Currently only an ``"lru"`` policy is implemented
+for the in-memory backend.
+
+.. code-block:: python
+
+  @cachier(cache_size_limit="100KB")
+  def heavy(x):
+      return x * 2
+
 Ignore Cache
 ~~~~~~~~~~~~
 

src/cachier/config.py

@@ -66,6 +66,8 @@ class Params:
     cleanup_stale: bool = False
     cleanup_interval: timedelta = timedelta(days=1)
     entry_size_limit: Optional[int] = None
+    cache_size_limit: Optional[int] = None
+    replacement_policy: str = "lru"
 
 
 _global_params = Params()

src/cachier/core.py

@@ -124,6 +124,8 @@ def cachier(
     cleanup_stale: Optional[bool] = None,
     cleanup_interval: Optional[timedelta] = None,
     entry_size_limit: Optional[Union[int, str]] = None,
+    cache_size_limit: Optional[Union[int, str]] = None,
+    replacement_policy: str = "lru",
 ):
     """Wrap as a persistent, stale-free memoization decorator.
 
@@ -196,6 +198,12 @@ def cachier(
         Maximum serialized size of a cached value. Values exceeding the limit
         are returned but not cached. Human readable strings like ``"10MB"`` are
         allowed.
+    cache_size_limit: int or str, optional
+        Maximum total size allowed for the cache. When exceeded, entries are
+        evicted according to ``replacement_policy``.
+    replacement_policy: str, optional
+        Cache replacement policy used when trimming the cache. Currently only
+        ``"lru"`` is supported.
 
     """
     # Check for deprecated parameters
@@ -212,6 +220,10 @@ def cachier(
     size_limit_bytes = parse_bytes(
         _update_with_defaults(entry_size_limit, "entry_size_limit")
     )
+    cache_limit_bytes = parse_bytes(
+        _update_with_defaults(cache_size_limit, "cache_size_limit")
+    )
+    policy = _update_with_defaults(replacement_policy, "replacement_policy")
     # Override the backend parameter if a mongetter is provided.
     if callable(mongetter):
         backend = "mongo"
@@ -224,33 +236,43 @@ def cachier(
             separate_files=separate_files,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=size_limit_bytes,
+            cache_size_limit=cache_limit_bytes,
+            replacement_policy=policy,
         )
     elif backend == "mongo":
         core = _MongoCore(
             hash_func=hash_func,
             mongetter=mongetter,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=size_limit_bytes,
+            cache_size_limit=cache_limit_bytes,
+            replacement_policy=policy,
         )
     elif backend == "memory":
         core = _MemoryCore(
             hash_func=hash_func,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=size_limit_bytes,
+            cache_size_limit=cache_limit_bytes,
+            replacement_policy=policy,
         )
     elif backend == "sql":
         core = _SQLCore(
             hash_func=hash_func,
             sql_engine=sql_engine,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=size_limit_bytes,
+            cache_size_limit=cache_limit_bytes,
+            replacement_policy=policy,
         )
     elif backend == "redis":
         core = _RedisCore(
             hash_func=hash_func,
             redis_client=redis_client,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=size_limit_bytes,
+            cache_size_limit=cache_limit_bytes,
+            replacement_policy=policy,
         )
     else:
         raise ValueError("specified an invalid core: %s" % backend)

src/cachier/cores/base.py

@@ -38,11 +38,15 @@ def __init__(
         hash_func: Optional[HashFunc],
         wait_for_calc_timeout: Optional[int],
         entry_size_limit: Optional[int] = None,
+        cache_size_limit: Optional[int] = None,
+        replacement_policy: str = "lru",
     ):
         self.hash_func = _update_with_defaults(hash_func, "hash_func")
         self.wait_for_calc_timeout = wait_for_calc_timeout
         self.lock = threading.RLock()
         self.entry_size_limit = entry_size_limit
+        self.cache_size_limit = cache_size_limit
+        self.replacement_policy = replacement_policy
 
     def set_func(self, func):
         """Set the function this core will use.

src/cachier/cores/memory.py

@@ -1,8 +1,9 @@
 """A memory-based caching core for cachier."""
 
 import threading
+from collections import OrderedDict
 from datetime import datetime, timedelta
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Optional, Tuple
 
 from .._types import HashFunc
 from ..config import CacheEntry
@@ -17,9 +18,18 @@ def __init__(
         hash_func: Optional[HashFunc],
         wait_for_calc_timeout: Optional[int],
         entry_size_limit: Optional[int] = None,
+        cache_size_limit: Optional[int] = None,
+        replacement_policy: str = "lru",
     ):
-        super().__init__(hash_func, wait_for_calc_timeout, entry_size_limit)
-        self.cache: Dict[str, CacheEntry] = {}
+        super().__init__(
+            hash_func,
+            wait_for_calc_timeout,
+            entry_size_limit,
+            cache_size_limit,
+            replacement_policy,
+        )
+        self.cache: "OrderedDict[str, CacheEntry]" = OrderedDict()
+        self._cache_size = 0
 
     def _hash_func_key(self, key: str) -> str:
         return f"{_get_func_str(self.func)}:{key}"
@@ -28,18 +38,22 @@ def get_entry_by_key(
         self, key: str, reload=False
     ) -> Tuple[str, Optional[CacheEntry]]:
         with self.lock:
-            return key, self.cache.get(self._hash_func_key(key), None)
+            hkey = self._hash_func_key(key)
+            entry = self.cache.get(hkey, None)
+            if entry is not None:
+                self.cache.move_to_end(hkey)
+            return key, entry
 
     def set_entry(self, key: str, func_res: Any) -> bool:
         if not self._should_store(func_res):
             return False
         hash_key = self._hash_func_key(key)
+        size = self._estimate_size(func_res)
         with self.lock:
             try:
-                # we need to retain the existing condition so that
-                # mark_entry_not_calculated can notify all possibly-waiting
-                # threads about it
                 cond = self.cache[hash_key]._condition
+                old_size = self._estimate_size(self.cache[hash_key].value)
+                self._cache_size -= old_size
             except KeyError:  # pragma: no cover
                 cond = None
             self.cache[hash_key] = CacheEntry(
@@ -50,6 +64,12 @@ def set_entry(self, key: str, func_res: Any) -> bool:
                 _condition=cond,
                 _completed=True,
             )
+            self.cache.move_to_end(hash_key)
+            self._cache_size += size
+            if self.cache_size_limit is not None:
+                while self._cache_size > self.cache_size_limit and self.cache:
+                    old_key, old_entry = self.cache.popitem(last=False)
+                    self._cache_size -= self._estimate_size(old_entry.value)
         return True
 
     def mark_entry_being_calculated(self, key: str) -> None:
@@ -101,6 +121,7 @@ def wait_on_entry_calc(self, key: str) -> Any:
     def clear_cache(self) -> None:
         with self.lock:
             self.cache.clear()
+            self._cache_size = 0
 
     def clear_being_calculated(self) -> None:
         with self.lock:
@@ -116,4 +137,5 @@ def delete_stale_entries(self, stale_after: timedelta) -> None:
                 k for k, v in self.cache.items() if now - v.time > stale_after
             ]
             for key in keys_to_delete:
-                del self.cache[key]
+                entry = self.cache.pop(key)
+                self._cache_size -= self._estimate_size(entry.value)

src/cachier/cores/mongo.py

@@ -41,6 +41,8 @@ def __init__(
         mongetter: Optional[Mongetter],
         wait_for_calc_timeout: Optional[int],
         entry_size_limit: Optional[int] = None,
+        cache_size_limit: Optional[int] = None,
+        replacement_policy: str = "lru",
     ):
         if "pymongo" not in sys.modules:
             warnings.warn(
@@ -53,6 +55,8 @@ def __init__(
             hash_func=hash_func,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=entry_size_limit,
+            cache_size_limit=cache_size_limit,
+            replacement_policy=replacement_policy,
         )
         if mongetter is None:
             raise MissingMongetter(

src/cachier/cores/pickle.py

@@ -79,8 +79,16 @@ def __init__(
         separate_files: Optional[bool],
         wait_for_calc_timeout: Optional[int],
         entry_size_limit: Optional[int] = None,
+        cache_size_limit: Optional[int] = None,
+        replacement_policy: str = "lru",
     ):
-        super().__init__(hash_func, wait_for_calc_timeout, entry_size_limit)
+        super().__init__(
+            hash_func,
+            wait_for_calc_timeout,
+            entry_size_limit,
+            cache_size_limit,
+            replacement_policy,
+        )
         self._cache_dict: Dict[str, CacheEntry] = {}
         self.reload = _update_with_defaults(pickle_reload, "pickle_reload")
         self.cache_dir = os.path.expanduser(

src/cachier/cores/redis.py

@@ -36,6 +36,8 @@ def __init__(
         wait_for_calc_timeout: Optional[int] = None,
         key_prefix: str = "cachier",
         entry_size_limit: Optional[int] = None,
+        cache_size_limit: Optional[int] = None,
+        replacement_policy: str = "lru",
     ):
         if not REDIS_AVAILABLE:
             warnings.warn(
@@ -49,6 +51,8 @@ def __init__(
             hash_func=hash_func,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=entry_size_limit,
+            cache_size_limit=cache_size_limit,
+            replacement_policy=replacement_policy,
         )
         if redis_client is None:
             raise MissingRedisClient(

src/cachier/cores/sql.py

@@ -64,6 +64,8 @@ def __init__(
         sql_engine: Optional[Union[str, "Engine", Callable[[], "Engine"]]],
         wait_for_calc_timeout: Optional[int] = None,
         entry_size_limit: Optional[int] = None,
+        cache_size_limit: Optional[int] = None,
+        replacement_policy: str = "lru",
     ):
         if not SQLALCHEMY_AVAILABLE:
             raise ImportError(
@@ -74,6 +76,8 @@ def __init__(
             hash_func=hash_func,
             wait_for_calc_timeout=wait_for_calc_timeout,
             entry_size_limit=entry_size_limit,
+            cache_size_limit=cache_size_limit,
+            replacement_policy=replacement_policy,
         )
         self._engine = self._resolve_engine(sql_engine)
         self._Session = sessionmaker(bind=self._engine)

tests/test_cache_size_limit.py

@@ -0,0 +1,25 @@
+import pytest
+
+import cachier
+
+
+@pytest.mark.memory
+def test_cache_size_limit_lru_eviction():
+    call_count = 0
+
+    @cachier.cachier(backend="memory", cache_size_limit="220B")
+    def func(x):
+        nonlocal call_count
+        call_count += 1
+        return "a" * 50
+
+    func.clear_cache()
+    func(1)
+    func(2)
+    assert call_count == 2
+    func(1)  # access to update LRU order
+    assert call_count == 2
+    func(3)  # should evict key 2
+    assert call_count == 3
+    func(2)
+    assert call_count == 4