Как правильно остановить корутину

Корутина работает благодаря тому, что event loop раз за разом вызывает её метод coroutine.send(…). Если перестать вызывать send, то и корутина перестанет работать. Это простой и изящный способ остановить корутину, но не все так просто.

Рассмотрим пример. У нас есть компьютерная игра, которая хранит в базе данных достижения игрока. Функция работает всю игру, и каждые 5 секунд сохраняет в БД список достижений:

async def save_achivements(user_id):
    connection = open_connection()

    while True:
        achivements = get_achivements(user_id)
        await connection.send(key=user_id, value=achivements)
        await asyncio.sleep(5)

    connection.close()

С этим кодом есть проблема. Если случится ошибка внутри while True:, то до закрытия соединения с БД дело так и не дойдет. Возникнет исключение, нормальный поток исполнения команд прервется и connection.close() никто не вызовет. От такой проблемы спасает try finally:

async def save_achivements(user_id):
    connection = open_connection()

    try:
        while True:
            achivements = get_achivements(user_id)
            await connection.send(key=user_id, value=achivements)
            await asyncio.sleep(5)
    finally:
        connection.close()

Теперь в случае ошибки сработает код внутри finally и соединение будет закрыто вызовом connection.close(), ни смотря ни на что. Это победа!

А что произойдет, если корутину save_achivements(…) просто перестанут вызывать? Игрок прервал игру, корутина save_achivements(…) перестала быть нужной, поэтому её выкинули из event loop и соединение осталось незакрытым. Простой и элегантный способ остановки корутины сломал нам try finally.

Из-за этой проблемы у корутин есть еще один метод — coroutine.throw(…). Он принимает объект исключения — coroutine.throw(exc_obj) — и забрасывает его внутрь корутины. Исключение всплывает по стеку вызовов корутин, как в обычных синхронных функциях.

CancelledError

Event loop, встроенный в библиотеку asyncio никогда не прерывает работу корутин, вместо этого он просит их остановиться самим. Его просьба — это исключение CancelledError. Его можно перехватить и обработать, как обычное исключение. Воспользуемся этим фактом, чтобы перед выходом из игры записать в БД дату последнего сохранения:

import time
import asyncio

async def save_achivements(user_id):
    connection = open_connection()

    try:
        while True:
            achivements = get_achivements(user_id)
            await connection.send(key=user_id, value=achivements)
            await asyncio.sleep(5)

    except asyncio.CancelledError:
        timestamp = time.time()
        await connection.send(key=user_id, value=achivements)
        await connection.send(key='last_update', value=timestamp)
        
        # отпускаем перехваченный CancelledError
        raise
    finally:
        connection.close()

Внутри asyncio even loop происходит примерно следующее:

coroutine = save_achivements(user_id)
coroutine.send(...)
coroutine.send(...)
...

coroutine.throw(asyncio.CancelledError())
coroutine.send(...)
coroutine.send(...)
...

Обратите внимание, что CancelledError не заблокировал работу event loop. Метод coroutine.send(…) продолжает вызываться, авейты исправно работают, а значит, завершение работы корутины тоже может быть асинхронным с вызовами await connection.send(…).

Никто не может остановить корутину, пока та сама не пожелает. Внешний код будет терпеливо ждать, когда же корутина закончит работу и, чтобы он узнал об этом, в конце отпускаем перехваченный CancelledError. Event loop asyncio поймает это исключение и уведомит всех, кто ждал этого завершения. Подробнее читайте в документации по asyncio.Task.

Остановить корутину может только сама корутина

StopIteration vs CancelledError

StopIteration случается, когда корутина уже дошла до return и успела завершить свою работу. Это исключение нельзя перехватить внутри корутины, оно существует только снаружи.

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

CancelledError просит об остановке, а StopIteration — её констатирует