Connection returned to pool in inconsistent state when `asyncio.CancelledError` occurs

[medikos]

I encountered an issue in the acquire decorator where a connection may be returned to the pool in an inconsistent state.

This happens because asyncio.CancelledError inherits from BaseException, not from Exception. As a result, it is not caught by the except Exception block.

Relevant code:
https://github.com/aio-libs/aiomcache/blob/master/aiomcache/client.py#L28

Because of this, when a coroutine is cancelled, conn[0].set_exception(exc) is not called before self._pool.release(conn).

As a result, a broken connection can be returned to the pool and reused later, which may lead to errors like:

ValueError: Separator is not found, and chunk exceed the limit

This can be fixed by explicitly handling asyncio.CancelledError and marking the connection as failed before returning it to the pool.

Example:

def acquire(
    func: Callable[Concatenate[_Client, Connection, _P], Awaitable[_T]]
) -> Callable[Concatenate[_Client, _P], Awaitable[_T]]:

    @functools.wraps(func)
    async def wrapper(self: _Client, *args: _P.args,  # type: ignore[misc]
                      **kwargs: _P.kwargs) -> _T:
        conn = await self._pool.acquire()
        try:
            return await func(self, conn, *args, **kwargs)
        except Exception as exc:
            conn[0].set_exception(exc)
            raise
        except asyncio.CancelledError as exc:
            # Got CancelledError from client code.
            conn[0].set_exception(exc)
            raise
        finally:
            self._pool.release(conn)

    return wrapper

[Dreamsorcerer]

Sounds like you've already gone through the code. Do you think you could create a PR with a regression test?

[medikos]

Thanks for the quick response! I’ll prepare a PR with a fix and a regression test.