Caching co-routines when working with asyncio.gather

[imxdn]

I was trying out @cached decorator for a simple experiment where I am invoking a co-routine with asyncio.gather(...). However, I could notice that in such scenarios, probably because the cache is not populated with the result, the coroutine is being called multiple times even with the same argument.

import asyncio
import logging

from aiocache import cached

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@cached()
async def my_async_function(arg1, arg2):
    logger.info(f"Invoked with {arg1} and {arg2}")
    await asyncio.sleep(2)
    return arg1 * arg2


async def main():
    tasks = [
        my_async_function(2, 3),
        my_async_function(2, 3),  # This should be cached
        my_async_function(4, 5),
        my_async_function(6, 7),
    ]
    results = await asyncio.gather(*tasks)
    print(results)

asyncio.run(main())

Output:

INFO:__main__:Invoked with 2 and 3
INFO:__main__:Invoked with 2 and 3
INFO:__main__:Invoked with 4 and 5
INFO:__main__:Invoked with 6 and 7
[6, 6, 20, 42]

I was researching around with approach of saving the unawaited result of coroutine and found out that while we cannot await the same coroutine multiple times, we can get around the limitation by creating a task using asyncio.create_task(...).

import asyncio
import logging

from cachetools import LFUCache
from functools import wraps

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class cached:

    def __init__(self):
        self.cache = LFUCache(maxsize=500)

    def __call__(self, func):

        @wraps(func)
        async def wrapper(*args, **kwargs):
            key = (args, tuple(sorted(kwargs.items())))
            if key not in self.cache:
                self.cache[key] = asyncio.create_task(func(*args, **kwargs))
            return await self.cache[key]

        return wrapper

@cached()
async def my_async_function(arg1, arg2):
    logger.info(f"Invoked with {arg1} and {arg2}")
    await asyncio.sleep(2)
    return arg1 * arg2


async def main():
    tasks = [
        my_async_function(2, 3),
        my_async_function(2, 3),  # This will be cached
        my_async_function(4, 5),
        my_async_function(6, 7),
    ]
    results = await asyncio.gather(*tasks)
    print(results)

asyncio.run(main())

Output:

INFO:__main__:Invoked with 2 and 3
INFO:__main__:Invoked with 4 and 5
INFO:__main__:Invoked with 6 and 7
[6, 6, 20, 42]

I tried to go over this repository and could only find one example with asyncio.gather(...) which was using locks instead. I would like to understand what am I missing with my approach and why is it not implemented with aiocache.cached decorator? Are there scenarios where the approach I wrote above fails that I am not looking at?

[Dreamsorcerer]

aiocache doesn't do that. Not sure if that's a deliberate choice or not, but I suspect it's of less relevance for something using an external cache.
async-lru does handle this though, so one possibility might be to combine them and use both decorators.