Python協程中使用上下文

在Python 3.7中，asyncio 協程加入了對上下文的支持。使用上下文就能夠在一些場景下隱式地傳遞變量，好比數據庫鏈接session等，而不須要在全部方法調用顯示地傳遞這些變量。使用得當的話，能夠提升接口的可讀性和擴展性。

基本使用方式

協和的上下文是經過 contextvars 中的 ContextVar 對象來管理的。最基本的使用方式是在某一調用層次中設置上下文，而後在後續調用中使用。以下例所示：

import asyncio
import contextvars
from random import randint
from unittest import TestCase

request_id_context = contextvars.ContextVar('request-id')


async def inner(x):
    request_id = request_id_context.get()
    if request_id != x:
        raise AssertionError('request_id %d from context does NOT equal with parameter x %d' % (request_id, x))

    print('start handling inner request-%d, with x: %d' % (request_id, x))
    await asyncio.sleep(randint(0, 3))
    print('finish handling inner request-%d, with x: %d' % (request_id, x))


async def outer(i):
    print('start handling outer request-%d' % i)
    request_id_context.set(i)
    await inner(i)
    print('finish handling outer request-%d with request_id in context %d' % (i, request_id_context.get()))


async def dispatcher():
    await asyncio.gather(*[
        outer(i) for i in range(0, 10)
    ])


class ContextTest(TestCase):

    def test(self):
        asyncio.run(dispatcher())

上例中，在最後定義了一個單元測試用例對象 ContextTest 。它的方法 test 是程序的入口，使用 asyncio.run 方法來在協程中執行被測試的異步方法 dispatcher 。dispatcher 則併發啓動10個異步方法 outer 。 outer方法首先將在模塊層定義的上下文變量 request_id_context 設置爲當前調用指定的值，這個值對於每一個 outer 的調用都是不一樣的。 而後在後續被調用的 inner 方法，以及 outer 方法內部訪問了這個上下文變動。在 inner 方法內容，則比較了顯示傳入的 i 和從上下文變量中取出的 request_id 。

測試用例的執行結果以下：

start handling outer request-0
start handling inner request-0, with x: 0
start handling outer request-1
start handling inner request-1, with x: 1
start handling outer request-2
start handling inner request-2, with x: 2
start handling outer request-3
start handling inner request-3, with x: 3
start handling outer request-4
start handling inner request-4, with x: 4
start handling outer request-5
start handling inner request-5, with x: 5
start handling outer request-6
start handling inner request-6, with x: 6
start handling outer request-7
start handling inner request-7, with x: 7
start handling outer request-8
start handling inner request-8, with x: 8
start handling outer request-9
start handling inner request-9, with x: 9
finish handling inner request-3, with x: 3
finish handling outer request-3 with request_id in context 3
finish handling inner request-7, with x: 7
finish handling outer request-7 with request_id in context 7
finish handling inner request-1, with x: 1
finish handling outer request-1 with request_id in context 1
finish handling inner request-4, with x: 4
finish handling outer request-4 with request_id in context 4
finish handling inner request-5, with x: 5
finish handling outer request-5 with request_id in context 5
finish handling inner request-9, with x: 9
finish handling outer request-9 with request_id in context 9
finish handling inner request-0, with x: 0
finish handling outer request-0 with request_id in context 0
finish handling inner request-2, with x: 2
finish handling outer request-2 with request_id in context 2
finish handling inner request-6, with x: 6
finish handling outer request-6 with request_id in context 6
finish handling inner request-8, with x: 8
finish handling outer request-8 with request_id in context 8

能夠看到，雖然每次 outer 方法對模塊層同定義的同一個上下文變量 request_id_context 設置了不一樣的值，但後續併發訪問相互之間並不會混淆或衝突。

不一樣調用層次間對上下文的修改

前一節展現了在設置了上下文變量後，在後續使用中讀取這個變量的狀況。這一節，咱們看一下不用調用層次間對同一個上下文變量進行修改的狀況。

在上一節代碼上作了一些調整後以下：

import asyncio
import contextvars
from random import randint
from unittest import TestCase

request_id_context = contextvars.ContextVar('request-id')

obj_context = contextvars.ContextVar('obj')


class A(object):

    def __init__(self, x):
        self.x = x

    def __repr__(self):
        return '<A|x: %d>' % self.x


async def inner(x):
    request_id = request_id_context.get()
    if request_id != x:
        raise AssertionError('request_id %d from context does NOT equal with parameter x %d' % (request_id, x))

    print('start handling inner request-%d, with x: %d' % (request_id, x))
    request_id_context.set(request_id * 10)
    await asyncio.sleep(randint(0, 3))

    obj = A(x)
    obj_context.set(obj)
    print('finish handling inner request-%d, with x: %d' % (request_id, x))


async def outer(i):
    print('start handling outer request-%d with request_id in context %d' % (i, request_id_context.get()))
    request_id_context.set(i)
    await inner(i)
    print('obj: %s in outer request-%d' % (obj_context.get(), i))
    print('finish handling outer request-%d with request_id in context %d' % (i, request_id_context.get()))


async def dispatcher():
    request_id_context.set(-1)
    await asyncio.gather(*[
        outer(i) for i in range(0, 10)
    ])
    print('finish all coroutines with request_id in context: %d' % (request_id_context.get()))


class ContextTest(TestCase):

    def test(self):
        asyncio.run(dispatcher())

具體調整

1. 在 dispatcher 中，開始啓動協程前，將 request_id_context 設置爲 -1 。 而後在全部的協程調用完畢後，再查看 request_context_id 的值。
2. 在 outer 中，在設置 request_id_context 以前，先查看它的值。
3. 在 inner 中，在檢查和查看 request_id_context 以後，將它修改成其原始值的10倍。
4. 定義了一個對象 A ，以及一個用來傳遞 A 對象實例的上下文變量 obj_context。
5. 在 inner 中，建立A的實例並保存到obj_context中。
6. 在outer中，調用完inner方法後，查看obj_context上下文變量。

代碼的執行結果以下：

start handling outer request-0 with request_id in context -1
start handling inner request-0, with x: 0
start handling outer request-1 with request_id in context -1
start handling inner request-1, with x: 1
start handling outer request-2 with request_id in context -1
start handling inner request-2, with x: 2
start handling outer request-3 with request_id in context -1
start handling inner request-3, with x: 3
start handling outer request-4 with request_id in context -1
start handling inner request-4, with x: 4
start handling outer request-5 with request_id in context -1
start handling inner request-5, with x: 5
start handling outer request-6 with request_id in context -1
start handling inner request-6, with x: 6
start handling outer request-7 with request_id in context -1
start handling inner request-7, with x: 7
start handling outer request-8 with request_id in context -1
start handling inner request-8, with x: 8
start handling outer request-9 with request_id in context -1
start handling inner request-9, with x: 9
finish handling inner request-6, with x: 6
obj: <A|x: 6> in outer request-6
finish handling outer request-6 with request_id in context 60
finish handling inner request-0, with x: 0
obj: <A|x: 0> in outer request-0
finish handling outer request-0 with request_id in context 0
finish handling inner request-2, with x: 2
obj: <A|x: 2> in outer request-2
finish handling outer request-2 with request_id in context 20
finish handling inner request-3, with x: 3
obj: <A|x: 3> in outer request-3
finish handling outer request-3 with request_id in context 30
finish handling inner request-5, with x: 5
obj: <A|x: 5> in outer request-5
finish handling outer request-5 with request_id in context 50
finish handling inner request-7, with x: 7
obj: <A|x: 7> in outer request-7
finish handling outer request-7 with request_id in context 70
finish handling inner request-8, with x: 8
obj: <A|x: 8> in outer request-8
finish handling outer request-8 with request_id in context 80
finish handling inner request-9, with x: 9
obj: <A|x: 9> in outer request-9
finish handling outer request-9 with request_id in context 90
finish handling inner request-1, with x: 1
obj: <A|x: 1> in outer request-1
finish handling outer request-1 with request_id in context 10
finish handling inner request-4, with x: 4
obj: <A|x: 4> in outer request-4
finish handling outer request-4 with request_id in context 40
finish all coroutines with request_id in context: -1

觀察執行結果，能夠看到對上下文變量的修改，有兩種狀況：

1. 對於已經設置過值的上下文變量，後續對其作的修改是單向傳播的。儘管每一個 outer 方法都 request_id_context 設置成了不一樣的值，但最後在 dispatcher 調用完全部的 outer 後，它取到的 request_id_context 仍然爲 -1。 一樣，inner方法雖然修改了request_id_context，但這個修改對調用它的outer是不可見的。另一個方向，outer能夠讀取到調用它的dispatcher修改的值，inner也能夠讀取到outer的修改。
2. 若是是新設置的上下文變量，它的值能夠傳遞到其所在方法的調用者。好比在inner中設置的obj_context，在outer中能夠讀取。

內存泄漏和上下文清理

根據Python文檔， ContextVar對象會持有變量值的強引用，因此若是沒有適當清理，會致使內存漏泄。咱們使用如下代碼演示這種問題。

import asyncio
import contextvars
from unittest import TestCase
import weakref

obj_context = contextvars.ContextVar('obj')
obj_ref_dict = {}


class A(object):

    def __init__(self, x):
        self.x = x

    def __repr__(self):
        return '<A|x: %d>' % self.x


async def inner(x):
    obj = A(x)
    obj_context.set(obj)
    obj_ref_dict[x] = weakref.ref(obj)


async def outer(i):
    await inner(i)
    print('obj: %s in outer request-%d from obj_ref_dict' % (obj_ref_dict[i](), i))


async def dispatcher():
    await asyncio.gather(*[
        outer(i) for i in range(0, 10)
    ])
    for i in range(0, 10):
        print('obj-%d: %s in obj_ref_dict' % (i, obj_ref_dict[i]()))


class ContextTest(TestCase):

    def test(self):
        asyncio.run(dispatcher())

和上一節中的代碼同樣，inner方法在調用棧的最內部設置了上下文變量obj_context。不一樣的是，在設置上下文的同時，也將保存在上下文中的對象A的實例保存到一個弱引用中，以便後續經過弱引用來檢查對象實例是否被回收。

代碼的執行結果以下：

obj: <A|x: 0> in outer request-0 from obj_ref_dict
obj: <A|x: 1> in outer request-1 from obj_ref_dict
obj: <A|x: 2> in outer request-2 from obj_ref_dict
obj: <A|x: 3> in outer request-3 from obj_ref_dict
obj: <A|x: 4> in outer request-4 from obj_ref_dict
obj: <A|x: 5> in outer request-5 from obj_ref_dict
obj: <A|x: 6> in outer request-6 from obj_ref_dict
obj: <A|x: 7> in outer request-7 from obj_ref_dict
obj: <A|x: 8> in outer request-8 from obj_ref_dict
obj: <A|x: 9> in outer request-9 from obj_ref_dict
obj-0: <A|x: 0> in obj_ref_dict
obj-1: <A|x: 1> in obj_ref_dict
obj-2: <A|x: 2> in obj_ref_dict
obj-3: <A|x: 3> in obj_ref_dict
obj-4: <A|x: 4> in obj_ref_dict
obj-5: <A|x: 5> in obj_ref_dict
obj-6: <A|x: 6> in obj_ref_dict
obj-7: <A|x: 7> in obj_ref_dict
obj-8: <A|x: 8> in obj_ref_dict
obj-9: <A|x: 9> in obj_ref_dict

能夠看到，不管是在outer中，仍是在dispatcher中，全部inner方法保存的上下文變量都被沒有被回收。因此咱們必需在使用完上下文變量後，顯示清理上下文，不然會致使內存泄漏。

這裏，咱們在inner方法的最後，將obj_context設置爲None，就能夠保證不會由於上下文而致使內存不會被回收：

async def inner(x):
    obj = A(x)
    obj_context.set(obj)
    obj_ref_dict[x] = weakref.ref(obj)
    obj_context.set(None)

修改後的代碼執行結果以下：

obj: None in outer request-0 from obj_ref_dict
obj: None in outer request-1 from obj_ref_dict
obj: None in outer request-2 from obj_ref_dict
obj: None in outer request-3 from obj_ref_dict
obj: None in outer request-4 from obj_ref_dict
obj: None in outer request-5 from obj_ref_dict
obj: None in outer request-6 from obj_ref_dict
obj: None in outer request-7 from obj_ref_dict
obj: None in outer request-8 from obj_ref_dict
obj: None in outer request-9 from obj_ref_dict
obj-0: None in obj_ref_dict
obj-1: None in obj_ref_dict
obj-2: None in obj_ref_dict
obj-3: None in obj_ref_dict
obj-4: None in obj_ref_dict
obj-5: None in obj_ref_dict
obj-6: None in obj_ref_dict
obj-7: None in obj_ref_dict
obj-8: None in obj_ref_dict
obj-9: None in obj_ref_dict

能夠看到，當outer和dispatcher嘗試經過弱引用來訪問曾經保存在上下文中的對象實例時，這些對象都已經被回收了。

總結

在協程中使用 contextvars 模塊中的_ContextVar_對象可讓咱們方便在協程間保存上下文數據。在使用時要注意如下幾點：

1. contextvars 對協程的支持是從Python 3.7纔開始的，使用時要注意Python版本。
2. ContextVar 應當在模塊級別定義和建立，必定不能在閉包中定義。
3. 保存在上下文中的變量必定要在使用完成後顯示清理，不然會致使內存泄漏。

參考資料

• https://docs.python.org/3/library/contextvars.html#asyncio-support
• https://docs.python.org/3/library/contextvars.html#contextvars.ContextVar