《asyncio 系列》14. asyncio 的一些高级用法
楔子
到现在我们已经了解了 asyncio 提供的大部分功能，使用前面文章中介绍的 asyncio 模块，你应该能完成几乎所有你需要完成的任务。但对于更高级的使用场景，你可能还需要使用一些鲜为人知的技术，在设计自己的 asyncio API 时尤其如此。
在本篇文章中，我们将学习 asyncio 中的更高级技术。比如如何设计可以同时处理协程和常规 Python 函数的 API，如何强制事件循环的迭代，以及如何在不传递参数的情况下，在任务之间传递状态。我们还将深入了解 asyncio 究竟如何使用生成器来充分了解幕后发生的事情，为此将实现定制的可等待对象，并使用它们来构建我们自己的事件循环实现，该实现可以同时运行多个协程。
除非你正在构建依赖于异步编程内部工作方式的新 API 或框架，否则在你的日常开发任务中不太可能需要使用本章涉及的内容。本章中提到的技术主要针对依赖于异步编程内部工作方式的新 API 或框架的应用程序，以及那些想要更深入地理解异步 Python 内部原理的读者。
带有协程和函数的 API
如果我们自己构建 API，如何设计一个既能接收协程又能接收普通 Python 函数的 API 呢？asyncio 提供了两个方便的函数来帮助我们实现这一点：asyncio.iscoroutine 和 asyncio.iscoroutinefuction。这些函数让我们判断可调用对象是否为协程和协程函数，从而对它们使用不同的逻辑。这些函数是 Django 无缝处理同步和异步视图的基础，有兴趣可以了解一下。
import asyncio
async def foo():
pass
print(asyncio.iscoroutinefunction(foo))
print(asyncio.iscoroutine(foo()))
"""
True
True
"""
为了解这一点，让我们构建一个同时接收函数和协程的类，该类允许用户将函数添加到一个内部列表中，当调用 start 方法时将并发运行（如果是协程）或串行运行（如果是普通函数）。
import asyncio
class TaskRunner:
def __init__(self):
self.loop = asyncio.new_event_loop()
self.tasks = []
def add_task(self, func):
self.tasks.append(func)
async def _run_all(self):
awaitable_tasks = []
for task in self.tasks:
if asyncio.iscoroutinefunction(task):
awaitable_tasks.append(asyncio.create_task(task()))
elif asyncio.iscoroutine(task):
awaitable_tasks.append(asyncio.create_task(task))
else:
self.loop.call_soon(task)
await asyncio.gather(*awaitable_tasks)
def run(self):
self.loop.run_until_complete(self._run_all())
def regular_function():
print("我是一个普通函数")
async def coroutine_function():
print("我是一个协程函数")
runner = TaskRunner()
runner.add_task(coroutine_function())
runner.add_task(coroutine_function)
runner.add_task(regular_function)
runner.run()
"""
我是一个协程函数
我是一个协程函数
我是一个普通函数
"""
在代码中，TaskRunner 创建了一个新的事件循环和一个空的任务列表。然后定义一个 add_task 方法，它只是将一个函数（或协程）添加到待处理的任务列表中。然后，一旦用户调用 run 方法，我们就在事件循环中启动 _run_all 方法。该方法将遍历任务列表，并检查所涉及的函数是否为协程。如果是协程，则创建一个任务，否则使用事件循环的 call_soon 方法来调度普通函数在事件循环的下一次迭代中运行。然后，一旦创建了所有需要的任务，调用 gather，并等待其全部完成。
然后定义两个函数：一个是普通的 Python 函数，另一个是协程。我们创建一个 TaskRunner 实例并添加三个任务，调用 coroutine_function 两次来演示可在 API 中引用协程的两种不同方式，最终得到以上输出。
这表明已经成功地运行了协程和普通的 Python 函数，现在已经构建了一个 API，它可以处理协程以及普通的 Python 函数，增加了最终用户使用 API 的方式。接下来我们将查看上下文变量，它使我们能够存储任务本地的状态，而无需将其作为函数参数显式传递。
contextvars 模块的使用方法
这里需要介绍一个模块叫 contextvars，先来看看它的用法，至于为什么介绍它后续会会解释。
contextvars 是 Python 在 3.7 的时候引入的一个模块，从名字上很容易看出它指的是上下文变量（Context Variables），所以在介绍 contextvars 之前我们需要先了解一下什么是上下文（Context）。
Context 是一个包含了相关信息内容的对象，举个例子："比如一部 13 集的动漫，你直接点进第八集，看到女主角在男主角面前流泪了"。相信此时你是不知道为什么女主角会流泪的，因为你没有看前面几集的内容，缺失了相关的上下文信息。所以 Context 并不是什么神奇的东西，它的作用就是携带一些指定的信息。
web 框架中的 request
我们以 fastapi 和 sanic 为例，看看当一个请求过来的时候，它们是如何解析的。
# fastapi
from fastapi import FastAPI, Request
import uvicorn
app = FastAPI()
@app.get("/index")
async def index(request: Request):
name = request.query_params.get("name")
return {"name": name}
uvicorn.run("__main__:app", host="127.0.0.1", port=5555)
# -------------------------------------------------------
# sanic
from sanic import Sanic
from sanic.request import Request
from sanic import response
app = Sanic("sanic")
@app.get("/index")
async def index(request: Request):
name = request.args.get("name")
return response.json({"name": name})
app.run(host="127.0.0.1", port=6666)
发请求测试一下，看看结果是否正确。
可以看到请求都是成功的，并且对于 fastapi 和 sanic 而言，其 request 和 视图函数是绑定在一起的。也就是在请求到来的时候，会被封装成一个 Request 对象、然后传递到视图函数中。但对于 flask 而言则不是这样子的，我们看一下 flask 是如何接收请求参数的。
from flask import Flask, request
app = Flask("flask")
@app.route("/index")
def index():
name = request.args.get("name")
return {"name": name}
app.run(host="127.0.0.1", port=7777)
我们看到对于 flask 而言则是通过 import request 的方式，如果不需要的话就不用 import，当然我这里并不是在比较哪种方式好，主要是为了引出我们今天的主题。首先对于 flask 而言，如果我再定义一个视图函数的话，那么获取请求参数依旧是相同的方式，但是这样问题就来了，不同的视图函数内部使用同一个 request，难道不会发生冲突吗？
显然根据我们使用 flask 的经验来说，答案是不会的，至于原因就是 ThreadLocal。
ThreadLocal
ThreadLocal，从名字上看可以得出它肯定是和线程相关的。没错，它专门用来创建局部变量，并且创建的局部变量是和线程绑定的。
import threading
# 创建一个 local 对象
local = threading.local()
def get():
name = threading.current_thread().name
# 获取绑定在 local 上的 value
value = local.value
print(f"线程: {name}, value: {value}")
def set_():
name = threading.current_thread().name
# 为不同的线程设置不同的值
if name == "one":
local.value = "ONE"
elif name == "two":
local.value = "TWO"
# 执行 get 函数
get()
t1 = threading.Thread(target=set_, name="one")
t2 = threading.Thread(target=set_, name="two")
t1.start()
t2.start()
"""
线程 one, value: ONE
线程 two, value: TWO
"""
可以看到两个线程之间是互不影响的，因为每个线程都有自己唯一的 id，在绑定值的时候会绑定在当前的线程中，获取也会从当前的线程中获取。可以把 ThreadLocal 想象成一个字典：
{
"one": {"value": "ONE"},
"two": {"value": "TWO"}
}
更准确的说 key 应该是线程的 id，为了直观我们就用线程的 name 代替了，但总之在获取的时候只会获取绑定在该线程上的变量的值。
而 flask 内部也是这么设计的，只不过它没有直接用 threading.local，而是自己实现了一个 Local 类，除了支持线程之外还支持 greenlet 的协程，那么它是怎么实现的呢？首先我们知道 flask 内部存在 "请求 context" 和 "应用 context"，它们都是通过栈来维护的（两个不同的栈）。
# flask/globals.py
_request_ctx_stack = LocalStack()
_app_ctx_stack = LocalStack()
current_app = LocalProxy(_find_app)
request = LocalProxy(partial(_lookup_req_object, "request"))
session = LocalProxy(partial(_lookup_req_object, "session"))
每个请求都会绑定在当前的 Context 中，等到请求结束之后再销毁，这个过程由框架完成，开发者只需要直接使用 request 即可。所以请求的具体细节流程可以点进源码中查看，这里我们重点关注一个对象：werkzeug.local.Local，也就是上面说的 Local 类，它是变量的设置和获取的关键。直接看部分源码：
# werkzeug/local.py
class Local(object):
__slots__ = ("__storage__", "__ident_func__")
def __init__(self):
# 内部有两个成员：__storage__ 是一个字典，值就存在这里面
# __ident_func__ 只需要知道它是用来获取线程 id 的即可
object.__setattr__(self, "__storage__", {})
object.__setattr__(self, "__ident_func__", get_ident)
def __call__(self, proxy):
"""Create a proxy for a name."""
return LocalProxy(self, proxy)
def __release_local__(self):
self.__storage__.pop(self.__ident_func__(), None)
def __getattr__(self, name):
try:
# 根据线程 id 得到 value（一个字典）
# 然后再根据 name 获取对应的值
# 所以只会获取绑定在当前线程上的值
return self.__storage__[self.__ident_func__()][name]
except KeyError:
raise AttributeError(name)
def __setattr__(self, name, value):
ident = self.__ident_func__()
storage = self.__storage__
try:
# 将线程 id 作为 key，然后将值设置在对应的字典中
# 所以只会将值设置在当前的线程中
storage[ident][name] = value
except KeyError:
storage[ident] = {name: value}
def __delattr__(self, name):
# 删除逻辑也很简单
try:
del self.__storage__[self.__ident_func__()][name]
except KeyError:
raise AttributeError(name)
所以我们看到 flask 内部的逻辑其实很简单，通过 ThreadLocal 实现了线程之间的隔离。每个请求都会绑定在各自的 Context 中，获取值的时候也会从各自的 Context 中获取，因为它就是用来保存相关信息的（重要的是同时也实现了隔离）。
相应此刻你已经理解了上下文，但是问题来了，不管是 threading.local 也好、还是类似于 flask 自己实现的 Local 也罢，它们都是针对线程的。如果是使用 async def 定义的协程该怎么办呢？如何实现每个协程的上下文隔离呢？所以终于引出了我们的主角：contextvars。
contextvars
该模块提供了一组接口，可用于在协程中管理、设置、访问局部 Context 的状态。
import asyncio
import contextvars
c = contextvars.ContextVar("只是一个标识, 用于调试")
async def get():
# 获取值
return c.get() + "~~~"
async def set_(val):
# 设置值
c.set(val)
print(await get())
async def main():
coro1 = set_("协程1")
coro2 = set_("协程2")
await asyncio.gather(coro1, coro2)
asyncio.run(main())
"""
协程1~~~
协程2~~~
"""
ContextVar 提供了两个方法，分别是 get 和 set，用于获取值和设置值。我们看到效果和 ThreadingLocal 类似，数据在协程之间是隔离的，不会受到彼此的影响。
但我们再仔细观察一下，我们是在 set_ 函数中设置的值，然后在 get 函数中获取值。可 await get() 相当于是开启了一个新的协程，那么意味着设置值和获取值不是在同一个协程当中。但即便如此，我们依旧可以获取到希望的结果。因为 Python 的协程是无栈协程，通过 await 可以实现级联调用。
我们不妨再套一层：
import asyncio
import contextvars
c = contextvars.ContextVar("只是一个标识, 用于调试")
async def get1():
return await get2()
async def get2():
return c.get() + "~~~"
async def set_(val):
# 设置值
c.set(val)
print(await get1())
print(await get2())
async def main():
coro1 = set_("协程1")
coro2 = set_("协程2")
await asyncio.gather(coro1, coro2)
asyncio.run(main())
"""
协程1~~~
协程1~~~
协程2~~~
协程2~~~
"""
我们看到不管是 await get1() 还是 await get2()，得到的都是 set_ 中设置的结果，说明它是可以嵌套的。并且在这个过程当中，可以重新设置值。
import asyncio
import contextvars
c = contextvars.ContextVar("只是一个标识, 用于调试")
async def get1():
c.set("重新设置")
return await get2()
async def get2():
return c.get() + "~~~"
async def set_(val):
# 设置值
c.set(val)
print("------------")
print(await get2())
print(await get1())
print(await get2())
print("------------")
async def main():
coro1 = set_("协程1")
coro2 = set_("协程2")
await asyncio.gather(coro1, coro2)
asyncio.run(main())
"""
------------
协程1~~~
重新设置~~~
重新设置~~~
------------
------------
协程2~~~
重新设置~~~
重新设置~~~
------------
"""
先 await get2() 得到的就是 set_ 函数中设置的值，这是符合预期的。但是我们在 get1 中将值重新设置了，那么之后不管是 await get1() 还是直接 await get2()，得到的都是新设置的值。
这也说明了，一个协程内部 await 另一个协程，另一个协程内部 await 另另一个协程，不管套娃（await）多少次，它们获取的值都是一样的。并且在任意一个协程内部都可以重新设置值，然后获取会得到最后一次设置的值。再举个栗子：
import asyncio
import contextvars
c = contextvars.ContextVar("只是一个标识, 用于调试")
async def get1():
return await get2()
async def get2():
val = c.get() + "~~~"
c.set("重新设置啦")
return val
async def set_(val):
# 设置值
c.set(val)
print(await get1())
print(c.get())
async def main():
coro = set_("古明地觉")
await coro
asyncio.run(main())
"""
古明地觉~~~
重新设置啦
"""
await get1() 的时候会执行 await get2()，然后在里面拿到 c.set 设置的值，打印 "古明地觉~~~"。但是在 get2 里面，又将值重新设置了，所以第二个 print 打印的就是新设置的值。
如果在 get 之前没有先 set，那么会抛出一个 LookupError，所以 ContextVar 支持默认值：
import asyncio
import contextvars
c = contextvars.ContextVar("只是一个标识, 用于调试",
default="哼哼")
async def set_(val):
print(c.get())
c.set(val)
print(c.get())
async def main():
coro = set_("古明地觉")
await coro
asyncio.run(main())
"""
哼哼
古明地觉
"""
除了在 ContextVar 中指定默认值之外，也可以在 get 中指定：
import asyncio
import contextvars
c = contextvars.ContextVar("只是一个标识, 用于调试",
default="哼哼")
async def set_(val):
print(c.get("古明地恋"))
c.set(val)
print(c.get())
async def main():
coro = set_("古明地觉")
await coro
asyncio.run(main())
"""
古明地恋
古明地觉
"""
所以结论如下，如果在 c.set 之前使用 c.get：
- 当 ContextVar 和 get 中都没有指定默认值，会抛出 LookupError；
- 只要有一方设置了，那么会得到默认值；
- 如果都设置了，那么以 get 为准；
如果 c.get 之前执行了 c.set，那么无论 ContextVar 和 get 有没有指定默认值，获取到的都是 c.set 设置的值。所以总的来说还是比较好理解的，并且 ContextVar 除了可以作用在协程上面，它也可以用在线程上面。没错，它可以替代 threading.local，我们来试一下：
import threading
import contextvars
c = contextvars.ContextVar("context_var")
def get():
name = threading.current_thread().name
value = c.get()
print(f"线程 {name}, value: {value}")
def set_():
name = threading.current_thread().name
if name == "one":
c.set("ONE")
elif name == "two":
c.set("TWO")
get()
t1 = threading.Thread(target=set_, name="one")
t2 = threading.Thread(target=set_, name="two")
t1.start()
t2.start()
"""
线程 one, value: ONE
线程 two, value: TWO
"""
和 threading.local 的表现是一样的，但是更建议使用 ContextVars。不过前者可以绑定任意多个值，而后者只能绑定一个值（可以通过传递字典的方式解决这一点）。
c.Token
当我们调用 c.set 的时候，其实会返回一个 Token 对象：
import contextvars
c = contextvars.ContextVar("context_var")
token = c.set("val")
print(token)
"""
<Token var=<ContextVar name='context_var' at 0x00..> at 0x00...>
"""
Token 对象有一个 var 属性，它是只读的，会返回指向此 token 的 ContextVar 对象。
import contextvars
c = contextvars.ContextVar("context_var")
token = c.set("val")
print(token.var is c) # True
print(token.var.get()) # val
print(
token.var.set("val2").var.set("val3").var is c
) # True
print(c.get()) # val3
Token 对象还有一个 old_value 属性，它会返回上一次 set 设置的值，如果是第一次 set，那么会返回一个 <Token.MISSING>。
import contextvars
c = contextvars.ContextVar("context_var")
token = c.set("val")
# 该 token 是第一次 c.set 所返回的
# 在此之前没有 set，所以 old_value 是 <Token.MISSING>
print(token.old_value) # <Token.MISSING>
token = c.set("val2")
print(c.get()) # val2
# 返回上一次 set 的值
print(token.old_value) # val
那么这个 Token 对象有什么作用呢？从目前来看貌似没太大用处啊，其实它最大的用处就是和 reset 搭配使用，可以对状态进行重置。
import contextvars
c = contextvars.ContextVar("context_var")
token = c.set("val")
# 显然是可以获取的
print(c.get()) # val
# 将其重置为 token 之前的状态
# 但这个 token 是第一次 set 返回的
# 那么之前就相当于没有 set 了
c.reset(token)
try:
c.get() # 此时就会报错
except LookupError:
print("报错啦") # 报错啦
# 但是我们可以指定默认值
print(c.get("默认值")) # 默认值
contextvars.Context
它负责保存 ContextVars 对象和设置的值之间的映射，但是我们不会直接通过 contextvars.Context 来创建，而是通过 contentvars.copy_context 函数来创建。
mport contextvars
c1 = contextvars.ContextVar("context_var1")
c1.set("val1")
c2 = contextvars.ContextVar("context_var2")
c2.set("val2")
# 此时得到的是所有 ContextVar 对象和设置的值之间的映射
# 它实现了 collections.abc.Mapping 接口
# 因此我们可以像操作字典一样操作它
context = contextvars.copy_context()
# key 就是对应的 ContextVar 对象，value 就是设置的值
print(context[c1]) # val1
print(context[c2]) # val2
for ctx, value in context.items():
print(ctx.get(), ctx.name, value)
"""
val1 context_var1 val1
val2 context_var2 val2
"""
print(len(context)) # 2
除此之外，context 还有一个 run 方法：
import contextvars
c1 = contextvars.ContextVar("context_var1")
c1.set("val1")
c2 = contextvars.ContextVar("context_var2")
c2.set("val2")
context = contextvars.copy_context()
def change(val1, val2):
c1.set(val1)
c2.set(val2)
print(c1.get(), context[c1])
print(c2.get(), context[c2])
# 在 change 函数内部，重新设置值
# 然后里面打印的也是新设置的值
context.run(change, "VAL1", "VAL2")
"""
VAL1 VAL1
VAL2 VAL2
"""
print(c1.get(), context[c1])
print(c2.get(), context[c2])
"""
val1 VAL1
val2 VAL2
"""
我们看到 run 方法接收一个 callable，如果在里面修改了 ContextVar 实例设置的值，那么对于 ContextVar 而言只会在函数内部生效，一旦出了函数，那么还是原来的值。但是对于 Context 而言，它是会受到影响的，即便出了函数，也是新设置的值，因为它直接把内部的字典给修改了。
以上就是 contextvars 模块的用法，在多个协程之间传递数据是非常方便的，并且也是并发安全的。如果你用过 Go 的话，你应该会发现和 Go 在 1.7 版本引入的 context 模块比较相似，当然 Go 的 context 模块功能要更强大一些，除了可以传递数据之外，对多个 goroutine 的级联管理也提供了非常清蒸的解决方案。
总之对于 contextvars 而言，它传递的数据应该是多个协程之间需要共享的数据，像 cookie, session, token 之类的，比如上游接收了一个 token，然后不断地向下透传（比使用函数参数简单的多）。但是不要把本应该作为函数参数的数据，也通过 contextvars 来传递，这样就有点本末倒置了。
上下文变量
假设我们正在使用一个基于线程请求的 Web 服务器 REST API，向 Web 服务器发出请求时会携带一些公共数据，并且还需要跟踪这些数据，如用户 ID、访问令牌或其他信息。然后你可能试图在 Web 服务器的所有线程中全局存储这些数据，但这存在问题。主要缺点是，需要处理从线程到其数据的映射，以及所有锁定，以防止竞态条件。更好的做法是通过使用线程局部变量的概念来解决这个问题，线程局部变量是特定于一个线程的全局状态。对于在线程本地设置的数据，只能被设置它的线程看到，从而避免了线程到数据的映射发生的竞态条件。
当然，在 asyncio 应用程序中，通常只有一个线程，所以作为线程本地存储的任何内容都可以在应用程序的其他地方使用。但在 PEP-567 中引入了上下文变量的概念，从而处理单线程并发模型中的局部线程问题。上下文变量类似于线程局部变量，区别在于它们是特定任务的局部变量，而不是线程的局部变量。这意味着，如果一个任务创建了一个上下文变量，初始任务中的任何内部协程或任务都可访问该变量，而其他任何任务都不能看到或修改该变量。这让我们可以跟踪特定任务的状态，而不必将其作为显式参数进行传递。
import asyncio
from asyncio import StreamReader, StreamWriter
from contextvars import ContextVar
class Server:
user_address = ContextVar("user_address")
def __init__(self, host: str, port: int):
self.host = host
self.port = port
async def start_server(self):
server = await asyncio.start_server(self._client_connected, self.host, self.port)
await server.serve_forever()
def _client_connected(self, reader: StreamReader, writer: StreamWriter):
# 当客户端连接时，将客户端的地址存储在上下文变量中
self.user_address.set(writer.get_extra_info("peername"))
asyncio.create_task(self.listen_for_message(reader))
async def listen_for_message(self, reader: StreamReader):
while data := await reader.read(1024):
print(f"从 {self.user_address.get()} 中获得数据 {data}")
async def main():
server = Server("localhost", 9999)
await server.start_server()
asyncio.run(main())
在代码中，首先创建一个 ContextVar 实例来保存用户的地址信息。上下文变量要求我们提供一个字符串名称，所以这里给它一个描述性名称 user_address，主要用于调试。然后在 _cient_connected 回调中，将上下文变量的数据设置为客户端的地址，这将允许从该父任务产生的任何任务都可以访问我们设置的信息。
在 listen_for_messages 协程方法中，我们监听来自客户端的数据：当得到数据时将它与我们存储在上下文变量中的地址一起输出。当运行此应用程序并连接多个客户端，并且发送一些消息时，应该会看到如下输出：
注意，地址的端口号不同，表明我们从 locahost 上的两个不同客户端获取了消息。即使只创建了一个上下文变量，我们仍能访问特定于每个客户端的唯一数据。这为我们提供了一种在任务之间传递数据的简洁方式，而不必显式地将数据传递给任务。
强制事件循环迭代
事件循环内部的运行方式在很大程度上超出了我们的控制范围，它决定何时以及如何执行协程和任务。也就是说，如有必要，有一种方法可以触发事件循环迭代。对于长时间运行的任务来说，这很方便，因为可避免阻塞事件循环（这种情况下，你还应该考虑线程）或确保任务立即启动。
回顾一下，如果我们要创建多个任务，在一个任务遇到 await 之前，它会一直运行下去。那么如何实现显式地切换呢？
import asyncio
async def foo():
print("foo1")
await asyncio.sleep(0)
print("foo2")
async def bar():
print("bar1")
await asyncio.sleep(0)
print("bar2")
async def main():
task1 = asyncio.create_task(foo())
task2 = asyncio.create_task(bar())
await asyncio.gather(task1, task2)
asyncio.run(main())
"""
foo1
bar1
foo2
bar2
"""
如果没有 asyncio.sleep(0)，那么每个任务会一直运行完毕，然后才会执行下一个任务。但通过 sleep(0) 可以让任务强行发生切换，这类似于 greenlet 里面的 switch。
使用不同的事件循环实现
asyncio 提供了事件循环的默认实现，我们到目前为止一直在使用该实现，但也可以使用具有不同特征的其他实现。有几种方法可以使用不同的实现：一个是子类化 AbstractEventLoop 类，并实现它的方法创建一个实例，然后用 asyncio.set_event_loop 将其设置为事件循环。如果正在构建自己的实现，可以使用这种方法，也可以使用现成的事件循环，让我们来看一个名为 uvloop 的实现。
那什么是 uvloop，为什么要使用它呢？uvloop 是一个事件循环的实现，它严重依赖于 libuv 库（https://libuv.org），libuv 是一个非常流行的事件驱动框架。并且由于 libuv 是用 C 实现的，因此比纯解释型 Python 代码具有更好的性能。uvloop 可以比默认的 asyncio 事件循环更快，在编写基于套接字和流的应用程序时，它往往表现得非常出众。你可以在该项目的 Github 站点上阅读有关基准测试的更多信息。注意：uvloop 仅在 Unix 平台上可用。
直接 pip install uvloop 即可。
import asyncio
from asyncio import StreamReader, StreamWriter
import uvloop
async def connected(reader: StreamReader, writer: StreamWriter):
line = await reader.readline()
writer.write(line)
await writer.drain()
writer.close()
await writer.wait_closed()
async def main():
# 启动一个服务器，客户端连接 "localhost:9999" 时执行 connected 协程
server = await asyncio.start_server(connected, port=9999)
await server.serve_forever()
# 需要安装 uvloop 事件循环
uvloop.install()
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
在代码中，我们调用了 uvloop.instal()，它将切换事件循环。如果愿意可使用以下代码手动执行此操作，而不是调用 install。
loop = uvloop.new_event_loop()
asyncio.set_event_loop(loop)
重要的是在调用 asyncio.run(main()) 之前调用它，在后台 asyncio.run 会调用 get_event_loop，如果不存在事件循环，则创建一个事件循环。所以一定要先调用 uvloop.install()，或者将默认的事件循环实现替换掉。
如果 uvloop 的事件循环有助于应用程序的性能提升，你可能需要进行一个基准测试，Github 上的 uvloop 项目的代码可以在吞吐量和每秒请求数方面运行基准测试。
我们现在已经看到了如何使用现有的事件循环实现，而不是默认的事件循环实现。接下来，将看到如何完全在 asyncio 之外创建自己的事件循环，这会使我们更深入地了解 asyncio 事件循环以及协程、任务和 future 在底层是如何工作的。
创建自定义事件循环
对于 asyncio，它在概念上不同于 async/await 语法和协程，协程类定义甚至不在 asyncio 库模块中。
协程、async/await 语法与 asyncio 是不同的概念，Python 带有一个默认的事件循环实现，即 asyncio，这是我们迄今为止一直用来运行事件循环的方法。当然我们可以使用任何事件循环实现，甚至是自定义事件循环实现。下面就来分析一下，如何构建自己的可以处理非阻塞套接字的简单事件循环实现。
协程和生成器
在 Python3.5 引入 async 和 await 语法之前，协程和生成器之间的关系是显而易见的。让我们使用装饰器和生成器构建一个简单协程，它使用旧语法休眠 1 秒。
import asyncio
@asyncio.coroutine
def coroutine():
print("开始修秒 1 秒")
yield from asyncio.sleep(1)
print("休眠结束")
asyncio.run(coroutine())
"""
开始修秒 1 秒
休眠结束
"""
我们使用 @asyncio.coroutine 装饰器代替 async 关键字来指定函数是协程函数，在驱动另一个协程执行时也是使用 yield from，而不是 await 关键字。目前，async 和 await 关键字只是围绕这个结构的语法糖（syntactic sugar）。
上面这种协程叫做生成器协程，而使用 async def 定义出来的是原生协程，但这两者可以混合使用。
import asyncio
async def native_coroutine():
return "native coroutine"
@asyncio.coroutine
def gen_coroutine():
return (yield from native_coroutine()) + "......"
async def main():
print(await gen_coroutine())
asyncio.run(main())
"""
native coroutine......
"""
main() 里面 await 一个生成器协程，生成器协程里面 yield from 原生协程，因此它们是可以混用的。然后注意：async def 里面不可以出现 yield from，但出现 yield 是可以的，因为这是异步生成器。
不建议使用基于生成器的协程
请注意，基于生成器的协程目前计划在 Python 3.10 版本中完全删除，你可能会在遗留代码库中遇到它们，但你不应该再以这种风格编写新的异步代码。
那为什么生成器对单线程并发模型有意义呢？回顾一下，协程在遇到阻塞操作时需要暂停执行，从而允许其他协程运行。生成器在达到 yield point 时会暂停执行，从而有效地在中途暂停它们。这意味着如果有两个生成器，可以交错执行它们。让第一个生成器运行直到它到达一个 yield point（在协程语言中则叫 await point），然后让第二个生成器运行到它的 yield point。重复这种方式，直到两个生成器都执行完成。为观察这一点，让我们构建一个非常简单的示例，该示例将两个生成器交错执行。
"""
补充：判断一个函数是什么类型的函数，可以有一个更快捷的方式
flags = func.__code__.co_flags
如果 flags & 0x20 为真，则 func 是生成器函数
如果 flags & 0x80 为真，则 func 是原生协程函数
如果 flags & 0x200 为真，则 func 是异步生成器函数
"""
from typing import Generator
def generator(start: int, end: int):
for i in range(start, end):
yield i
one_to_five = generator(1, 6)
six_to_ten = generator(6, 11)
def run_generator_step(gen: Generator):
try:
return gen.send(None)
except StopIteration:
print("生成器执行完毕")
while True:
one_to_five_result = run_generator_step(one_to_five)
six_to_ten_result = run_generator_step(six_to_ten)
print(one_to_five_result)
print(six_to_ten_result)
if one_to_five_result is None and six_to_ten_result is None:
break
"""
1
6
2
7
3
8
4
9
5
10
生成器执行完毕
生成器执行完毕
None
None
"""
想象一下，我们没有得到数字，而是采取了一些缓慢的操作。当缓慢的操作完成后，可以恢复生成器，从之前停止的地方继续执行，而其他没有被暂停的生成器可运行其他代码，这是事件循环工作的核心。我们跟踪在慢速操作中暂停执行的生成器，然后，任何其他生成器都可在另一个生成器暂停时运行。一旦慢速操作完成，可通过再次调用 send 来唤醒上一个生成器，并推进到它的下一个 yield point。
其实 async 和 await 只是生成器周围的语法糖，可通过创建一个协程实例并在其上调用 send 来演示这一点。
import asyncio
async def say_hello():
print("Hello")
async def say_goodbye():
print("Goodbye")
async def meet_and_greet():
await say_hello()
await say_goodbye()
coro = meet_and_greet()
coro.send(None)
"""
Hello
Goodbye
Traceback (most recent call last):
File ".../main.py", line 14, in <module>
coro.send(None)
StopIteration
"""
协程上调用 send 会运行 meet_and_greet 中的所有协程，因为在等待结果时，我们实际上并没有暂停什么，因为所有代码都是立即运行的，即使在 await 语句中也是如此。
那么如何让协程暂停并在慢速操作中唤醒呢？为此，让我们自定义一个可等待对象，这样就可以使用 await 语法，而不是生成器风格的语法。
自定义可等待对象
如果想自定义可等待对象，那么要实现 __await__ 方法，该方法的唯一要求是它返回一个迭代器。
class CustomFuture:
"""自定义 Future"""
def __init__(self):
self._result = None
self._is_finished = False
self._done_callback = None
def result(self):
return self._result
def is_finished(self):
return self._is_finished
def set_result(self, result):
self._result = result
self._is_finished = True
if self._done_callback:
self._done_callback(result)
def add_done_callback(self, fn):
self._done_callback = fn
def __await__(self):
if not self._is_finished:
yield self
return self._result
在代码清单创建了一个 CustomFuture 类，其中定义了 __await__ 以及设置结果、获取结果和添加回调的方法。 __await__ 方法用来检查 future 是否完成，如果是，我们只返回结果，迭代器就完成了。如果没有完成，我们返回 self，这意味着迭代器将继续无限地返回自身，直到值被设置。就生成器而言，这意味着可以一直调用 await ，直到值被设置为止。
future = CustomFuture()
i = 0
while True:
try:
print("checking future")
# 执行 __await__ 方法，返回一个生成器
gen = future.__await__()
# 相当于 __next__()，如果 future 没有完成，那么会 yield self
gen.send(None)
print("future is not done")
if i == 1:
print("Setting future value")
future.set_result("Finished")
i += 1
# 如果 future 完成了，那么会直接 return
# 那么 gen.send(None) 就会 raise 一个 StopIteration
except StopIteration as si:
print(f"value is {si.value}")
break
"""
checking future
future is not done
checking future
future is not done
Setting future value
checking future
value is Finished
"""
在代码中我们创建了一个自定义 future 和一个调用 __await__ 方法的循环，然后尝试推进迭代器。如果 future 完成了，则会抛出一个 StopIteration 异常，并带有 future 的结果，否则继续执行循环的下一次迭代。
这个例子只是为了强化对可等待对象的理解，在现实工作中我们不会编写这样的代码，因为我们通常希望用其他内容来设置 future 的结果。接下来对其进行扩展，使用套接字和选择器模块做一些更有用的事情。
使用带有 future 的套接字
在前面的文章中，我们学习了一些关于 selectors 模块的知识，它允许我们在套接字事件（如连接建立或准备读取数据）发生时注册回调函数。现在，将通过使用自定义 future 类与 selectors 选择器交互来扩展这一知识，在套接字事件发生时设置 future 的结果。
回顾一下，选择器允许注册回调函数，以便在套接字上发生事件（如读或写）时运行这些函数。这个概念非常适合我们构建 future，对应它的 set_result 方法。
from functools import partial
import socket
import selectors
from selectors import BaseSelector
class CustomFuture:
"""自定义 Future"""
def __init__(self):
self._result = None
self._is_finished = False
self._done_callback = None
def result(self):
return self._result
def is_finished(self):
return self._is_finished
def set_result(self, result):
self._result = result
self._is_finished = True
if self._done_callback:
self._done_callback(result)
def add_done_callback(self, fn):
self._done_callback = fn
def __await__(self):
if not self._is_finished:
yield self
return self._result
def accept_connection(future: CustomFuture, connection: socket.socket):
# 当客户端连接时，为 future 设置连接套接字
print(f"连接建立: {connection}")
future.set_result(connection)
async def socket_accept(sel: BaseSelector, server: socket.socket) -> socket.socket:
# 向选择器注册 accept_connection 函数，并暂停以等待客户端连接。
print("注册套接字")
future = CustomFuture()
# register 方法接收三个参数
# 参数一：传一个套接字即可
# 参数二：事件类型，这里是读事件
# 参数三：事件发生时，执行的回调函数
sel.register(server, selectors.EVENT_READ, partial(accept_connection, future))
print("暂停以等待连接")
# 什么时候 future 解除阻塞呢？显然是当客户端来连接时
# 会执行 accept_connection，然后设置和客户端用于通信的连接
connection: socket.socket = await future
return connection
async def main(sel: BaseSelector):
server = socket.socket()
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(("localhost", 9999))
server.listen()
# IO 多路复用一定要搭配非阻塞 IO
server.setblocking(False)
await socket_accept(sel, server)
selector = selectors.DefaultSelector()
coro = main(selector)
while True:
try:
state = coro.send(None)
events = selector.select()
for key, mask in events:
print("处理 selector 事件...")
callback = key.data
callback(key.fileobj)
except StopIteration:
print("引用程序结束")
break
我们创建客户端连接一下，然后查看服务端输出：
现在，我们已经构建了一个基本的异步应用程序，只使用 async 和 await 关键字而不使用任何 asyncio 内容，最后的 while 循环也是一个简单的事件循环，并演示了 asyncio 事件循环如何工作的关键概念。当然，如果没有创建任务的能力，就不能同时做很多事情。
任务的实现
任务是 future 和协程的组合，当它包装的协程完成时，任务的 future 就完成了。可通过继承 CustomFuture 类，并编写一个接收协程的构造函数来包装 future 和协程，但我们仍然需要一种方法来运行该协程。
from functools import partial
import socket
import selectors
from selectors import BaseSelector
class CustomFuture:
"""自定义 Future"""
def __init__(self):
self._result = None
self._is_finished = False
self._done_callback = None
def result(self):
return self._result
def is_finished(self):
return self._is_finished
def set_result(self, result):
self._result = result
self._is_finished = True
if self._done_callback:
self._done_callback(result)
def add_done_callback(self, fn):
self._done_callback = fn
def __await__(self):
if not self._is_finished:
yield self
return self._result
class CustomTask(CustomFuture):
"""自定义 Task"""
def __init__(self, coro, loop):
super().__init__()
self._coro = coro
self._loop = loop
self._current_result = None
self._task_state = None
# 用事件循环注册任务
loop.register_task(self)
def step(self):
# 运行协程的一个步骤
try:
if self._task_state is None:
self._task_state = self._coro.send(None)
if isinstance(self._task_state, CustomFuture):
# 如果协程产生一个 future，则添加一个 done 回调
self._task_state.add_done_callback(self._future_done)
except StopIteration as e:
self.set_result(e.value)
def _future_done(self, result):
self._current_result = result
try:
self._task_state = self._coro.send(self._current_result)
except StopIteration as e:
self.set_result(e.value)
在代码中，创建了 CustomFuture 的子类，并创建了一个接收协程和事件循环的构造函数，通过调用 loop.register_task 将任务注册到循环中。然后在 step 方法中调用协程的 send 方法，如果协程产生 CustomFuture，添加一个 done 回调。这种情况下，done 回调将获取 future 的结果，将其发送到我们包装的协程，并在 future 完成时推进它。
实现事件循环
我们现在知道了如何运行协程，并创建了 future 和任务的实现，这为我们提供了构建事件循环需要的所有构建块。但要构建异步套接字应用程序，事件 API 需要是什么样子的？
需要一个方法接收主入口协程，就像 asyncio.run 一样需要一些方法来接收连接、接收数据和关闭套接字，这些方法将使用选择器注册和注销套接字需要一个方法来注册 CustomTask
首先谈谈主要切入点，我们将这个方法称为 run，这是事件循环的原动力。这个方法将接收一个主入口协程，并调用它的 send，然后在一个无限循环中跟踪生成器的结果。如果主协程产生了一个 future，将添加一个 done 回调，以便在 future 完成时跟踪其结果。之后，将运行所有已注册任务的 step 方法，然后调用选择器，等待任何套接字事件被触发。一旦它们运行，将运行相关的回调函数，并触发循环的另一次迭代。如果主协程在任何时候抛出一个 Stoplteration 异常，我们就知道应用程序已经完成，并且可以退出以及返回异常中的值。
接下来，需要使用协程方法来接收套接字连接，并从客户端套接字接收数据。这里的策略是创建一个 CustomFuture 实例，回调将设置其结果，并将这个回调注册到选择器中，以便在读取事件时触发。然后，我们将等待这个 future。
最后，需要一个方法向事件循环注册任务，此方法简单地接收一个任务并将其添加到列表中。然后，在事件循环的每次迭代中，将对在事件循环中注册的任何任务调用 step 方法，如果它们准备好，则推进它们。实现以上这些，将产生一个最小的可行事件循环。
from functools import partial
from typing import List
import selectors
class CustomFuture:
"""自定义 Future"""
def __init__(self):
self._result = None
self._is_finished = False
self._done_callback = None
def result(self):
return self._result
def is_finished(self):
return self._is_finished
def set_result(self, result):
self._result = result
self._is_finished = True
if self._done_callback:
self._done_callback(result)
def add_done_callback(self, fn):
self._done_callback = fn
def __await__(self):
if not self._is_finished:
yield self
return self._result
class CustomTask(CustomFuture):
"""自定义 Task"""
def __init__(self, coro, loop):
super().__init__()
self._coro = coro
self._loop = loop
self._current_result = None
self._task_state = None
# 用事件循环注册任务
loop.register_task(self)
def step(self):
# 运行协程的一个步骤
try:
if self._task_state is None:
self._task_state = self._coro.send(None)
if isinstance(self._task_state, CustomFuture):
# 如果协程产生一个 future，则添加一个 done 回调
self._task_state.add_done_callback(self._future_done)
except StopIteration as e:
self.set_result(e.value)
def _future_done(self, result):
self._current_result = result
try:
self._task_state = self._coro.send(self._current_result)
except StopIteration as e:
self.set_result(e.value)
class EventLoop:
_task_to_run: List[CustomTask] = []
def __init__(self):
self.selector = selectors.DefaultSelector()
self.current_result = None
def _register_socket_to_read(self, sock, callback):
future = CustomFuture()
try:
self.selector.get_key(sock)
except KeyError:
sock.setblocking(False)
self.selector.register(sock, selectors.EVENT_READ, partial(callback, future))
else:
self.selector.modify(sock, selectors.EVENT_READ, partial(callback, future))
return future
def _set_current_result(self, result):
self.current_result = result
async def sock_accept(self, sock):
print("注册套接字，接收客户端连接")
return await self._register_socket_to_read(sock, self.accept_connection)
async def sock_recv(self, sock):
print("注册套接字，接收客户端数据")
return await self._register_socket_to_read(sock, self.received_data)
def sock_close(self, sock):
self.selector.unregister(sock)
sock.close()
def register_task(self, task):
# 向事件循环注册任务
self._task_to_run.append(task)
def received_data(self, future, sock):
data = sock.recv(1024)
future.set_result(data)
def accept_connection(self, future, sock):
conn, _ = sock.accept()
future.set_result(conn)
def run(self, coro):
# 运行一个主协程，直到它完成，在每次选代中执行任何待处理的任务。
self.current_result = coro.send(None)
while True:
try:
if isinstance(self.current_result, CustomFuture):
self.current_result.add_done_callback(self._set_current_result)
if self.current_result.result() is not None:
coro.send(self.current_result.result())
else:
self.current_result = coro.send(self.current_result)
except StopIteration as e:
return e.value
for task in self._task_to_run:
task.step()
self._task_to_run = [task for task in self._task_to_run if not task.is_finished()]
events = self.selector.select()
print("有事件发生, 开始处理")
for key, mask in events:
callback = key.data
callback(key.fileobj)
首先定义一个 _register_socket_to_read 简便方法，此方法接收一个套接字和一个回调，如果套接字尚未注册，则将它们注册到选择器。如果套接字已注册，将替换回调，回调的第一个参数需要一个 future，在这个方法中我们创建一个 future。最后返回绑定到回调的 future，这意味着方法的调用者现在可以等待它并暂停执行，直到回调完成。
然后定义协程方法来接收套接字数据和接收新的客户端连接，这些方法分别为 sock_recv 和 sock_accept。这些方法调用我们刚定义的 _register_socket_to_read 方法，传入处理数据和新连接可用的回调（这些方法只是将这些数据设置为 future）。
接着构建 run 方法，这个方法接收主入口协程，并在它上面调用 send 方法，将它推进到第一个暂停点，并存储来自 send 方法的结果。然后开始一个无限循环，首先检查主协程的当前结果是不是 CustomFuture，如果是，注册一个回调来存储结果，如有必要，可将其发送回主协程。如果结果不是 CustomFuture，只需要将其发送给协程。一旦控制了主协程的流程，就可通过调用 step 方法来运行在事件循环中注册的任何任务。一旦任务运行完毕，就会从任务列表中删除。
然后调用 selector.select，会一直阻塞，直到我们注册的套接字上触发任何事件为止。一旦有了一个套接字事件或一组事件，就会遍历它们，调用在 _register_socket_to_read 中为该套接字注册的回调。在实现中，任何套接字事件都会触发事件循环的迭代。现在已经实现了 EventLoop 类，我们来创建第一个没有 asyncio 的异步应用程序。
使用自定义事件循环实现服务器
现在有了一个事件循环，我们将构建一个非常简单的服务器来记录从连接的客户端收到的消息。会创建一个服务器套接字，并编写一个协程函数在无限循环中监听连接。建立连接后，将创建一个任务从该客户端读取数据，直到它们断开连接为止。
from functools import partial
from typing import List
import socket
import selectors
class CustomFuture:
"""自定义 Future"""
def __init__(self):
self._result = None
self._is_finished = False
self._done_callback = None
def result(self):
return self._result
def is_finished(self):
return self._is_finished
def set_result(self, result):
self._result = result
self._is_finished = True
if self._done_callback:
self._done_callback(result)
def add_done_callback(self, fn):
self._done_callback = fn
def __await__(self):
if not self._is_finished:
yield self
return self._result
class CustomTask(CustomFuture):
"""自定义 Task"""
def __init__(self, coro, loop):
super().__init__()
self._coro = coro
self._loop = loop
self._current_result = None
self._task_state = None
# 用事件循环注册任务
loop.register_task(self)
def step(self):
# 运行协程的一个步骤
try:
if self._task_state is None:
self._task_state = self._coro.send(None)
if isinstance(self._task_state, CustomFuture):
# 如果协程产生一个 future，则添加一个 done 回调
self._task_state.add_done_callback(self._future_done)
except StopIteration as e:
self.set_result(e.value)
def _future_done(self, result):
self._current_result = result
try:
self._task_state = self._coro.send(self._current_result)
except StopIteration as e:
self.set_result(e.value)
class EventLoop:
_task_to_run: List[CustomTask] = []
def __init__(self):
self.selector = selectors.DefaultSelector()
self.current_result = None
def _register_socket_to_read(self, sock, callback):
future = CustomFuture()
try:
self.selector.get_key(sock)
except KeyError:
sock.setblocking(False)
self.selector.register(sock, selectors.EVENT_READ, partial(callback, future))
else:
self.selector.modify(sock, selectors.EVENT_READ, partial(callback, future))
return future
def _set_current_result(self, result):
self.current_result = result
async def sock_accept(self, sock):
print("注册套接字，接收客户端连接")
return await self._register_socket_to_read(sock, self.accept_connection)
async def sock_recv(self, sock):
print("注册套接字，接收客户端数据")
return await self._register_socket_to_read(sock, self.received_data)
def sock_close(self, sock):
self.selector.unregister(sock)
sock.close()
def register_task(self, task):
# 向事件循环注册任务
self._task_to_run.append(task)
def received_data(self, future, sock):
data = sock.recv(1024)
future.set_result(data)
def accept_connection(self, future, sock):
conn, _ = sock.accept()
future.set_result(conn)
def run(self, coro):
# 运行一个主协程，直到它完成，在每次选代中执行任何待处理的任务。
self.current_result = coro.send(None)
while True:
try:
if isinstance(self.current_result, CustomFuture):
self.current_result.add_done_callback(self._set_current_result)
if self.current_result.result() is not None:
coro.send(self.current_result.result())
else:
self.current_result = coro.send(self.current_result)
except StopIteration as e:
return e.value
for task in self._task_to_run:
task.step()
self._task_to_run = [task for task in self._task_to_run if not task.is_finished()]
events = self.selector.select()
print("有事件发生, 开始处理")
for key, mask in events:
callback = key.data
callback(key.fileobj)
async def read_from_client(conn, loop: EventLoop):
print(f"从 {conn.getpeername()} 里面读取数据")
try:
while data := await loop.sock_recv(conn):
print(f"从客户端读取到数据 {data}")
finally:
loop.sock_close(conn)
async def listen_for_connection(sock, loop: EventLoop):
# 监听客户端连接，创建一个任务在客户端连接时读取数据
while True:
print("等待客户端连接")
conn = await loop.sock_accept(sock)
print(f"从 {conn.getpeername()} 获取一个新的连接")
CustomTask(read_from_client(conn, loop), loop)
async def main(loop: EventLoop):
server = socket.socket()
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
server.bind(("localhost", 9999))
server.listen()
server.setblocking(False)
await listen_for_connection(server, loop)
event_loop = EventLoop()
event_loop.run(main(event_loop))
我们先来启动服务，会有以下输出：
程序启动时，会执行 listen_for_connection，在执行到 await loop.sock_accept(sock) 时，会向多路复用器注册一个事件，并返回一个 CustomFuture 对象。如果有客户端来连接，则执行 EventLoop 的 accept_connection 方法，给 future 设置结果集（这里是 conn）。
我们创建一个客户端，然后发一行数据：
建立完连接之后，创建一个 CustomTask 来同时监听来自该客户端的数据，然后立马监听下一个连接。
虽然这不是一个在生产环境中使用的事件循环（我们并没有正确地处理异常，且只允许套接字事件触发事件循环迭代，当然还有其他缺点），但这应该有助于你了解 Python 中事件循环和异步编程的内部工作原理。可采用此处的概念，构建一个用于生产的事件循环。
小结
在本篇文章中，我们学习了以下内容：
- 可以检查一个对象是不是协程，从而创建同时处理协程和常规函数的 API。
- 如果有需要在协程之间传递的状态，但希望这个状态独立于参数，可使用上下文局部变量。
- asyncio 的 sleep 协程可用于强制事件循环的迭代，当需要触发事件循环来完成些工作，但没有一个合适的等待点（await point）时，这很有帮助。
- asyncio 只是 Python 对事件循环的标准实现，除了 asyncio 之外还有其他实现，比如 uvloop，我们可按自己的意愿修改它们，但仍使用 async 和 await 语法。如果想设计一些具有不同特征的程序，从而更好地满足我们的需求，也可以创建自己的事件循环。
如果觉得文章对您有所帮助，可以请囊中羞涩的作者喝杯柠檬水，万分感谢，愿每一个来到这里的人都生活愉快，幸福美满。
微信赞赏
支付宝赞赏
