Fix several lock, async, exception-handling and counting bugs (+ tests)

cachebox/_wrappers.py

@@ -18,8 +18,10 @@ def __init__(self, lock: AbstractContextManager) -> None:
 
     def __enter__(self) -> None:
         self.waiters += 1
-        self._lock.__enter__()
-        self.waiters -= 1
+        try:
+            self._lock.__enter__()
+        finally:
+            self.waiters -= 1
 
     def __exit__(self, *_) -> None:
         self._lock.__exit__(*_)
@@ -34,8 +36,10 @@ def __init__(self, lock: AbstractAsyncContextManager) -> None:
 
     async def __aenter__(self) -> None:
         self.waiters += 1
-        await self._lock.__aenter__()
-        self.waiters -= 1
+        try:
+            await self._lock.__aenter__()
+        finally:
+            self.waiters -= 1
 
     async def __aexit__(self, *_) -> None:
         await self._lock.__aexit__(*_)
@@ -102,7 +106,7 @@ def _wrapped(*args, **kwds):
 
         result = func(*args, **kwds)
         _cache.insert(key, result)
-        hits += 1
+        misses += 1
         if callback is not None:
             callback(EVENT_MISS, key, result)
 
@@ -159,7 +163,7 @@ async def _wrapped(*args, **kwds):
         # Passing `cachebox__ignore=True` bypasses the cache and
         # calls the function directly.
         if kwds.pop("cachebox__ignore", False):
-            return func(*args, **kwds)
+            return await func(*args, **kwds)
 
         _cache: BaseCacheImpl = cache(args[0]) if cache_is_fn else cache  # type: ignore[arg-type]
         key = _make_key(args, kwds)
@@ -266,17 +270,19 @@ def _wrapped(*args, **kwds):
             except KeyError:
                 try:
                     result = func(*args, **kwds)
-                except Exception as exc:
+                except BaseException as exc:
                     if lock.waiters > 0:
                         pending_errors[key] = exc
+
                     raise
                 else:
                     _cache[key] = result
                     misses += 1
                     event = EVENT_MISS
 
-            if lock.waiters == 0:
-                locks.pop(key, None)
+            finally:
+                if lock.waiters == 0:
+                    locks.pop(key, None)
 
         if callback is not None:
             callback(event, key, result)
@@ -319,29 +325,7 @@ def _async_cached_wrapper(
     hits = 0
     misses = 0
 
-    # See _cached_wrapper
     locks: Cache[typing.Hashable, _AsyncLock] = Cache(0)
-
-    # if lock_type is asyncio.Lock:
-
-    #     async def _get_lock(key: typing.Hashable):
-    #         return locks.setdefault(key, _AsyncLock(lock_type()))
-
-    # else:
-    #     _lock_creation_guard = asyncio.Lock()
-
-    #     async def _get_lock(key: typing.Hashable):
-    #         lock = locks.get(key)
-    #         if lock is not None:
-    #             return lock
-
-    #         async with _lock_creation_guard:
-    #             lock = locks.get(key)
-    #             if lock is None:
-    #                 locks[key] = lock = _AsyncLock(lock_type())
-
-    #         return lock
-
     pending_errors: dict[typing.Hashable, BaseException] = {}
 
     async def _wrapped(*args, **kwds):
@@ -381,17 +365,17 @@ async def _wrapped(*args, **kwds):
             except KeyError:
                 try:
                     result = await func(*args, **kwds)
-                except Exception as exc:
+                except BaseException as exc:
                     if lock.waiters > 0:
                         pending_errors[key] = exc
                     raise
                 else:
                     _cache[key] = result
                     misses += 1
                     event = EVENT_MISS
-
-            if lock.waiters == 0:
-                locks.pop(key, None)
+            finally:
+                if lock.waiters == 0:
+                    locks.pop(key, None)
 
         await _call_async_callback(callback, event, key, result)
 

cachebox/utils.py

@@ -73,12 +73,14 @@ def make_key(*args, **kwds) -> typing.Hashable:
     if not kwds:
         if len(args) == 1 and type(args[0]) in _FAST_TYPES:
             return args[0]
+
         return args
 
     key = args + (_KWDS_MARK,)
     for item in kwds.items():
         key += item
-    return key[0] if len(key) == 1 and type(key[0]) in _FAST_TYPES else key
+
+    return key
 
 
 def make_hash_key(*args, **kwds) -> int:
@@ -155,8 +157,13 @@ def __init__(self, cls: BaseCacheImpl[KT, VT], ignore: bool = False) -> None:
             ignore: If ``True``, silently ignores modification attempts; if ``False``, raises
                 ``TypeError`` when modification is attempted. Default is ``False``.
         """
-        assert isinstance(cls, BaseCacheImpl)
-        assert type(cls) is not Frozen
+        if not isinstance(cls, BaseCacheImpl):
+            raise TypeError(
+                f"expected a BaseCacheImpl instance, got {type(cls).__name__!r}"
+            )
+
+        if type(cls) is Frozen:
+            raise TypeError("cannot wrap an already-frozen cache")
 
         self.__cache = cls
         self.ignore = ignore
@@ -366,9 +373,16 @@ def __repr__(self) -> str:
 def _cast_lock(
     iscoroutinefunction: bool,
     lock: (
-        typing.Type[AbstractContextManager] | typing.Type[AbstractAsyncContextManager] | bool | None
+        typing.Type[AbstractContextManager]
+        | typing.Type[AbstractAsyncContextManager]
+        | bool
+        | None
     ) = True,
-) -> typing.Type[AbstractContextManager] | typing.Type[AbstractAsyncContextManager] | None:
+) -> (
+    typing.Type[AbstractContextManager]
+    | typing.Type[AbstractAsyncContextManager]
+    | None
+):
     import _thread
     import asyncio
     import threading
@@ -381,12 +395,18 @@ def _cast_lock(
 
     if iscoroutinefunction:
         if not hasattr(lock, "__aenter__"):
-            raise TypeError("For async functions, you cannot use a regular synchronous lock.")
+            raise TypeError(
+                "For async functions, you cannot use a regular synchronous lock."
+            )
 
         return typing.cast(typing.Type[AbstractAsyncContextManager], lock)
 
     # threading.Lock, threading.RLock and _thread.allocate_lock are function
-    if lock is threading.Lock or lock is threading.RLock or lock is _thread.allocate_lock:
+    if (
+        lock is threading.Lock
+        or lock is threading.RLock
+        or lock is _thread.allocate_lock
+    ):
         return typing.cast(typing.Type[AbstractContextManager], lock)
 
     if not hasattr(lock, "__enter__"):
@@ -403,7 +423,10 @@ def cached(
     copy_level: int = 1,
     postprocess: _PostProcess | None = postprocess_copy_mutables,
     lock: (
-        typing.Type[AbstractContextManager] | typing.Type[AbstractAsyncContextManager] | bool | None
+        typing.Type[AbstractContextManager]
+        | typing.Type[AbstractAsyncContextManager]
+        | bool
+        | None
     ) = True,
 ) -> typing.Callable[[FT], FT]:
     """
@@ -480,7 +503,9 @@ def decorator(func: FT) -> FT:
         lock_type = _cast_lock(iscoroutinefunction, lock)
 
         if not iscoroutinefunction and inspect.iscoroutinefunction(callback):
-            raise TypeError("For sync functions, you cannot use a asynchronous callback")
+            raise TypeError(
+                "For sync functions, you cannot use a asynchronous callback"
+            )
 
         if lock_type:
             builder = _async_cached_wrapper if iscoroutinefunction else _cached_wrapper

tests/test_utils.py

@@ -742,3 +742,124 @@ async def calc(a, b, c, exception: bool):
         await calc(1, 2, 3, exception=True)
 
     assert await calc(1, 2, 3, exception=False) == 6
+
+
+def test_misses_should_increment_without_lock():
+    calls = {"n": 0}
+
+    @cachebox.cached(cachebox.LRUCache(0), lock=False)
+    def f(x):
+        calls["n"] += 1
+        return x * 2
+
+    assert f(1) == 2  # miss
+    assert f(1) == 2  # hit
+    assert f(2) == 4  # miss
+
+    info = cachebox.get_cached_cache_info(f)
+    print(info)
+
+    assert info.misses == 2
+    assert info.hits == 1
+
+
+@pytest.mark.asyncio
+async def test_ignore_path_should_return_awaited_value():
+    @cachebox.cached(cachebox.LRUCache(0), lock=False)
+    async def f(x):
+        await asyncio.sleep(0)
+        return x * 2
+
+    result = await f(
+        5,
+        cachebox__ignore=True,  # type: ignore
+    )
+    assert not asyncio.iscoroutine(result)
+    assert result == 10
+
+
+def test_locks_should_not_leak_on_exception():
+    @cachebox.cached(cachebox.LRUCache(0))  # lock پیش‌فرض True است
+    def f(x):
+        raise ValueError("boom")
+
+    for i in range(50):
+        with pytest.raises(ValueError):
+            f(i)
+
+    locks_cache = None
+    for cell in f.__closure__:  # type: ignore
+        try:
+            val = cell.cell_contents
+        except ValueError:
+            continue
+
+        if isinstance(val, cachebox.Cache):
+            locks_cache = val
+            break
+
+    assert locks_cache is not None, "We couldn't find locks_cache"
+    print("locks length after 50 failing calls:", len(locks_cache))
+
+    assert len(locks_cache) == 0
+
+
+@pytest.mark.asyncio
+async def test_waiters_should_not_leak_on_cancel():
+    started = asyncio.Event()
+    release = asyncio.Event()
+
+    @cachebox.cached(cachebox.LRUCache(0))
+    async def f(x):
+        started.set()
+        await release.wait()
+        return x
+
+    first = asyncio.create_task(f(1))
+    await started.wait()
+
+    second = asyncio.create_task(f(1))
+    await asyncio.sleep(0.05)
+
+    second.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await second
+
+    release.set()
+    result = await first
+    assert result == 1
+
+    locks_cache = None
+    for cell in f.__closure__:  # type: ignore
+        try:
+            val = cell.cell_contents
+        except ValueError:
+            continue
+        if isinstance(val, cachebox.Cache):
+            locks_cache = val
+            break
+    assert locks_cache is not None
+
+    print("locks length after cancel scenario:", len(locks_cache))
+    assert len(locks_cache) == 0
+
+
+@pytest.mark.asyncio
+async def test_cancelled_error_should_propagate_to_waiters_too():
+    gate = asyncio.Event()
+
+    @cachebox.cached(cachebox.LRUCache(0))
+    async def f(x):
+        await gate.wait()
+        raise asyncio.CancelledError()
+
+    t1 = asyncio.create_task(f(1))
+    t2 = asyncio.create_task(f(1))
+    await asyncio.sleep(0.01)
+    gate.set()
+
+    with pytest.raises(asyncio.CancelledError):
+        await t1
+
+    with pytest.raises(asyncio.CancelledError):
+        await t2