1#!/usr/bin/env python
2
3import asyncio
4import statistics
5import tracemalloc
6
7import websockets
8from websockets.extensions import permessage_deflate
9
10
11CLIENTS = 10
12INTERVAL = 1 / 10  # seconds
13
14MEM_SIZE = []
15
16
17async def mem_client(client):
18    # Space out connections to make them sequential.
19    await asyncio.sleep(client * INTERVAL)
20
21    tracemalloc.start()
22
23    async with websockets.connect(
24        "ws://localhost:8765",
25        extensions=[
26            permessage_deflate.ClientPerMessageDeflateFactory(
27                server_max_window_bits=10,
28                client_max_window_bits=10,
29                compress_settings={"memLevel": 3},
30            )
31        ],
32    ) as ws:
33        await ws.send("hello")
34        await ws.recv()
35
36        await ws.send(b"hello")
37        await ws.recv()
38
39        MEM_SIZE.append(tracemalloc.get_traced_memory()[0])
40        tracemalloc.stop()
41
42        # Hold connection open until the end of the test.
43        await asyncio.sleep(CLIENTS * INTERVAL)
44
45
46asyncio.get_event_loop().run_until_complete(
47    asyncio.gather(*[mem_client(client) for client in range(CLIENTS + 1)])
48)
49
50# First connection incurs non-representative setup costs.
51del MEM_SIZE[0]
52
53print(f"µ = {statistics.mean(MEM_SIZE) / 1024:.1f} KiB")
54print(f"σ = {statistics.stdev(MEM_SIZE) / 1024:.1f} KiB")
55