@cached does not coalesce concurrent calls when using Semaphore which leads to cache stampede in first batch.

[0x1618]

When using @cached with a concurrency limit via asyncio.Semaphore, the first batch of tasks all miss the cache and execute the underlying function in parallel. This results in multiple uncached executions instead of a single in-flight call being shared.

import asyncio
import random
from aiocache import cached, caches

caches.set_config({
    "default": {
        "cache": "aiocache.SimpleMemoryCache",
        "ttl": 300
    }
})

async def with_semaphore(semaphore, coro):
    async with semaphore:
        return await coro

@cached()
async def long_job():
    await asyncio.sleep(2)
    return random.randint(0, 1000)

async def main():
    s = asyncio.Semaphore(4)
    tasks = [with_semaphore(s, long_job()) for _ in range(8)]
    print(await asyncio.gather(*tasks))

asyncio.run(main())

output:

[177, 972, 52, 418, 418, 418, 418, 418]

The first caller (within the semaphore limit) should execute the function, and all other concurrent callers should await the same in-flight result, not trigger additional executions.

fix:

from functools import wraps

def _default_key_builder(func, *args, **kwargs):
    return f"{func.__module__}:{func.__qualname__}:{args}:{tuple(sorted(kwargs.items()))}"


def coalesce(key_builder):
    lock: Optional[asyncio.Lock] = None
    inflight: dict[str, asyncio.Task] = {}

    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            nonlocal lock

            if lock is None:
                lock = asyncio.Lock()

            key: str = key_builder(func, *args, **kwargs)

            async with lock:
                task: Optional[asyncio.Task] = inflight.get(key)

                if task is None:
                    task = asyncio.create_task(func(*args, **kwargs))

                    inflight[key] = task

            try:
                return await task
            finally:
                if task.done():
                    async with lock:
                        if inflight.get(key) is task:
                            inflight.pop(key, None)

        return wrapper

    return decorator


@coalesce(key_builder=_default_key_builder)
@cached(alias="default", key_builder=_default_key_builder)
async def long_job():
    await asyncio.sleep(2)
    return random.randint(0, 1000)

[Dreamsorcerer]

After #625 is done, we could maybe add some extra logic here, probably pinching some code from async-lru: https://github.com/aio-libs/async-lru/blob/5f18bfe448a454aaa317d1f67fb204dae0424c12/async_lru/__init__.py#L203-L210