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())
具體調整
- 在
dispatcher
中,開始啟動協程前,將request_id_context
設定為-1
。 然後在所有的協程呼叫完畢後,再檢視request_context_id
的值。 - 在
outer
中,在設定request_id_context
之前,先檢視它的值。 - 在
inner
中,在檢查和檢視request_id_context
之後,將它修改為其原始值的10倍。 - 定義了一個物件
A
,以及一個用來傳遞A
物件例項的上下文變數obj_context
。 - 在
inner
中,建立A
的例項並儲存到obj_context
中。 - 在
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
觀察執行結果,可以看到對上下文變數的修改,有兩種情況:
- 對於已經設定過值的上下文變數,後續對其做的修改是單向傳播的。儘管每個
outer
方法都request_id_context
設定成了不同的值,但最後在dispatcher
呼叫完所有的outer
後,它取到的request_id_context
仍然為-1
。 同樣,inner
方法雖然修改了request_id_context
,但這個修改對呼叫它的outer
是不可見的。另外一個方向,outer
可以讀取到呼叫它的dispatcher
修改的值,inner
也可以讀取到outer
的修改。 - 如果是新設定的上下文變數,它的值可以傳遞到其所在方法的呼叫者。比如在
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_物件可以讓我們方便在協程間儲存上下文資料。在使用時要注意以下幾點:
- contextvars 對協程的支援是從Python 3.7才開始的,使用時要注意Python版本。
- ContextVar 應當在模組級別定義和建立,一定不能在閉包中定義。
- 儲存在上下文中的變數一定要在使用完成後顯示清理,否則會導致記憶體洩漏。