Async TTL Cache with Stampede Lock

import asyncio
import time
from collections.abc import Awaitable, Callable
from typing import Any


class AsyncTTLCache:
    def __init__(self, ttl: float, clock: Callable[[], float] = time.monotonic) -> None:
        self._ttl = ttl
        self._clock = clock
        self._values: dict[str, tuple[float, Any]] = {}
        self._locks: dict[str, asyncio.Lock] = {}

    def _fresh(self, key: str) -> Any | None:
        item = self._values.get(key)
        if item is None:
            return None
        expires_at, value = item
        if self._clock() >= expires_at:
            self._values.pop(key, None)
            return None
        return value

    async def get_or_load(self, key: str, loader: Callable[[], Awaitable[Any]]) -> Any:
        value = self._fresh(key)
        if value is not None:
            return value
        lock = self._locks.setdefault(key, asyncio.Lock())
        async with lock:
            value = self._fresh(key)
            if value is not None:
                return value
            value = await loader()
            self._values[key] = (self._clock() + self._ttl, value)
            return value

    def invalidate(self, key: str) -> None:
        self._values.pop(key, None)


async def main() -> None:
    cache = AsyncTTLCache(ttl=0.2)
    calls = 0

    async def load() -> dict[str, str]:
        nonlocal calls
        calls += 1
        await asyncio.sleep(0.05)
        return {"id": "42", "name": "Ada"}

    first, second = await asyncio.gather(
        cache.get_or_load("user:42", load),
        cache.get_or_load("user:42", load),
    )
    print("same:", first == second, "loader calls:", calls)
    await asyncio.sleep(0.25)
    await cache.get_or_load("user:42", load)
    print("after expiry:", calls)


asyncio.run(main())