Merge pull request #39 from stalkerg/memory-limit

Max memory limit

Author: awolverp

python/cachebox/_cachebox.py

@@ -27,14 +27,19 @@ class BaseCacheImpl(typing.Generic[KT, VT]):
     def __init__(
         self,
         maxsize: int,
-        iterable: typing.Union[typing.Iterable[typing.Tuple[KT, VT]], typing.Dict[KT, VT]] = ...,
+        iterable: typing.Union[
+            typing.Iterable[typing.Tuple[KT, VT]], typing.Dict[KT, VT]
+        ] = ...,
         *,
         capacity: int = ...,
+        maxmemory: int = ...,
     ) -> None: ...
     @staticmethod
     def __class_getitem__(*args: typing.Any) -> None: ...
     @property
     def maxsize(self) -> int: ...
+    @property
+    def maxmemory(self) -> int: ...
     def __len__(self) -> int: ...
     def __sizeof__(self) -> int: ...
     def __bool__(self) -> bool: ...
@@ -47,29 +52,42 @@ class BaseCacheImpl(typing.Generic[KT, VT]):
     def __eq__(self, other: typing.Any) -> bool: ...
     def __ne__(self, other: typing.Any) -> bool: ...
     def capacity(self) -> int: ...
+    def memory(self) -> int: ...
     def is_full(self) -> bool: ...
     def is_empty(self) -> bool: ...
     def insert(
         self, key: KT, value: VT, *args: typing.Any, **kwargs: typing.Any
     ) -> typing.Optional[VT]: ...
-    def get(self, key: KT, default: typing.Optional[DT] = None) -> typing.Union[VT, DT]: ...
-    def pop(self, key: KT, default: typing.Optional[DT] = None) -> typing.Union[VT, DT]: ...
+    def get(
+        self, key: KT, default: typing.Optional[DT] = None
+    ) -> typing.Union[VT, DT]: ...
+    def pop(
+        self, key: KT, default: typing.Optional[DT] = None
+    ) -> typing.Union[VT, DT]: ...
     def setdefault(
-        self, key: KT, default: typing.Optional[DT] = None, *args: typing.Any, **kwargs: typing.Any
+        self,
+        key: KT,
+        default: typing.Optional[DT] = None,
+        *args: typing.Any,
+        **kwargs: typing.Any,
     ) -> typing.Optional[VT | DT]: ...
     def popitem(self) -> typing.Tuple[KT, VT]: ...
     def drain(self, n: int) -> int: ...
     def clear(self, *, reuse: bool = False) -> None: ...
     def shrink_to_fit(self) -> None: ...
     def update(
         self,
-        iterable: typing.Union[typing.Iterable[typing.Tuple[KT, VT]], typing.Dict[KT, VT]],
+        iterable: typing.Union[
+            typing.Iterable[typing.Tuple[KT, VT]], typing.Dict[KT, VT]
+        ],
         *args: typing.Any,
         **kwargs: typing.Any,
     ) -> None: ...
     def keys(self) -> typing.Iterable[KT]: ...
     def values(self) -> typing.Iterable[VT]: ...
     def items(self) -> typing.Iterable[typing.Tuple[KT, VT]]: ...
     def __copy__(self) -> "BaseCacheImpl[KT, VT]": ...
-    def __deepcopy__(self, memo: typing.Dict[str, object]) -> "BaseCacheImpl[KT, VT]": ...
+    def __deepcopy__(
+        self, memo: typing.Dict[str, object]
+    ) -> "BaseCacheImpl[KT, VT]": ...
     def copy(self) -> "BaseCacheImpl[KT, VT]": ...

python/cachebox/_core.pyi

@@ -1,11 +1,11 @@
-from ._cachebox import BaseCacheImpl, FIFOCache
-from collections import namedtuple, defaultdict
-import functools
-import asyncio
 import _thread
+import asyncio
+import functools
 import inspect
 import typing
+from collections import defaultdict, namedtuple
 
+from ._cachebox import BaseCacheImpl, FIFOCache
 
 KT = typing.TypeVar("KT")
 VT = typing.TypeVar("VT")
@@ -46,6 +46,10 @@ def cache(self) -> BaseCacheImpl[KT, VT]:
     def maxsize(self) -> int:
         return self.__cache.maxsize
 
+    @property
+    def maxmemory(self) -> int:
+        return self.__cache.maxmemory
+
     def __len__(self) -> int:
         return len(self.__cache)
 
@@ -85,6 +89,9 @@ def __richcmp__(self, other: typing.Any, op: int) -> bool:
     def capacity(self) -> int:
         return self.__cache.capacity()
 
+    def memory(self) -> int:
+        return self.__cache.memory()
+
     def is_full(self) -> bool:
         return self.__cache.is_full()
 
@@ -140,7 +147,9 @@ def shrink_to_fit(self) -> None:
 
     def update(
         self,
-        iterable: typing.Union[typing.Iterable[typing.Tuple[KT, VT]], typing.Dict[KT, VT]],
+        iterable: typing.Union[
+            typing.Iterable[typing.Tuple[KT, VT]], typing.Dict[KT, VT]
+        ],
         *args,
         **kwargs,
     ) -> None:
@@ -270,16 +279,22 @@ def _cached_wrapper(
     cache: typing.Union[BaseCacheImpl, typing.Callable],
     key_maker: typing.Callable[[tuple, dict], typing.Hashable],
     clear_reuse: bool,
-    callback: typing.Optional[typing.Callable[[int, typing.Any, typing.Any], typing.Any]],
+    callback: typing.Optional[
+        typing.Callable[[int, typing.Any, typing.Any], typing.Any]
+    ],
     copy_level: int,
     is_method: bool,
 ):
     is_method = cache_is_function = inspect.isfunction(cache)
-    _key_maker = (lambda args, kwds: key_maker(args[1:], kwds)) if is_method else key_maker
+    _key_maker = (
+        (lambda args, kwds: key_maker(args[1:], kwds)) if is_method else key_maker
+    )
 
     hits = 0
     misses = 0
-    locks: defaultdict[typing.Hashable, _LockWithCounter] = defaultdict(_LockWithCounter)
+    locks: defaultdict[typing.Hashable, _LockWithCounter] = defaultdict(
+        _LockWithCounter
+    )
     exceptions: typing.Dict[typing.Hashable, BaseException] = {}
 
     def _wrapped(*args, **kwds):
@@ -308,7 +323,9 @@ def _wrapped(*args, **kwds):
 
         with locks[key]:
             if exceptions.get(key, None) is not None:
-                cached_error = exceptions[key] if locks[key].waiters > 1 else exceptions.pop(key)
+                cached_error = (
+                    exceptions[key] if locks[key].waiters > 1 else exceptions.pop(key)
+                )
                 raise cached_error
 
             try:
@@ -337,7 +354,7 @@ def _wrapped(*args, **kwds):
     if not cache_is_function:
         _wrapped.cache = cache
         _wrapped.cache_info = lambda: CacheInfo(
-            hits, misses, cache.maxsize, len(cache), cache.capacity()
+            hits, misses, cache.maxsize, len(cache), cache.memory()
         )
 
     _wrapped.callback = callback
@@ -362,12 +379,16 @@ def _async_cached_wrapper(
     cache: typing.Union[BaseCacheImpl, typing.Callable],
     key_maker: typing.Callable[[tuple, dict], typing.Hashable],
     clear_reuse: bool,
-    callback: typing.Optional[typing.Callable[[int, typing.Any, typing.Any], typing.Any]],
+    callback: typing.Optional[
+        typing.Callable[[int, typing.Any, typing.Any], typing.Any]
+    ],
     copy_level: int,
     is_method: bool,
 ):
     is_method = cache_is_function = inspect.isfunction(cache)
-    _key_maker = (lambda args, kwds: key_maker(args[1:], kwds)) if is_method else key_maker
+    _key_maker = (
+        (lambda args, kwds: key_maker(args[1:], kwds)) if is_method else key_maker
+    )
 
     hits = 0
     misses = 0
@@ -404,7 +425,9 @@ async def _wrapped(*args, **kwds):
 
         async with locks[key]:
             if exceptions.get(key, None) is not None:
-                cached_error = exceptions[key] if locks[key].waiters > 1 else exceptions.pop(key)
+                cached_error = (
+                    exceptions[key] if locks[key].waiters > 1 else exceptions.pop(key)
+                )
                 raise cached_error
 
             try:
@@ -435,7 +458,7 @@ async def _wrapped(*args, **kwds):
     if not cache_is_function:
         _wrapped.cache = cache
         _wrapped.cache_info = lambda: CacheInfo(
-            hits, misses, cache.maxsize, len(cache), cache.capacity()
+            hits, misses, cache.maxsize, len(cache), cache.memory()
         )
 
     _wrapped.callback = callback
@@ -459,7 +482,9 @@ def cached(
     cache: typing.Union[BaseCacheImpl, dict, None],
     key_maker: typing.Callable[[tuple, dict], typing.Hashable] = make_key,
     clear_reuse: bool = False,
-    callback: typing.Optional[typing.Callable[[int, typing.Any, typing.Any], typing.Any]] = None,
+    callback: typing.Optional[
+        typing.Callable[[int, typing.Any, typing.Any], typing.Any]
+    ] = None,
     copy_level: int = 1,
 ) -> typing.Callable[[FT], FT]:
     """
@@ -532,7 +557,9 @@ def cachedmethod(
     cache: typing.Union[BaseCacheImpl, dict, None],
     key_maker: typing.Callable[[tuple, dict], typing.Hashable] = make_key,
     clear_reuse: bool = False,
-    callback: typing.Optional[typing.Callable[[int, typing.Any, typing.Any], typing.Any]] = None,
+    callback: typing.Optional[
+        typing.Callable[[int, typing.Any, typing.Any], typing.Any]
+    ] = None,
     copy_level: int = 1,
 ) -> typing.Callable[[FT], FT]:
     """

python/cachebox/utils.py

@@ -1,8 +1,9 @@
-from cachebox import BaseCacheImpl, TTLCache
 import dataclasses
-import pytest
-import typing
 import sys
+import typing
+
+import pytest
+from cachebox import BaseCacheImpl, TTLCache
 
 
 @dataclasses.dataclass
@@ -26,6 +27,28 @@ def __hash__(self) -> int:
         return self.val
 
 
+@dataclasses.dataclass
+class Sized:
+    size: int
+    key: int
+
+    def __sizeof__(self) -> int:
+        return self.size
+
+    def __hash__(self) -> int:
+        return self.key
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Sized):
+            return False
+        return self.key == other.key
+
+
+class SizeError:
+    def __sizeof__(self) -> int:
+        raise ValueError("boom")
+
+
 def getsizeof(obj, use_sys=True):  # pragma: no cover
     try:
         if use_sys:
@@ -71,6 +94,56 @@ def test_overflow(self):
         with pytest.raises(OverflowError):
             cache["new-key"] = "new-value"
 
+    def test_maxmemory_config(self):
+        cache = self.CACHE(10, **self.KWARGS, maxmemory=128)
+        assert cache.maxmemory == 128
+        assert cache.memory() == 0
+
+    def test_maxmemory_enforced(self):
+        cache = self.CACHE(0, **self.KWARGS, maxmemory=100)
+
+        k1 = Sized(10, 1)
+        v1 = Sized(80, 101)
+        cache[k1] = v1
+
+        k2 = Sized(10, 2)
+        v2 = Sized(80, 102)
+
+        if self.NO_POLICY:
+            with pytest.raises(OverflowError):
+                cache[k2] = v2
+            assert k1 in cache
+        else:
+            cache[k2] = v2
+            assert k2 in cache
+            assert cache.memory() <= cache.maxmemory
+
+    def test_update_overflow_preserves_entry(self):
+        cache = self.CACHE(0, **self.KWARGS, maxmemory=60)
+
+        key = Sized(10, 1)
+        value = Sized(10, 101)
+        cache[key] = value
+
+        too_big = Sized(100, 102)
+        with pytest.raises(OverflowError):
+            cache[key] = too_big
+
+        assert cache[key].key == 101
+        assert cache.memory() <= cache.maxmemory
+
+    def test_update_sizeof_error_preserves_entry(self):
+        cache = self.CACHE(0, **self.KWARGS, maxmemory=60)
+
+        key = Sized(10, 1)
+        value = Sized(10, 101)
+        cache[key] = value
+
+        with pytest.raises(ValueError):
+            cache[key] = SizeError()
+
+        assert cache[key].key == 101
+
     def test___len__(self):
         cache = self.CACHE(10, **self.KWARGS, capacity=10)
 

python/tests/mixin.py

@@ -1,16 +1,18 @@
+import time
+from datetime import timedelta
+
+import pytest
 from cachebox import (
     Cache,
     FIFOCache,
-    RRCache,
-    LRUCache,
     LFUCache,
+    LRUCache,
+    RRCache,
     TTLCache,
     VTTLCache,
 )
-from datetime import timedelta
-import pytest
-from .mixin import _TestMixin
-import time
+
+from .mixin import Sized, _TestMixin
 
 
 class TestCache(_TestMixin):
@@ -64,6 +66,23 @@ def test_policy(self):
 
         assert cache.popitem() == (10, 10)
 
+    def test_update_can_evict_self_on_maxmemory(self):
+        cache = FIFOCache(0, maxmemory=50)
+
+        k1 = Sized(10, 1)
+        v1 = Sized(10, 101)
+        k2 = Sized(10, 2)
+        v2 = Sized(10, 102)
+
+        cache[k1] = v1
+        cache[k2] = v2
+
+        cache[k1] = Sized(40, 103)
+
+        assert k1 not in cache
+        assert k2 in cache
+        assert cache.memory() <= cache.maxmemory
+
     def test_ordered_iterators(self):
         obj = self.CACHE(100, **self.KWARGS, capacity=100)