Асинхронные контекстные менеджеры в Python: где они текут

Сервис держит пул на 20 соединений к PostgreSQL. Под нагрузкой в 300 одновременных запросов пул исчерпывается за секунды, хотя каждый обработчик оформлен через async with pool.acquire(). Логи чистые, исключений нет. Причина не в утечке памяти и не в медленной базе — соединения не возвращаются в пул, потому что задачи отменяются по таймауту, а __aexit__ в этот момент ведёт себя не так, как думает автор кода.

Асинхронные контекстные менеджеры выглядят как обычные, но живут в мире, где корутину могут прервать между двумя await. Здесь разбираем именно поведение async with под отменой и при вложенности — синхронные менеджеры и общий синтаксис with оставляем за скобками. Всё проверено на CPython 3.11 и 3.12.

Как async with устроен внутри

Протокол простой: объект реализует __aenter__ и __aexit__, обе — корутины. Компилятор разворачивает async with cm as x: примерно в следующее (схематично):

manager = cm
x = await manager.__aenter__()
try:
    # тело блока
finally:
    await manager.__aexit__(exc_type, exc, tb)

Ключевое слово здесь — await внутри finally. Освобождение ресурса — это не мгновенная операция, а корутина, которой нужен шанс доработать до конца. И вот тут начинаются проблемы, потому что цикл событий может доставить в эту корутину CancelledError.

Отмена задачи бьёт по __aexit__

Когда задачу отменяют (по таймауту asyncio.wait_for, при task.cancel(), при выходе из asyncio.timeout), CancelledError прилетает в ближайшую точку await. Если это await внутри __aexit__, освобождение обрывается на полуслове.

import asyncio

class Connection:
    async def __aenter__(self):
        print("взяли соединение")
        return self

    async def __aexit__(self, exc_type, exc, tb):
        # если здесь произойдёт await и задачу отменят —
        # release может не выполниться
        await asyncio.sleep(0.1)  # имитация сетевого RELEASE
        print("вернули соединение")

async def worker():
    async with Connection():
        await asyncio.sleep(1)

async def main():
    # отменяем через 0.5 секунды — worker ещё в теле блока
    await asyncio.wait_for(worker(), timeout=0.5)

asyncio.run(main())

Здесь при таймауте CancelledError прилетит в await asyncio.sleep(1) внутри тела, затем управление уйдёт в __aexit__, и await asyncio.sleep(0.1) тоже получит отмену. Строка «вернули соединение» может не напечататься. Соединение утекло.

Асинхронный __aexit__ — не гарантия освобождения, а всего лишь ещё одна корутина, которую цикл событий вправе прервать. Ресурс освобождён тогда, когда завершилась освобождающая корутина, а не когда управление вошло в блок finally.

Порядок выхода при вложенности

Несколько менеджеров в одной строке выходят в обратном порядке — как и в синхронном случае:

async with outer() as a, inner() as b:
    ...
# __aexit__ у inner, затем __aexit__ у outer

Проблема в том, что если __aexit__ внешнего менеджера сам делает await и в этот момент прилетает отмена, внутренний уже освобождён, а внешний — нет. Разбираться, что успело закрыться, а что нет, приходится по логам. Поэтому ресурсы с зависимостью (транзакция поверх соединения) лучше не разворачивать в одну строку, а вкладывать явно — так видно, что от чего зависит.

Три способа не потерять ресурс при отмене

Способ 1: shield на освобождение

asyncio.shield защищает корутину от внешней отмены. Оборачиваем внутренности __aexit__:

async def __aexit__(self, exc_type, exc, tb):
    # release доработает, даже если задачу отменяют
    await asyncio.shield(self._release())

async def _release(self):
    await asyncio.sleep(0.1)
    print("вернули соединение")

Цена: если release сам зависнет, shield не даст его прервать — рискуете подвесить выход из блока. Годится только для быстрых, гарантированно завершающихся операций освобождения.

Способ 2: освобождение с собственным таймаутом

Даём release ограниченное время и логируем, если не уложился:

async def __aexit__(self, exc_type, exc, tb):
    try:
        async with asyncio.timeout(2.0):
            await self._release()
    except (asyncio.CancelledError, TimeoutError):
        # release не успел — соединение считаем битым,
        # закрываем принудительно, а не возвращаем в пул
        self._force_close()
        raise

Синтаксис asyncio.timeout доступен с Python 3.11. Цена: усложнение и необходимость иметь _force_close, который работает синхронно и без await.

Способ 3: не делать await в освобождении вообще

Самый надёжный вариант для критичных ресурсов — синхронное освобождение. Пул возвращает соединение через синхронный вызов, а фактический сетевой ping/reset откладывает на следующее использование. Тогда __aexit__ не содержит точек, где может прилететь отмена.

Ещё три места, где async with подводит

• Нельзя вызвать в синхронном коде. async with работает только внутри async def. Попытка использовать асинхронный менеджер через обычный with даст AttributeError: __enter__ — объект не реализует синхронный протокол. Одна из частых ошибок при переносе кода.
• __aenter__ тоже может быть отменён. Если задачу отменяют во время входа в блок, ресурс мог быть частично захвачен, а __aexit__ при этом не вызовется — тело блока не начиналось. Захват ресурса в __aenter__ должен быть атомарным либо откатываемым внутри самого __aenter__.
• @asynccontextmanager и генераторы. При декорировании через contextlib.asynccontextmanager код после yield — это ваш __aexit__. Если между try и yield нет обработки GeneratorExit/CancelledError, тот же сценарий с отменой оставит вас без освобождения.

Как проверить свой менеджер на устойчивость к отмене

1. Оберните вызов блока в asyncio.wait_for с таймауту меньше, чем длительность тела.
2. Добавьте счётчик захваченных и освобождённых ресурсов (или логи).
3. Прогоните 200–500 итераций с отменой и сравните счётчики. Расхождение = утечка при отмене.
4. Повторите с отменой в момент __aenter__: таймаут меньше времени захвата.

Если пул всё равно исчерпывается

Типичные симптомы и что смотреть:

• TimeoutError на pool.acquire() при живой базе — соединения не возвращаются, ищите await в __aexit__ и проверьте поведение под отменой по чек-листу выше.
• Счётчик активных соединений на стороне PostgreSQL (pg_stat_activity) растёт и не падает — release не доходит до сервера.
• В логах есть вход в блок, но нет выхода — __aexit__ прерывается отменой.

Что делать дальше

Проверьте каждый асинхронный менеджер, который держит внешний ресурс (соединение, файл, блокировку, транзакцию), на устойчивость к отмене по чек-листу из четырёх шагов. Для пулов измеряйте разницу между числом acquire и release под нагрузкой с отменами — это самый прямой индикатор проблемы. Поведение отмены заметно менялось между версиями asyncio, поэтому фиксируйте версию Python в требованиях и перечитайте раздел про asyncio.timeout и asyncio.CancelledError в официальной документации перед тем, как править продовый код.