skip to content
back to til

A quick note on async cache

til

Async code lets a web application handle other work while it waits for a task to finish. In Python, asyncio lets functions pause and resume this way. Sometimes we need to cache an async function's result so later calls can reuse it. The LRU (least recently used) cache in functools works for regular functions, but async functions return a coroutine, or work to be awaited, before they produce a result. We need a cache that waits for that result before storing it. A decorator lets us add this behavior without changing each function.

Keeping recently used results

First, we need a key that identifies the arguments passed to a function. The key must be hashable so we can use it in a dictionary. An LRU cache removes the least recently used result when it runs out of space. This is a common approach, also used in memory paging.

class HashableKey:
    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs
        self.kwargs.pop("use_cache", None)
 
    def __eq__(self, other: Any) -> bool:
        return hash(self) == hash(other)
 
    def __hash__(self) -> int:
        def _hash(param: Any):
            if isinstance(param, tuple):
                return tuple(map(_hash, param))
            if isinstance(param, dict):
                return tuple(map(_hash, param.items()))
            if hasattr(param, "__dict__"):
                return str(vars(param))
            return str(param)
        return hash(_hash(self.args) + _hash(self.kwargs))

The cache keeps recently read entries at the end of an ordered dictionary:

class LeastRecentlyUsedCache(OrderedDict):
    def __init__(self, max_size: Optional[int] = None, *args, **kwargs):
        self.max_size = max_size
        super().__init__(*args, **kwargs)
 
    def __getitem__(self, key: Any) -> Any:
        value = super().__getitem__(key)
        self.move_to_end(key)
        return value
 
    def __setitem__(self, key: Any, value: Any):
        super().__setitem__(key, value)
        if self.max_size and len(self) > self.max_size:
            # pop the least recently used result
            oldest = next(iter(self))
            del self[oldest]

Expiring old results

A time-to-live (TTL) cache gives each result an expiration time. We can extend the LRU cache to check that time before reusing a result.

class TimeToLiveCache(LeastRecentlyUsedCache):
    def __init__(
        self, ttl: Optional[int] = None, max_size: Optional[int] = None
    ) -> None:
        super().__init__(max_size=max_size)
 
        self.ttl = datetime.timedelta(seconds=ttl) if ttl else None
 
    def __contains__(self, key: Any) -> bool:
        if key not in self.keys():
            return False
        expiration = super().__getitem__(key)[1]
        if expiration and expiration < datetime.datetime.now():
            del self[key]
            return False
        return True
 
    def __getitem__(self, key: Any) -> Any:
        value = super().__getitem__(key)[0]
        return value
 
    def __setitem__(self, key: Any, value: Any) -> None:
        ttl = (datetime.datetime.now() + self.ttl) if self.ttl else None
        super().__setitem__(key, (value, ttl))

Wrapping an async function

This decorator uses the key and cache classes above. It checks the cache first, then waits for the function's result if needed.

class async_lru_cache:
    def __init__(self, max_size: int = 128):
        self.lru = LeastRecentlyUsedCache(max_size=max_size)
 
    def __call__(self, func):
        async def wrapper(*args, **kwargs):
            key = HashableKey(*args, **kwargs)
            if key in self.lru:
                return self.lru[key]
            self.lru[key] = await func(*args, **kwargs)
            return self.lru[key]
        return wrapper

For a TTL version, replace self.lru with an instance of the TTL cache above. Apply the decorator to an async function like this:

class UserService:
    def __init__(self):
        ...
 
    @async_lru_cache()
    async def get_user(self, id: int) -> User:
        ...

comments