Как ускорить асинхронные генераторы в CPython на 40%, выкинув скрытые исключения

Недавно я гонял бенчмарки на свежей сборке CPython (ветка с tail-call интерпретатором, PGO + LTO, GCC) и снял флеймграф стандартного бенчмарка async_generators из набора pyperformance (рекурсивный обход дерева на 100 000 нод с глубиной вложенности ~17).

Картина на профиле оказалась фееричной: почти треть всего процессорного времени (29%) сжиралась кодом, который вообще не делал никакой полезной работы.

Рантайм создавал, настраивал, а затем тут же уничтожал объекты исключений StopIteration, о существовании которых Python-код даже не догадывался.

Ниже о том, откуда растут ноги у этой проблемы в Си-коде CPython и как пара правок в genobject.c дали ускорение в 1.33x — 1.53x.

Улики: флеймграф

Вот профиль выполнения бенчмарка (Perforator, 11 200 сэмплов):

Если просуммировать inclusive time по функциям обработки исключений (схлопнув рекурсию):

В чем абсурд происходящего в Си?

В CPython асинхронный генератор под капотом использует вспомогательный объект-обёртку PyAsyncGenASend. Когда в Python выполняется async for или await agen._anext__(), вызывается Си-функция PyAsyncGenASendSend().

До патча этот процесс выглядел как перекладывание бумажек через три инстанции:

1. async_gen_asend_send() дёргает нативный gen_send(), забирает значение из генератора и передает его в async_gen_unwrap_value().

2. async_gen_unwrap_value() видит, что генератор сделал yield, и вызывает PyGenSetStopIterationValue(value). В этот момент CPython честно аллоцирует инстанс исключения StopIteration, аллоцирует кортеж args, упаковывает туда наше значение и выставляет ошибку в тред-стейт.

3. Функция уровнем выше (_PyAsyncGenASend_Send()) сразу же вызывает PyGenFetchStopIterationValue(): Она проверяет тип исключения, вытаскивает из него значение обратно, сбрасывает ошибку в треде и… вызывает StopIteration_dealloc, уничтожая только что созданное исключение.

Это исключение никогда не предназначалось для передачи в Python. Оно использовалось тупо как костыльный Си-интерфейс для передачи одного указателя между двумя функциями внутри рантайма.

И если у вас генератор завернут в генератор (например, пайплайн обработки или рекурсия), эта бессмысленная аллокация и очистка происходили на каждом уровне вложенности для каждого элемента.

Что делаем в патче?

Код патча можно разделить на две части:

1. Избавляемся от StopIteration в горячем пути

Мы разделяем распаковку значения на два пути:

• Добавляем Си-хелпер async_gen_unwrap_send(), который возвращает статус через нативный enum PySendResult (PYGEN_RETURN, PYGEN_NEXT, PYGEN_ERROR), а само значение отдает через указатель PyObject **presult. Никаких исключений в треде.

• Переписываем PyAsyncGenASendSend на прямую работу с этим хелпером. Забираем strong reference через Py_NewRef, прибиваем внутреннюю обертку _PyAsyncGenWrappedValue и отдаем результат.

• Для совместимости питонячий метод asend.send() (async_gen_asend_send) оставляем тонкой оберткой, которая при необходимости всё ещё поднимает StopIteration, чтобы не ломать публичное API.

2. Микрооптимизация деаллокатора

Заодно в async_gen_asend_dealloc нашлась еще одна мелкая свинья:

// До:
if (PyObject_CallFinalizerFromDealloc(self)) return;

// После:
if (ags->ags_state == AWAITABLE_STATE_INIT
    && PyObject_CallFinalizerFromDealloc(self))
{
    return;
}

Финализатор у asend объекта нужен ровно для одного: бросить RuntimeWarning, если корутину создали, но ни разу не заэвейтили. Если объект уже крутится в итерации (AWAITABLE_STATE_ITER) или закрыт (CLOSED), звать тяжелый PyObject_CallFinalizerFromDealloc нет никакого смысла — варнингов там быть не может.

Результаты

Бенчмарк async_generators из pyperformance. Тот же коммит, то же железо, единственная разница — патч в genobject.c:

Остальные тесты из pyperformance остались в рамках шума измерений (что ожидаемо, так как патч изолирован в кишках PyAsyncGen). Полный сьют тестов Python (test -j8, 491 тест) проходит без ошибок.

PR заслан в апстрим CPython: gh-158315 (ссылка на PR).

Мораль

Если ваш рантайм на C/C++ использует механизм исключений для регулярной передачи данных в горячем цикле — рано или поздно в профилировщике вы увидите не свою логику, а работу аллокатора. Даже пара простых правок по замене псевдо-исключений на возврат enum-статуса может выкинуть 30% оверхеда на ровном месте.