[FIXED] Why sending a message to a web socket does not yield control to the event loop?

Issue

Consider the following code:

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())

By first running main.py and then running in parallel other.py, I get:

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

But if I uncomment await asyncio.sleep(0.) before await ws.send(json.dumps('hello there!')), I get:

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

which is somehow what I would expect.

So apparently sending the message to the web socket does not yield control to the event loop, and the run_job coroutine does not have the opportunity to be run. However, asyncio.sleep effectively suspends the current task and gives an opportunity to run_job to be executed.

Why that happens?

Solution

TL;DR: indeed - await is not in itself sufficient to warrant control is passed back to the cooperative part of the asyncio loop. To ensure that, a low-level callback has to be scheduled somewhere down the call chain (and asyncio.sleep does that).

Long answer:

I had to debug that out - It happens that although everything is in place for websockets clients being async, it boils down to write all data it is sending in the Selector socket at once - since it is non-constrained otherwise.

In other words, ws.send will eventually call this line synchronously: https://github.com/python/cpython/blob/f2e5a6ee628502d307a97f587788d7022a200229/Lib/asyncio/selector_events.py#L1071

And then, the big surprise is that for raw-coroutines (co-routines not wrapped in tasks or futures), whenever they return their value, the execution does not yield to the asyncio-loop - and even if they contain other "await" statements, if the nested awaited co-routines do not actually would "block" in an await, the asyncio loop is never reached back. Internally, in the C code handling co-routines, whenever a co-routine actually "blocks" after a send, a callback is scheduled: this callback is wrapped in a Python asyncio.events.Handle object, and the control is then returned to the asyncio loop.

The code for that is in this C function: https://github.com/python/cpython/blob/f2e5a6ee628502d307a97f587788d7022a200229/Modules/_asynciomodule.c#L2696 . If you follow along there, you can see that if the co-routine returned a value, its result is set in the low-level object. If it returned from send with another awaitable, that one is scheduled in the loop, and the function returns.

It may be (and probably is) an implementation choice of asyncio Loops: any so called co-routine that actually resolves synchronously is run synchronously, for performance reasons.

Only when any of the nested co-routines called in an await ends up scheduling a future with a callback, in the leaf call, (which asyncio.sleep does), do the asyncio default loop runs through all other ready tasks. (another full execution of its ._run_once method in asyncio.base_events.BaseEventLoop).

Answered By - jsbueno