О полезности contextvars

В Python есть множество возможностей и языковых конструкций. Какие-то мы используем каждый день, а о некоторых даже опытные программисты узнают с удивлением после нескольких лет работы с языком (привет, Ellipsis!). Совсем недавно вышел Python 3.9, но в этой статье я расскажу о функциональности, представленной еще в версии 3.7. На мой взгляд, она совершенно незаслуженно обделена пристальным вниманием. Речь, конечно же, о contextvars.

В ДомКлике огромная кодовая база на асинхронном Python. С уверенностью можно сказать, что это лидирующая компетенция в нашей компании: разработчиков на Python даже больше, чем фронтендеров. Обычно release notes очередной версии пристально изучаются на предмет того, что из новых фич можно будет попробовать. Описание же contextvars, как и примеры, совершенно не впечатлило. Зачем нужно передавать значение между функциями в настолько странно объявленной переменной? Давайте разбираться: рассмотрим несколько способов работы с глобальным контекстом в Python-приложениях.

Глобальные переменные

Старый, как мир, подход, хотя и считающийся ужасным антипаттерном, работает:

a = 0
def x():
    global a
    for i in range(100000):
        a += 1

… до тех пор, пока наше приложение не становится многопоточным. Неcмотря на наличие GIL, инкремент в Python не является атомарной операцией:

import dis; dis.dis(x)
>>>
# Цикл убран для наглядности
14 LOAD_GLOBAL              1 (a)   # Загружаем в стек глобальную переменную a
16 LOAD_CONST               2 (1)   # Загружаем в стек 1
18 INPLACE_ADD                      # Сложение верхних элементов в стеке

Между каждой инструкцией байт-кода может переключиться контекст, что сделает значение переменной некорректной в этом потоке. Проверим:

import threading

threads = []

for j in range(5):
    thread = threading.Thread(target=x)
    threads.append(thread)
    thread.start()

for thread in threads:
    thread.join()

assert a == 500000
>>> AssertionError

Значение a будет плавать от запуска к запуску. На помощь могут прийти примитивы синхронизации (например, RLock) или, в зависимости от задачи, threading.local

Контекстные переменные

Прогресс не стоит на месте, и сейчас Python уверенно поддерживает асинхронные паттерны программирования. Теперь даже в рамках одного процесса нет защиты от переключения контекста выполнения, что приводит к использованию знакомых по многопоточности локов и семафоров. Но как быть с локальным хранилищем? Ведь теперь оно должно быть привязано к каждой вызываемой корутине, и к тому же быть доступным по всему стеку вызовов? Вот здесь на помощь и приходят contextvars, работающие единообразно при любых переключениях контекста:

• Разные потоки.
• Цепочки вызовов асинхронных функций.
• Создание новых задач на event loop (ensure_future / create_task).
• Создание генераторов.

Применение на практике

И всё же, какую конкретно пользу можно из этого извлечь? Рассмотрим цепочку вызовов, которая есть почти в любом микросервисе:



Из сервиса A вызывается сервис B, при этом по цепочке необходимо передать информацию об исходном запросе для трекинга (а service mesh не завезли). Клиент к стороннему сервису — это абстракция, которая может не иметь информации о текущем запросе. Также он может вызываться в отдельной корутине и вообще не иметь доступа к контексту текущего запроса. Можно передавать request_id каждый раз при вызове функции service_client, но расширение передаваемых данных будет затруднительным.

Используем contextvars:

import asyncio
import random
from contextvars import ContextVar

from aiohttp import web

request_id: ContextVar[int] = ContextVar('request_id')


async def perform_external_request():
    # Cозданная задача всегда будет иметь контекст родительской
    await asyncio.sleep(5)
    print('request_id =', request_id.get())
    # Здесь выполняем запрос к стороннему сервису


async def test_handler(request):
    r = random.randint(1, 100)
    request_id.set(r)
    asyncio.ensure_future(perform_external_request())
    return web.Response(text='ok')


app = web.Application()
app.router.add_route('GET', '/test', test_handler)
web.run_app(app, port=8000)

Так удобно хранить данные, определяющие контекст вызова: информацию о пользователе, метрики времени ответа и другие. Например, в логах:

import uuid
import logging
from contextvars import ContextVar

from aiohttp import web

request_id: ContextVar[str] = ContextVar('request_id')


class RequestIdFilter(logging.Filter):

    def filter(self, record):
        # Добавление нужного поля в запись
        record.request_id = request_id.get()
        return True


logger = logging.getLogger(__name__)
ch = logging.StreamHandler()
# Все сообщения от этого логгера будут иметь текущий X-Request-Id, 
# вне зависимости от места вызова!
ch.setFormatter(logging.Formatter('%(request_id)s: %(message)s'))
logger.addFilter(RequestIdFilter())
logger.addHandler(ch)


async def test_handler(request):

    logger.warning('Calling test handler')

    return web.Response(text='OK')


@web.middleware
async def request_id_middleware(request, handler):
    # Установка / чтение request_id
    request_id.set(request.headers.get('X-Request-Id', str(uuid.uuid4())))
    response = await handler(request)
    return response

app = web.Application(middlewares=[request_id_middleware])
app.router.add_route('GET', '/test', test_handler)
web.run_app(app, port=8000)

Что еще полезно знать

contextvars — это одна из немногих возможностей языка, для знакомства с которой мне пришлось глубоко погрузиться в соответствующий PEP из-за весьма скудной основной документации. Например, переменные контекста весьма интересно ведут себя с генераторами. Правила следующие:

• Изменения «внутри» генератора не видны в вызывающем коде.
• Переменная не может быть изменена между итерациями генератора.
• Изменения «снаружи» видны «внутри», если они не были изменены «внутри».

Я оставил комментарии на основе примера из исходного PEP:

var1 = contextvars.ContextVar('var1')
var2 = contextvars.ContextVar('var2')

def gen():
    var1.set('gen')
    assert var1.get() == 'gen'
    assert var2.get() == 'main'
    yield 1

    # Это изменение не будет применено, так как между итерациями модификации запрещены 
    var1.set('genXXXX')

    # var1 модифицируется снаружи, но внутри генератора изменение не видно, 
    # так как в нем эта переменная была изменена
    assert var1.get() == 'gen'

    # var2 меняется "снаружи" без изменения "внутри", поэтому оно доступно
    assert var2.get() == 'main modified'
    yield 2

def main():
    g = gen()

    var1.set('main')
    var2.set('main')
    next(g)

    # Модификация "изнутри" не доступна "снаружи"
    assert var1.get() == 'main'

    var1.set('main modified')
    var2.set('main modified')
    next(g)

По аналогии с генераторами, для корутин тоже действуют некоторые правила:

• Если одна функция ожидает другую через await, то изменения переменной видны и в «родительской», и в «дочерней».
• Если одна функция вызвала другую через создание задачи (ensure_future / create_task), то изменения переменной между ними не передаются.

Вы еще не обновились до 3.7?

Похожий функционал предоставляет библиотека aiotask-context. Она работает медленнее, чем нативная реализация в 3.7, а также требует дополнительной инициализации:

import asyncio
import aiotask_context as context

async def test():
    print(context.get('some_data', default='not set'))

loop = asyncio.get_event_loop()
loop.set_task_factory(context.task_factory)
loop.run_until_complete(test())

Заключение

contextvars — это не фича, которую нужно брать в каждый проект. Однако она способна сделать код значительно проще и чище, если правильно проектировать архитектуру сервиса.