Python asyncio と ContextVar

Python 3.7 で contextvars モジュールが追加され、asyncio に対応した ContextVar クラスが導入されました。

スレッド毎に固有のデータを持てる Thread Local (Python では threading.local()) のように、コルーチン毎に固有のデータを持てるようになっているのが ContextVar の特徴です。

実際に ContextVar を使ってみて、ContextVar にセットしたデータが消えてしまうような現象に見舞われ、改めて検証コードを書いて挙動を調べてみました。

以下のコードでは、parent_await_coroutine() と parent_create_new_task() の二つのコルーチン関数を実行しており、それぞれの関数内では ContextVar に値をセットし、child() 関数を呼び出しています。
この child() 関数では ContextVar から取り出した値を変更しています。

二つの parent 関数では child 関数の呼び出し方が異なります。前者はコルーチンを await し、後者はコルーチンをラップした新しい Task を生成して実行しています。

後者の新しい Task として実行した場合には、child 関数で行われた ContextVar への変更の一部が parent 関数には反映されていません。具体的に言うと、child 関数で number_var ContextVar の値を加算して再セットしていますが、parent 関数ではその変更が読み取れていません(child 関数を呼ぶ前の状態のまま)。
一方で、msg_var ContextVar の Msg オブジェクトへの変更は parent 関数からも見えています。

これは新しいタスクを生成した際に contextvars の内容がコピーされているためです。PEP 567 からそのことが読み取れます。
このコピー処理で number_var の int の場合は値コピーが行われ、msg_var の Msg オブジェクトは参照コピーが行われているため(つまり Shallow Copy)、上記の挙動になっていると考えられます。

import asyncio
import contextvars


class Msg:
    """単なるテキストの入れものクラス。
    contextvars の Shallow copy を確認するために使います。
    """

    def __init__(self, text: str):
        self._text = text

    @property
    def text(self) -> str:
        return self._text

    @text.setter
    def text(self, val):
        self._text = val


msg_var: contextvars.ContextVar[Msg] = contextvars.ContextVar('msg_var')
number_var: contextvars.ContextVar[int] = contextvars.ContextVar('number_var')


async def child():
    # ContextVar から数値を取得して 1 加算する
    n = number_var.get()
    print(f'child: number={n}')  # child: number=1
    n += 1
    number_var.set(n)

    # ContextVar から Msg オブジェクトを取得してテキストを変更する
    msg = msg_var.get()
    print(f'child: msg="{msg.text}"')  # child: msg="msg created by parent"
    msg.text = 'msg changed by child'

    # この child() を async 関数にするための処理
    await asyncio.sleep(0.1)


async def parent_await_coroutine():
    n = 1
    number_var.set(n)
    m = Msg('msg created by parent')
    msg_var.set(m)
    print(f'parent: number={n}')  # parent: number=1
    print(f'parent: msg="{m.text}"')  # parent: msg="msg created by parent"

    await child()

    n = number_var.get()
    m = msg_var.get()
    print(f'parent: number={n}')  # parent: number=2
    print(f'parent: msg="{m.text}"')  # parent: msg="msg changed by child"


async def parent_create_new_task():
    n = 1
    number_var.set(n)
    m = Msg('msg created by parent')
    msg_var.set(m)
    print(f'parent: number={n}')  # parent: number=1
    print(f'parent: msg="{m.text}"')  # parent: msg="msg created by parent"

    await asyncio.create_task(child())

    n = number_var.get()
    m = msg_var.get()
    print(f'parent: number={n}')  # parent: number=1
    print(f'parent: msg="{m.text}"')  # parent: msg="msg changed by child"


if __name__ == '__main__':
    asyncio.run(parent_create_new_task())
    asyncio.run(parent_await_coroutine())