TestClient slows with multiple calls to websocket_connect

[NPrescott]

I noticed an anomalous slow down in a test in which multiple websocket_connect calls are made using TestClient. In particular this was most egregious on a single machine but I've validated a more modest slow down across multiple machines.

I've pared back the application and test to this:

import unittest

import starlette.applications
import starlette.websockets
import starlette.routing
import starlette.testclient


async def websocket_endpoint(websocket):
    await websocket.accept()
    try:
        while True:
            await websocket.receive_json()
    except starlette.websockets.WebSocketDisconnect as e:
        print(e)
    finally:
        await websocket.close()

app = starlette.applications.Starlette(debug=True,
                                       routes=[starlette.routing.WebSocketRoute('/ws', websocket_endpoint)],
                                       on_startup=[])


class TestWebsocket(unittest.TestCase):

    def test_websocket(self):

        client = starlette.testclient.TestClient(app)

        with (client.websocket_connect("/ws"),
              client.websocket_connect("/ws")):
            self.assertTrue(True)

In the case where I noticed the problem there is a dramatic increase in execution time between a test with one call to client.websocket_connect and two. For the above:

$ ./venv/bin/python -m unittest test/test_two_connects.py
(1000, None)
(1000, None)
.
----------------------------------------------------------------------
Ran 1 test in 19.902s

OK

If I remove the second call to websocket_connect:

$ ./venv/bin/python -m unittest test/test_one_connect.py
(1000, None)
.
----------------------------------------------------------------------
Ran 1 test in 0.122s

OK

Specifics of the environment:

$ ./venv/bin/pip freeze
anyio==4.3.0
certifi==2024.2.2
h11==0.14.0
httpcore==1.0.5
httpx==0.27.0
idna==3.7
sniffio==1.3.1
starlette==0.37.2

$ ./venv/bin/python -V
Python 3.12.2

$ uname -a
Linux fedora39-host 6.7.10-200.fc39.x86_64 #1 SMP PREEMPT_DYNAMIC Mon Mar 18 18:56:52 UTC 2024 x86_64 GNU/Linux

I've been unable to identify what makes the machine notable compared to others that don't demonstrate the same severity. In order to observe an increase in execution time on other machines I've had to include 20 calls to websocket_connect but an impact is observable.

I spent some time trying to investigate the source of the slowness and turned up mostly time spent on futex, I think linked to the queues backing send and receive in the TestClient. I was able to at least partly confirm issues with the scheduling by pinning the test to a single CPU (where the machine has 8 total), no particular CPU matters for this case:

$ taskset -c 0 ./venv/bin/python -m unittest test/test_two_connects.py 
(1000, None)
(1000, None)
.
----------------------------------------------------------------------
Ran 1 test in 0.005s

OK

Trying to identify potential sources of CPU contention or difficulty I turned up starlette.testclient.WebSocketTestSession._asgi_receive which is testing the receive queue and conditionally invoking anyio.sleep(0). In my case I believe this will resolve to a special coded case of asyncio that immediately yields to the TestClient threaded event loop.

I don't totally grasp the interaction causing a problem here but suspecting the loop was spinning on the lock backing the Queue I tried modifying _asgi_receive like this:

import anyio


async def _asgi_receive(self):
    while self._receive_queue.empty():
        await anyio.sleep(0.001)    # avoid special-coding on asyncio.sleep(0)
    return self._receive_queue.get()

starlette.testclient.WebSocketTestSession._asgi_receive = _asgi_receive

And observed an improvement in execution speed to a level I expect:

$ ./venv/bin/python -m unittest test/test_two_connects.py 
.
----------------------------------------------------------------------
Ran 1 test in 0.008s

OK

As mentioned, this effect issue seems less pronounced on other machines but observable with tens of calls to websocket_connect. In such a case on a less problematic machine the above patch produces a more modest speed up in the range of 10x.

A search of the discussion history turned up Testing interaction of WebSockets with TestClient imposes thread safety? which did not seem directly applicable and has been noted to no longer be a problem with recent releases.

I am filing this under "Potential Issue" because I am unable to identify why one machine in particular is so susceptible to the problem and I have two different solutions available. I suspect this can mostly serve as documentation in case anyone else tries searching for the problem like I did.

[Olocool17]

Hi! We've noticed the exact same issue here: slowdowns on websocket tests that are particularly noticeable on one specific machine ( in our case Linux 5.15.0-94-generic #104-Ubuntu). All of our websocket tests are performed with two websockets, so I did not think quantity of websocket connections would matter, but it clearly does.

I can verify everything here, though I too haven't been able to identify exactly why particular machines are more susceptible than others. Are the machines you've experienced less/no slowdown on all Windows machines by chance? That has certainly been our experience.

To my knowledge, asyncio.sleep(0) is often interpreted as 'yield to another task', but there is no guarantee that the event loop will not just pick the same task again, effectively causing the 'spinning' behavior we've observed. The solution is not to rely on asyncio.sleep(0) but use synchronisation primitives instead.

I've opened a PR (#2597) with my fix, but unfortunately only found your discussion after the fact. Thank you for the write up!

[NPrescott]

Thanks for the patch.

The machine I have experienced issues on runs Fedora 39 on an Intel i7-9700 (cpu family 6, model 158).
The reference machine without issues I've used runs Fedora 39 on an AMD Ryzen 7 5700U (cpu family 23, model 104).
The output of python -m sysconfig is identical between machines.
The problem machine does not support SMT but disabling it on a "fast" machine has no negative effect.

My best guess is that the performance difference comes from a combination of CPU thread scheduling and frequency. If I raise the CPU frequency of the problem machine by running stress -c 8 the performance improves (though not completely).