contextvars & async REPL

貓橘毛 aka Lanfon
8 min readMar 7, 2021

--

踩雷王是我…QQ
總之想在 python -m asyncio 測一下 aiogram 的各種 API 就踩到 contextvars 在 async REPL 的問題…用 aioconsole 測了一下發現是 asyncio 提供的 REPL 本身的問題,總之發了個 PR ,但修過之後在某些 case 下還是有問題 XDDD ((但這真的不是我沒修好是就現行的 API 我修不到…總之等等解釋!!

先把 REPL 的問題放一邊,先討論一下 contextvars (new in 3.7 via PEP567),stackoverflow 上面有一篇給的 code 還蠻有趣的,可以參考一下。(如果是找中文的…看了幾篇之後我寧可跑去看 source code 了ㄅ欠)

btw ,會摸到 threading State 所以 library 是直接實作在 C 那層,Python 版的可以參考 MagicStack/contextvars,CPython 的 PR 在 #5027

好的,那接下來建議參考 python 版的 code 比較容易理解(?)

可用的 APIs 什麼的就不說了(文件有寫稍微瞄一下吧), ContextVar 在同步的情況下,除了 API 不太一樣外,基本上和 Thread-Local Storage (threading.local()) 沒什麼太大的差別,就 threading 目前支援的情況,如果要在 threads 之間「共用」,只能把 target function 外再包一層context.run ,像是 PEP567 Offloading execution to other threads 的範例:

executor = ThreadPoolExecutor()
current_context = contextvars.copy_context()

executor.submit(current_context.run, some_function)

好的,回過頭來說說所謂「共用」的部份, contextvarsthread-local 比較大的差別是你沒辦法摸到目前的 Context A,只能拿到複本 B。

在 MainThread 的情況下,所有的 ContextVar 預設是寫到 Context A ,以上面的例子「共用」的情況, write 的行為只會發生在複本 B,也就是說 contextvar.set() 的行為都不會影響到 Context A 。(目前的 API 沒辦法拿到 ContextA)

當然有幾個常見的「例外」就是 immutability 的問題,這個在 thread-local 也一樣就沒啥好解釋的了…(然後請不要想著把它用在 multiprocessing 上好嗎..它本來就限制是 thread scope 的東西了…)

接下來就是 contextvars 主要影響到的部份 — asyncio ,在 3.7 的 what’s new 是這麼說的:

asyncio gained support for contextvars. loop.call_soon(), loop.call_soon_threadsafe(), loop.call_later(), loop.call_at(), and Future.add_done_callback() have a new optional keyword-only context parameter. Tasks now track their context automatically. See PEP 567 for more details. (Contributed by Yury Selivanov in bpo-32436.)

主要其實就一件事: copy_context() will be called automatically (if no context propagated).

白話點說就是,如果你用 .call_{at|soon|later|call_soon_threadsafe} 或是 .add_done_callback 的時候沒有指定 context 的話,default 會用 copy_context() 拿複本B 來執行;Tasks 的部份目前沒有辦法指定 context ,所以全部會產生 Task instance 的行為都不會影響到外層的 context

會產生 Tasks 的行為包括像 asyncio 下的 ensure_future, gather, wait_for, shield…etc ,基本上要丟 coroutine or future 當參數的,實作上都會用到 ensure_future (或是 loop.create_task, 基本上一樣…)。

只能說這用在 library 上其實蠻容易會有因為設計而產生的 bug ,例如 encode/databases 就有個 issue 在講這個 XD

這也是為什麼上面那個 stackoverflow 裡面的範例,拿到的 id 跟傳進去的一樣。完整一點的 example 像下面 gist 的例子,也可以直接開 repl.it 來玩看看。

簡單解釋一下幾個部份:

Context in Executor

Executor 的實作是開 n 個 workers 用 queue.get 的方式去跑傳進去的 functions ,所以每個 Thread 內的 Context 都是獨立被上一個 function 用過之後的狀態。(突然覺得有點可憐XD)

Context in Main()

asyncio.run 等同於 loop.run_until_complete ,傳進去的 coroutine 會被 ensure_future 包成 future 執行,所以在 main() 內的 Context 就已經是MainThread 的複本B 了。

Context in some_outer_{…}()

兩個 async functions 被 scheduled 的方式是透過 ensure_future ,所以目前的 Context 已經是複本B 的複本C 了。

Context in inner_coroutine()

這邊就比較有趣了,inner_coroutine()outer_with_await() 的情況會是共用 Context (複本C);在 outer_coroutine() 的情況下,因為先被 ensure_future wrapped 過,所以內部的 Context 會是複本D 。

就目前來說,task instance 沒辦法傳 Context 進去,所以 inner_coroutine() 內的情況其實蠻容易踩到的(畢竟沒辦法限制 caller 不要先 schedule 再 await..)

但老實說就算 tasks 可以傳 context 了,就目前 contextvars 的 API 並沒有辦法拿到 current context ,所以到頭來還是得用複本(copy_context())….

至於為什麼 await coroawait task 為什麼會是不同的 context ,我是覺得設計上其實蠻合理的, task 在被 scheduled 的想法上就是開另一個 mini-thread 同步在執行,就 thread-local 的概念上,會有 context 的副本還蠻正常的(?)

但這就又有另一個問題是,就問前 contextvars 在設計上,每個 thread 是獨立的而且不會用 parent thread 的複本(雖然有一個 PR 在改這個),那 schedule sub-task 的時候到底….???

回過頭來說說我發的那個 PR 為什麼沒完整修好 async REPL 的細節…

async REPL 的實作是開一個 REPLThread 來讀 input ,然後 MainThread 負責跑 event loop ,但對 user 來說兩個應該要是同一個 thread 才對,所以 context 需要 sync 。

解決方案就是弄一個 repl_context 讓執行的 code 都跑在這個 context 下就行了!!!

但實際上修不好的部份(真的不是我的問題啊大大)…如果 input 是用 await ... 的方式餵進來的話, REPLThread 會拿到一個 coroutine ,用 loop.create_task 讓 event loop 跑它(ref),在 task instance 會拿到 context 複本的前提下,REPL 的環境不會有像上面 outer & inner 共用 context 的情況出現。(但實際上,如果直接 await coro 的 context 應該是要 share 的…)

aioconsole 反而沒有這個問題,稍微瞄了一下好像是因為本質上就是在同一個 thread 的關係…

雜七雜八, contextvars 雖然有 backport 的版本,但等同於沒有….不管是 MagicStack 的 pure python 版,或是另一個 aiocontextvars ,基本上都沒有辦法 patch 到跟 3.7 之後的版本一樣 — — 因為 loop 的 behavior 不一樣。

混上 uvloop for 3.6 的話就更 WTF 了,不過 3.6 也快 EOL 了,還好還好。

--

--

貓橘毛 aka Lanfon
貓橘毛 aka Lanfon

Written by 貓橘毛 aka Lanfon

知,不知,上;不知,知,病。

No responses yet