为何向WebSocket发送消息不会将控制权交还给事件循环？

代码示例 #

main.py #

import asyncio
import websockets


async def echo(websocket):
    async for message in websocket:
        await websocket.send(message)
        print(message)


async def main():
    async with websockets.serve(echo, "localhost", 8765):
        await asyncio.Future()  # run forever

if __name__ == '__main__':
    asyncio.run(main())

other.py #

import asyncio
import json

import websockets

tasks = set()


async def run_job(i):
    await asyncio.sleep(0.)
    print(f"I'm job number {i}")


async def bunch_of_tasks(ws):
    for i in range(10):
        task = asyncio.create_task(run_job(i), name=f'job-{i}')
        tasks.add(task)
        task.add_done_callback(tasks.discard)
        print(f'had a nice sleep! now my value is {i}')
        # await asyncio.sleep(0.)
        await ws.send(json.dumps('hello there!'))
    await asyncio.gather(*tasks)
    print(f'tasks done')


async def do_stuff():
    async with websockets.connect("ws://localhost:8765") as websocket:
        await bunch_of_tasks(websocket)
        await websocket.recv()

if __name__ == '__main__':
    asyncio.run(do_stuff())

两种运行结果 #

未取消注释await asyncio.sleep(0.)时的输出 #

had a nice sleep! now my value is 0
had a nice sleep! now my value is 1
had a nice sleep! now my value is 2
had a nice sleep! now my value is 3
had a nice sleep! now my value is 4
had a nice sleep! now my value is 5
had a nice sleep! now my value is 6
had a nice sleep! now my value is 7
had a nice sleep! now my value is 8
had a nice sleep! now my value is 9
I'm job number 0
I'm job number 1
I'm job number 2
I'm job number 3
I'm job number 4
I'm job number 5
I'm job number 6
I'm job number 7
I'm job number 8
I'm job number 9
tasks done

取消注释await asyncio.sleep(0.)时的输出 #

had a nice sleep! now my value is 0
had a nice sleep! now my value is 1
I'm job number 0
had a nice sleep! now my value is 2
I'm job number 1
had a nice sleep! now my value is 3
I'm job number 2
had a nice sleep! now my value is 4
I'm job number 3
had a nice sleep! now my value is 5
I'm job number 4
had a nice sleep! now my value is 6
I'm job number 5
had a nice sleep! now my value is 7
I'm job number 6
had a nice sleep! now my value is 8
I'm job number 7
had a nice sleep! now my value is 9
I'm job number 8
I'm job number 9
tasks done

问题核心 #

为什么WebSocket发送消息时不会把控制权交还给事件循环，导致run_job协程没法执行，而asyncio.sleep(0.)却能挂起当前任务，让run_job有执行机会？

解答 #

原因得从AsyncIO的任务调度规则和WebSocket发送操作的实际执行逻辑来解释：

1. AsyncIO的调度是“协作式”的
   AsyncIO不会主动抢占任务的控制权，只有当前任务主动“喊停”，事件循环才会切换去执行其他已准备好的任务。asyncio.sleep(0.)就是个专门用来触发切换的操作——它不会真的让程序休眠，但会明确告诉事件循环：“我现在可以暂停，你先去跑别的任务吧”，直接触发一次任务调度。

2. WebSocket发送操作多数情况下是同步完成的
   websockets.send()的底层实现中，如果发送缓冲区还有空闲空间，消息会直接写入缓冲区，整个过程瞬间完成，不会触发异步等待。这种情况下，当前的bunch_of_tasks协程根本不会挂起，事件循环也就没机会执行之前创建好的run_job任务。直到10次循环全部跑完，执行到await asyncio.gather(*tasks)时，才会等待所有run_job任务完成，这时它们就批量执行了。

3. 加上sleep后的变化
   一旦取消注释await asyncio.sleep(0.)，每次循环迭代都会主动触发事件循环切换任务。这时候事件循环会发现run_job任务已经就绪（因为run_job里的await asyncio.sleep(0.)也会让出控制权，第一次切换时就会被执行），所以会交替运行bunch_of_tasks的迭代和run_job任务，也就出现了交错输出的情况。

总结一下：ws.send()在缓冲区够用时是同步完成的，没给事件循环切换任务的机会；而asyncio.sleep(0.)是强制让当前任务暂停，给了其他任务执行的时间窗口。

内容的提问来源于stack exchange，提问作者OES