本文详细介绍了Python中的asyncio库和websockets库如何协同工作，包括异步编程的基本原理、asyncio的事件循环、协程的使用、Future和Task的概念，以及websockets在服务器和客户端的具体应用。
Python 中的 asyncio 和 websockets
这两天碰到了有关 websocket 开发的应用，使用的是 python 中的 websockets 来实现的。
因为大多数教学中的 websockets 使用异步方式工作，而 python 中的 websockets 是通过 asyncio 来实现异步操作的，而这两者在使用的时候往往由于教程过于简单，导致在面对一些复杂应用的时候自己会弄不清异步的处理逻辑。
因此记录一下 asyncio 和 websockets 工作的基本原理，便于以后使用两者开发更复杂的异步应用，或者理解复杂异步应用的逻辑。
引例
使用 websockets 的一个经典例子如下所示：
# server.py
import asyncio
import websockets
async def on_message(websocket, path):
async for message in websocket:
print(f"收到消息: {message}")
await websocket.send(f"你发送的消息是: {message}")
if __name__ == "__main__":
host = "localhost"
port = 8888
asyncio.get_event_loop().run_until_complete(websockets.serve(on_message, host , port))
asyncio.get_event_loop().run_forever()
# client.py
import asyncio
import websockets
async def hello(uri):
async with websockets.connect(uri) as websocket:
message = input("请输入消息: ")
await websocket.send(message)
response = await websocket.recv()
print(f"服务器响应: {response}")
if __name__ == "__main__":
host = "localhost"
port = 8888
asyncio.get_event_loop().run_until_complete(hello(f"ws://{host}:{port}))
从表现看来上面的例子非常直观：服务端启动，客户端连接客户端，并向客户端发送消息，服务端接收后向客户端回传一个消息，结束。
那具体发生了什么呢？首先要从 asyncio 说起
asyncio
asyncio 是 Python 3.6 引入的标准库内容 中用于实现异步编程的标准库，而协程则是异步编程中的一种核心概念。asyncio 提供了一个事件循环（Event Loop），允许在单线程中执行多个协程，处理异步任务、事件和 I/O 操作。
协程本质上仍然是单线程，比较适合需要大量等待的任务（网络通讯、文件 IO 等），否则不能提升程序的效率。
协程
Python 3.5 之前协程主要是生成器协程，使用生成器定义，用 yield 关键字让出控制权，并且可以暂停和恢复。
Python 3.6 之后协程主要是函数协程，使用 async def 关键字定义的函数对象，用 await 关键字等待异步操作完成。
本着用新不用旧的原则，下面的提及的都是 Python 3.6 之后的用法，但函数协程本质上也能当成生成器看待，只是 python 解释器提供更高级的方法来控制运行。
注：
await 之后只能接awaitable 对象，python 原生提供三类awaitable 对象：
- Coroutine - 协程，使用
async def 声明的异步函数- Future -
asyncio.Future 是一个实现了awaitable 接口的对象，代表了一个异步操作的结果。可以使用asyncio.ensure_future 或asyncio.create_task 创建Future 对象- Task -
asyncio.Task 是Future 的子类，表示一个异步操作的执行。通常使用asyncio.create_task 创建任务另外，任何实现了
__await__ 方法的类实例都是awaitable 对象。
举个🌰
下面是一个利用 asyncio 使用协程的简单例子[1]：
import asyncio
import time
async def say_after(delay, what):
await asyncio.sleep(delay)
print(what)
async def main():
print(f"started at {time.strftime('%x')}")
await say_after(1, 'hello')
await say_after(2, 'world')
print(f"finished at {time.strftime('%x')}")
asyncio.run(main())
其中的 asyncio.run 是程序从同步模式变成异步模式的入口，执行时会将 main()（协程对象）变成 task 注册到 event loop 中。而 event loop 会从中选取任务执行，当前只有 main，于是将控制权交给 main，执行 main 函数。
在遇到 await 时，会将 await 后的协程对象变成 task 再次注册到 event loop 中，控制权交还给 event loop。现在 event loop 中有两个任务，main 和 say_after(1, 'hello')，但 main 的继续执行要等 say_after(1, 'hello') 执行完成，say_after(1, 'hello') 可以继续执行，于是控制权交给 say_after(1, 'hello')，并执行它。
重复上述过程，say_after(1, 'hello') 中遇到 await 时，将 asyncio.sleep(1) 注册为 task，控制权交还给 event loop。现在 event loop 中有三个任务，main、say_after(1, 'hello') 和 asyncio.sleep(1)，但 main 的继续执行要等 say_after(1, 'hello') 执行完成，say_after(1, 'hello') 要等 asyncio.sleep(1) 执行完成。asyncio.sleep(1) 可以继续执行，于是控制权交给 asyncio.sleep(1)，并执行它。
asyncio.sleep(1) 1s 后结束，控制权再次交还给 event loop，根据上面的任务的依赖关系，继续执行 say_after(1, 'hello') 直到结束。控制权再次交还给 event loop，执行 main，遇到 await say_after(2, 'world')，重复上述过程，直到整个程序结束，耗时 3s。
所有控制权的转移都是显式的，event loop 不能强制从 task 中获取控制权。task 交还控制权的方式有两种：
-
await 一个awaitable 对象时 - 函数运行完毕
当然，上面的例子没有体现出异步的优势。因此可以结合 create_task 和 gather 一起使用，首先创建好任务，然后一起等待执行的结果，如下所示：
import asyncio
import time
async def say_after(delay, what):
await asyncio.sleep(delay)
return f"{what} - {delay}"
async def main():
task1 = asyncio.create_task(say_after(1, 'hello'))
task2 = asyncio.create_task(say_after(2, 'world'))
print(f"started at {time.strftime('%x')}")
ret = await asyncio.gather(task1, task2) # 可以获取 await 的返回值
print(ret)
print(f"finished at {time.strftime('%x')}")
asyncio.run(main())
整个程序耗时 2s。
asyncio.sleep 和time.sleep 不一样，后者是系统级的睡眠，是阻塞的，而前者是 asyncio 库实现的睡眠，可以多个asyncio.sleep 一起执行。或许这就是asyncio 可以异步以及比较适合需要等待的任务的原因？（个人推测）
核心源码
通过上面的例子理解了 event loop 控制异步的基本逻辑，然后再来深入一点了解一下 asyncio 库的核心源码[2][3]。
首先从上面例子的 asyncio.run，也就是程序变成异步的入口开始。大多数教程中，会使用以下方式作为异步入口使用：
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
事实上 asyncio.run 就是对上面内容的包装。
asyncio.run
asyncio.run 的核心逻辑如下：
首先创建了新的 event loop，然后用 run_until_complete 方法执行协程，
注意：该方法只能调用一次生成新的 event loop，每次只允许有一个 event loop 在运行
而 run_until_complete 方法核心逻辑如下：
主要有 4 步操作：
-
tasks.ensure_future() 转换成 Future 对象，如果是 Coroutine，会调用create_task 方法生成一个 Task 对象 - 将
_run_until_complete_cb 方法注册为 future 的回调方法，future 执行完成时会回调该方法，将当前 event loop 设置为结束状态，这样 event loop 在当前循环执行完成之后就会退出 - 调用
run_forever() （后面详细解释）进入当前 event loop 的主循环中 - 该协程执行完毕，当前 event loop 的主循环已跳出，在最后通过
future.result() 返回协程的执行结果
在深入了解 asyncio 源码前，先介绍其中几个关键概念：
- Future -
asyncio 的核心类之一，是 task 类的父类 - Task - event loop 的最小执行单位，所有 task 执行完毕，event loop 也就退出了
- EventLoop -
asyncio 的核心，也可以看成是一个“大脑”。对于其中注册的众多 task，event loop 决定着执行哪一个任务，并将程序控制权交给该任务运行
Future
主要包含三个属性和三个方法：
-
_loop - 所属的 event loop -
_callbacks - 保存添加到该 Future 对象中的回调方法 -
_result - 保存该任务的执行结果 -
add_done_callback 方法 - 用于注册完成后的回调函数 -
set_result 方法 - 用于在任务执行完成后设置执行结果，并且会调用__schedule_callbacks()方法 -
__schedule_callbacks() 方法 - 将注册的回调方法添加到 event loop 最近一次循环会执行的任务队列中
__schedule_callbacks() 通过call_soon 方法将 future 中注册的回调函数在 event loop 中注册一个 handle，这样在 event loop 的最近一次循环中会执行该回调方法。
Future 类的 __await__() 方法如下所示：
await 一个 future 对象时会返回这个对象本身，在 future 对象执行结束后会返回其中的 result。
Task
每个 task 类都会包括一个 coroutine，其中最重要的 __step 方法就是驱动协程一步步执行的，关键部分如下所示：
__step 方法通过调用 core.send() 触发协程开始执行或者从上一次 await 的地方继续执行。直到遇到下一个 await future，将 future 对象返回并保存在 result 中，然后将 __wakeup 方法注册为该 future 对象的回调方法。
当 future 完成之后就会继续回调这里 __wakeup 方法，__wakeup 方法内部其实就是再次调用了一次 __step 方法。如此循环直到协程执行完成，抛出 StopIteration 异常，将 future 的值通过 set_result 方式保存到 _result 中，该 task 执行完成。
也就是说当该
future 对象执行结束后（StopIteration），会唤醒调用await future 的这个awaitable（依赖于该future、等它结束后才能开始），实际上是通过set_result 方法，将该future 对象上的callbacks（这里的__wakeup）都放入到 event loop 的最近一次循环会执行的任务队列中（后面提及的_ready）
而 Task 的构造方法就将 __step 方法注册到了 event loop 中，这样在 event loop 的下一个循环就会触发该 task 中 awaitable 的执行。
那 EventLoop 内部具体是什么什么执行流程呢？看一下 EventLoop 内部的实现
EventLoop
EventLoop 中维护了两个队列，_scheduled 和 _ready，前者存放未来某个时间会执行的任务，通过 EventLoop 的 call_later、call_at 等方法添加任务，后者存放了最近一次循环会执行的任务，可通过 call_soon 等方法直接往其中添加任务。
EventLoop 还支持通过
add_reader 和add_writer 方法向selector 中注册读写事件和回调，事件发生时就会触发回调方法的调用
回到前面提到的调用 run_forever，核心源码如下所示：
该方法会一直循环执行 _run_once 方法，直到 event loop 的 _stopping 属性被设置为 True 就跳出循环，而 _run_once 方法源码如下所示：
主要执行了以下几个步骤：
-
先将 event loop 中用于存放待调度的任务队列
_scheduled 进行建堆操作，这样可以按照任务优先级进行任务调度 -
找出最优先任务的等待时间 timeout，调用
event_list = self._selector.select(timeout) 等待事件的发生或者超时 -
调用
self._process_events(event_list) 方法处理返回的事件列表，该方法中会依次将每个事件对应的回调方法添加到 event loop 的_ready 任务队列中 -
有的回调并不是通过向
selector 中注册事件的方式，而是根据时间直接添加到任务队列中的（如asyncio.sleep），所以还需要将_scheduled 队列中的满足执行时间要求的任务添加到_ready 队列中，_ready 队列存放的即为当前循环需要执行的任务 -
针对
_ready 队列中的每个 handle 任务均执行handle._run() 方法，触发回调函数执行handle 对象除了封装回调函数及其方法参数，还维护了该任务的状态等，调用 handle 对象的
_run 方法即可执行回调
至此就梳理了协程的大致工作流程，通过 Task、Future 和 EventLoop 三者配合，驱动 awaitable 一步步执行的。
websockets
python 中 websockets 的服务器和客户端都是使用 asyncio 实现的[5]，也就是说其中的 server 和 client 都是实现了 __await__ 方法的 awaitable 对象。
Server
websockets 的 server 一般使用以下方式启动：
stop = asyncio.Future() # set this future to exit the server
server = await serve(...)
await stop
await server.close()
server 也实现了异步上下文管理方法（__aenter__ 和 __aexit__），因此也能使用以下方法启动：
stop = asyncio.Future() # set this future to exit the server
async with serve(...):
await stop
在 asyncio 中提到，Future 对象被设置 result 之后就会唤醒 await 该 future 对象的 awaitable，因此能通过这种方式结束服务器的运行。
下面就是利用这种方式主动停止服务器的例子：
#!/usr/bin/env python
import asyncio
import signal
import websockets
async def echo(websocket):
async for message in websocket:
await websocket.send(message)
async def server():
# Set the stop condition when receiving SIGTERM.
loop = asyncio.get_running_loop()
stop = loop.create_future()
loop.add_signal_handler(signal.SIGINT, stop.set_result, None) # ^C 停掉服务器
async with websockets.serve(echo, "localhost", 8765):
await stop
asyncio.run(server())
通过给服务器指定 ws_handler，可以处理客户端连接和发送的消息：
async def handler(websocket, path):
async for message in websocket:
print(message)
stop = asyncio.Future() # set this future to exit the server
async with serve(handler, host, port):
await stop
上面 handler 中的 webscoket 即建立的 websocket 连接， async for 是异步循环，会从连接的客户端中不断获取信息，直到连接断开或发生错误才会退出循环。
Client
websockets 的 client 一般使用以下方式连接：
stop = asyncio.Future() # set this future to exit the client
client = await websockets.connect(...)
await stop
await client.close()
同理，client 也实现了异步上下文管理方法，因此也能使用以下方法连接：
async with websockets.connect(...) as websocket:
...
同样也可以用作无限异步迭代器，以便在出现错误时自动重新连接：
async for websocket in websockets.connect(...):
try:
...
except websockets.ConnectionClosed:
continue
回望
回到开始的引例，具体发生了什么呢？
首先看服务端，line 13 中 websockets.serve(on_message, host , port) 会返回一个 awaitable。本行的 asyncio.get_event_loop() 获取了当前的 event loop（没有的话创建一个），并用 run_until_complete 方法这个对象注册成了 task，并开始运行。
在 event loop 中调用 websockets.serve(on_message, host , port) 就返回了一个 websocket server 对象，也即在给定 host 和 port 运行了一个服务器，其收到连接的回调函数是 on_message。
line 13 行仅仅进行了上述动作就结束了，并没有让服务器一直运行，如果没有后续代码，就结束了，因为 awaitable 形成的 task 执行结束了，event loop 中没有其他 task 存在，因此会直接结束。但 line 14 运行了 asyncio.get_event_loop().run_forever()。在前面提到了 run_forever 源码，解释了它会不断运行 _run_once 直到 _stopping 这个标志位为 True 的时候才会退出（如调用 stop()）方法。
由于 event loop 执行完建立服务器的 task 之后，其中没有其他任务可以运行了，因此 _stopping 这个标志位一直为 False，使得 event loop 一直在运行，其中返回的 websocket server 对象也就一直存在，形成了服务器一直在运行状态的现象。
这种方法的缺点也很明显：
- 由于
_stopping 标志位是 event loop 自己管理设置的变量，run_forever 之后没法手动关闭服务器了（Linux 系统）- 服务器只能处理一个 client 连接
再来看客户端，同上面的道理分析，在接收到服务器返回的消息过后，event loop 中的 task 运行结束后，没有其余操作，因此 event loop 直接结束，程序完成。
事实上，get_event_loop()、run_until_complete、run_forever() 这些都是 asyncio 的低层级 API，在开发应用时应该尽量使用高层级 API，如 asyncio.run，这样可以让自己避开这些操作，而让 asyncio 自己决定底层逻辑的跳转[4]。
使用 asyncio 和 websockets 的简单聊天室应用
服务端代码：
# server.py
import signal
import asyncio
import platform
import websockets
class ChatServer:
def __init__(self):
self.clients = set()
async def handle_connection(self, websocket, path):
# 将新连接的客户端添加到集合中
self.clients.add(websocket)
try:
await websocket.send("欢迎加入聊天室！")
# 处理消息
async for message in websocket:
# 广播消息给所有连接的客户端
await asyncio.gather(*[client.send(message) for client in self.clients])
except websockets.exceptions.ConnectionClosed:
pass
finally:
# 移除断开连接的客户端
self.clients.remove(websocket)
async def close_server(self):
# 关闭所有连接
await asyncio.gather(*[client.close() for client in self.clients])
# 关闭 WebSocket 服务器
await self.server.close()
async def run_server(self):
# 启动 WebSocket 服务器，监听在 localhost 的 8765 端口
self.server = await websockets.serve(
self.handle_connection, "localhost", 8765
)
print("WebSocket 服务器启动在 ws://localhost:8765")
try:
# 等待关闭服务器的信号
if platform.system().lower() == "windows":
while True:
await asyncio.sleep(0.1)
else:
# 注册关闭服务器的信号处理函数
loop = asyncio.get_event_loop()
stop = loop.create_future()
loop.add_signal_handler(signal.SIGINT, stop.set_result, None)
await stop
finally:
await self.close_server()
if __name__ == "__main__":
chat_server = ChatServer()
# 启动服务器
asyncio.run(chat_server.run_server())
客户端代码：
# client.py
import asyncio
import websockets
async def read_messages(websocket):
try:
while True:
message = await websocket.recv()
print(f"收到消息: {message}")
if message.lower() == "exit":
break
except websockets.exceptions.ConnectionClosed:
print("与服务器的连接已关闭.")
async def send_messages(websocket):
try:
loop = asyncio.get_event_loop()
print("请输入消息 (输入 'exit' 退出): \n")
while True:
message = await loop.run_in_executor(None, input)
await websocket.send(message)
if message.lower() == "exit":
break
except websockets.exceptions.ConnectionClosed:
print("与服务器的连接已关闭.")
async def main():
uri = "ws://localhost:8765"
async with websockets.connect(uri) as websocket:
await asyncio.gather(
read_messages(websocket),
send_messages(websocket)
)
if __name__ == "__main__":
asyncio.run(main())
参考资料
- 【python】asyncio的理解与入门，搞不明白协程？看这个视频就够了。_哔哩哔哩_bilibili
- 【python】await机制详解。再来个硬核内容，把并行和依赖背后的原理全给你讲明白_哔哩哔哩_bilibili
- Python asyncio库核心源码解析_深入理解 python 异步编程 中-CSDN博客
- 事件循环 — Python 3.12.1 文档
- API reference - websockets 12.0 documentation
扫一扫
1422
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
