⚡️ Speed up method WebSocket.receive by 6%

[codeflash-ai]

📄 6% (0.06x) speedup for WebSocket.receive in starlette/websockets.py

⏱️ Runtime : 2.49 milliseconds → 2.35 milliseconds (best of 131 runs)

📝 Explanation and details

The optimization achieves a 6% runtime improvement and 4% throughput increase by restructuring the conditional logic in the WebSocket's receive method to reduce branching overhead and improve CPU instruction flow.

Key optimizations applied:

1. Eliminated elif chain: Changed elif self.client_state == WebSocketState.CONNECTED: to a separate if statement, reducing nested branching that can hurt CPU branch prediction.

2. Optimized hot path for CONNECTED state: In the CONNECTED state (the most frequent case based on profiler data showing 3,747 hits), the code now uses two separate if statements for message type checking instead of a combined not in check followed by another if. This creates a more direct execution path:

   • First checks for websocket.disconnect (less common, ~203 hits)
   • Then checks for websocket.receive (most common, ~3,540 hits)
   • Only raises the error if neither condition matches

3. Reduced computational overhead: The original code used message_type not in {"websocket.receive", "websocket.disconnect"} which requires set membership testing, followed by another message_type == "websocket.disconnect" check. The optimized version eliminates the set lookup and performs direct string comparisons.

Performance impact by test type:

• Basic operations (connect/receive/disconnect sequences): The optimization particularly benefits the most common case of receiving regular messages in the CONNECTED state
• Large-scale tests (100+ messages): Show consistent improvements as the optimized CONNECTED state path is exercised repeatedly
• Concurrent operations: Benefit from reduced branching overhead when multiple WebSocket instances process messages simultaneously

The line profiler confirms that the CONNECTED state processing (lines handling message_type checks) shows reduced execution time, directly contributing to the overall throughput improvement of ~20,000 additional operations per second.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 1583 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

from __future__ import annotations

import asyncio  # used to run async functions
from collections.abc import Iterator, Mapping
from typing import TYPE_CHECKING, Any, Callable, Dict

import pytest  # used for our unit tests
from starlette.websockets import WebSocket


# Minimal WebSocketState enum for testing
class WebSocketState:
    CONNECTING = "connecting"
    CONNECTED = "connected"
    DISCONNECTED = "disconnected"


class HTTPConnection(Mapping[str, Any]):
    def __init__(self, scope, receive=None):
        self.scope = scope

    def __getitem__(self, key: str) -> Any:
        return self.scope[key]

    def __iter__(self) -> Iterator[str]:
        return iter(self.scope)

    def __len__(self) -> int:
        return len(self.scope)

    __eq__ = object.__eq__
    __hash__ = object.__hash__
from starlette.websockets import WebSocket

# --- Helper functions for test receive coroutines ---

def make_receive_coroutine(messages):
    """
    Returns an async function that yields the next message from the provided list each time it's called.
    """
    messages_iter = iter(messages)
    async def _receive():
        try:
            return next(messages_iter)
        except StopIteration:
            # Simulate ASGI server behavior: raise error if no more messages
            raise RuntimeError("No more messages")
    return _receive

# --- Basic Test Cases ---

@pytest.mark.asyncio
async def test_receive_connect_message():
    """Test that receive returns the connect message and transitions state."""
    scope = {"type": "websocket"}
    messages = [{"type": "websocket.connect"}]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    msg = await ws.receive()

@pytest.mark.asyncio
async def test_receive_receive_message_after_connect():
    """Test that receive returns a receive message after a connect."""
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.receive", "text": "hello"}
    ]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    msg1 = await ws.receive()
    msg2 = await ws.receive()

@pytest.mark.asyncio
async def test_receive_disconnect_message_after_connect():
    """Test that receive returns a disconnect message and transitions state."""
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.disconnect", "code": 1000}
    ]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    msg1 = await ws.receive()
    msg2 = await ws.receive()

# --- Edge Test Cases ---

@pytest.mark.asyncio
async def test_receive_invalid_first_message_type():
    """Test that receive raises if the first message is not websocket.connect."""
    scope = {"type": "websocket"}
    messages = [{"type": "websocket.receive"}]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_invalid_message_type_after_connect():
    """Test that receive raises if a message after connect is not receive or disconnect."""
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.foo"}
    ]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    await ws.receive()  # Should succeed
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_after_disconnect_raises():
    """Test that receive raises after disconnect message has been received."""
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.disconnect"}
    ]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    await ws.receive()  # connect
    await ws.receive()  # disconnect
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_no_messages_left():
    """Test that receive raises if no messages are left to receive."""
    scope = {"type": "websocket"}
    messages = [{"type": "websocket.connect"}]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    await ws.receive()
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_concurrent_execution():
    """Test concurrent execution of multiple WebSocket.receive coroutines."""
    scope = {"type": "websocket"}
    messages1 = [
        {"type": "websocket.connect"},
        {"type": "websocket.receive", "text": "msg1"},
    ]
    messages2 = [
        {"type": "websocket.connect"},
        {"type": "websocket.disconnect", "code": 1001},
    ]
    ws1 = WebSocket(scope, make_receive_coroutine(messages1), None)
    ws2 = WebSocket(scope, make_receive_coroutine(messages2), None)
    # Run both receive sequences concurrently
    results = await asyncio.gather(ws1.receive(), ws2.receive())
    # Next messages concurrently
    results2 = await asyncio.gather(ws1.receive(), ws2.receive())

# --- Large Scale Test Cases ---

@pytest.mark.asyncio
async def test_receive_large_number_of_messages():
    """Test receive with a large number of websocket.receive messages."""
    scope = {"type": "websocket"}
    N = 100
    messages = [{"type": "websocket.connect"}] + [
        {"type": "websocket.receive", "text": f"msg{i}"} for i in range(N)
    ]
    ws = WebSocket(scope, make_receive_coroutine(messages), None)
    # First message should be connect
    msg = await ws.receive()
    # Next N messages should be websocket.receive
    for i in range(N):
        msg = await ws.receive()
    # After all messages, calling receive should raise
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_concurrent_many_websockets():
    """Test concurrent execution of many WebSocket.receive coroutines."""
    scope = {"type": "websocket"}
    N = 50
    websockets = []
    for i in range(N):
        messages = [
            {"type": "websocket.connect"},
            {"type": "websocket.receive", "text": f"ws{i}"},
            {"type": "websocket.disconnect", "code": 1000 + i}
        ]
        ws = WebSocket(scope, make_receive_coroutine(messages), None)
        websockets.append(ws)
    # Connect all
    connects = await asyncio.gather(*(ws.receive() for ws in websockets))
    # Receive all
    receives = await asyncio.gather(*(ws.receive() for ws in websockets))
    for i, msg in enumerate(receives):
        pass
    # Disconnect all
    disconnects = await asyncio.gather(*(ws.receive() for ws in websockets))
    for i, msg in enumerate(disconnects):
        pass

# --- Throughput Test Cases ---

@pytest.mark.asyncio
async def test_receive_throughput_small_load():
    """Test throughput with a small number of concurrent WebSocket connections."""
    scope = {"type": "websocket"}
    N = 10
    websockets = []
    for i in range(N):
        messages = [
            {"type": "websocket.connect"},
            {"type": "websocket.receive", "text": f"small{i}"},
            {"type": "websocket.disconnect", "code": 2000 + i}
        ]
        ws = WebSocket(scope, make_receive_coroutine(messages), None)
        websockets.append(ws)
    # Run all receive sequences concurrently and check results
    results = []
    for ws in websockets:
        results.append(await ws.receive())
        results.append(await ws.receive())
        results.append(await ws.receive())
    for i in range(N):
        pass

@pytest.mark.asyncio
async def test_receive_throughput_medium_load():
    """Test throughput with a medium number of concurrent WebSocket connections."""
    scope = {"type": "websocket"}
    N = 50
    websockets = []
    for i in range(N):
        messages = [
            {"type": "websocket.connect"},
            {"type": "websocket.receive", "text": f"medium{i}"},
            {"type": "websocket.disconnect", "code": 3000 + i}
        ]
        ws = WebSocket(scope, make_receive_coroutine(messages), None)
        websockets.append(ws)
    # Connect all concurrently
    connects = await asyncio.gather(*(ws.receive() for ws in websockets))
    # Receive all concurrently
    receives = await asyncio.gather(*(ws.receive() for ws in websockets))
    for i, msg in enumerate(receives):
        pass
    # Disconnect all concurrently
    disconnects = await asyncio.gather(*(ws.receive() for ws in websockets))
    for i, msg in enumerate(disconnects):
        pass

@pytest.mark.asyncio

#------------------------------------------------
import asyncio  # used to run async functions
from typing import Any, Callable, Dict

import pytest  # used for our unit tests
from starlette.websockets import WebSocket


# Minimal Enum for WebSocketState to match function requirements
class WebSocketState:
    CONNECTING = "connecting"
    CONNECTED = "connected"
    DISCONNECTED = "disconnected"

# Minimal Message type for type checking
Message = Dict[str, Any]
Receive = Callable[[], asyncio.Future]
Send = Callable[[Message], asyncio.Future]
Scope = Dict[str, Any]

# --- Function under test (EXACT COPY) ---
class HTTPConnection:
    def __init__(self, scope: Scope, receive: Receive | None = None) -> None:
        self.scope = scope

    def __getitem__(self, key: str) -> Any:
        return self.scope[key]

    def __iter__(self):
        return iter(self.scope)

    def __len__(self) -> int:
        return len(self.scope)

    __eq__ = object.__eq__
    __hash__ = object.__hash__
from starlette.websockets import WebSocket

# --- Helper functions for mocking ASGI receive/send ---

def make_receive(messages):
    """
    Returns an async function that yields messages from the provided list in order.
    """
    messages = list(messages)
    async def _receive():
        if not messages:
            raise RuntimeError("No more messages to receive")
        return messages.pop(0)
    return _receive

async def dummy_send(message):
    pass  # Not used in these tests

# --- UNIT TESTS ---

# 1. Basic Test Cases

@pytest.mark.asyncio
async def test_receive_connect_message():
    """
    Test that receive returns the 'websocket.connect' message and transitions state to CONNECTED.
    """
    scope = {"type": "websocket"}
    messages = [{"type": "websocket.connect"}]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    msg = await ws.receive()

@pytest.mark.asyncio
async def test_receive_receive_message_after_connect():
    """
    Test that after connect, receive returns 'websocket.receive' message and stays CONNECTED.
    """
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.receive", "text": "hello"}
    ]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    msg1 = await ws.receive()
    msg2 = await ws.receive()

@pytest.mark.asyncio
async def test_receive_disconnect_message_after_connect():
    """
    Test that after connect, receive returns 'websocket.disconnect' message and transitions to DISCONNECTED.
    """
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.disconnect"}
    ]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    msg1 = await ws.receive()
    msg2 = await ws.receive()

# 2. Edge Test Cases

@pytest.mark.asyncio
async def test_receive_invalid_connect_message_type():
    """
    Test that receive raises RuntimeError if first message is not 'websocket.connect'.
    """
    scope = {"type": "websocket"}
    messages = [{"type": "websocket.receive"}]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_invalid_message_type_after_connect():
    """
    Test that receive raises RuntimeError if message type after connect is not valid.
    """
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.foobar"}
    ]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    await ws.receive()  # valid connect
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_after_disconnect_raises():
    """
    Test that calling receive after disconnect raises RuntimeError.
    """
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.disconnect"}
    ]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    await ws.receive()  # connect
    await ws.receive()  # disconnect
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()

@pytest.mark.asyncio
async def test_receive_concurrent_connect_and_receive():
    """
    Test concurrent calls to receive: each instance maintains its own state.
    """
    scope = {"type": "websocket"}
    messages1 = [
        {"type": "websocket.connect"},
        {"type": "websocket.receive", "text": "one"}
    ]
    messages2 = [
        {"type": "websocket.connect"},
        {"type": "websocket.receive", "text": "two"}
    ]
    ws1 = WebSocket(scope, make_receive(list(messages1)), dummy_send)
    ws2 = WebSocket(scope, make_receive(list(messages2)), dummy_send)
    # Run both receive sequences concurrently
    results = await asyncio.gather(ws1.receive(), ws2.receive())
    # Next, both should be able to receive their own messages
    results2 = await asyncio.gather(ws1.receive(), ws2.receive())

@pytest.mark.asyncio
async def test_receive_multiple_disconnects_edge_case():
    """
    Test that multiple disconnects are handled: only first transition is allowed.
    """
    scope = {"type": "websocket"}
    messages = [
        {"type": "websocket.connect"},
        {"type": "websocket.disconnect"},
        {"type": "websocket.disconnect"}
    ]
    ws = WebSocket(scope, make_receive(messages), dummy_send)
    await ws.receive()  # connect
    await ws.receive()  # disconnect
    with pytest.raises(RuntimeError) as excinfo:
        await ws.receive()  # second disconnect not allowed

# 3. Large Scale Test Cases

@pytest.mark.asyncio
async def test_receive_large_number_of_messages():
    """
    Test that receive can handle a large sequence of valid messages.
    """
    scope = {"type": "websocket"}
    # First message is connect, followed by many receives, then disconnect
    messages = [{"type": "websocket.connect"}] + [
        {"type": "websocket.receive", "text": f"msg{i}"} for i in range(100)
    ] + [{"type": "websocket.disconnect"}]
    ws = WebSocket(scope, make_receive(list(messages)), dummy_send)
    # First receive is connect
    msg = await ws.receive()
    # Next 100 receives
    for i in range(100):
        msg = await ws.receive()
    # Final disconnect
    msg = await ws.receive()

@pytest.mark.asyncio
async def test_receive_concurrent_large_scale():
    """
    Test multiple WebSocket instances concurrently, each handling their own message sequence.
    """
    scope = {"type": "websocket"}
    n = 10
    ws_list = []
    for i in range(n):
        messages = [{"type": "websocket.connect"}] + [
            {"type": "websocket.receive", "text": f"ws{i}-msg{j}"} for j in range(10)
        ] + [{"type": "websocket.disconnect"}]
        ws = WebSocket(scope, make_receive(list(messages)), dummy_send)
        ws_list.append(ws)
    # Run all connect receives concurrently
    results = await asyncio.gather(*(ws.receive() for ws in ws_list))
    # Run all 10 message receives concurrently for each ws
    for j in range(10):
        results = await asyncio.gather(*(ws.receive() for ws in ws_list))
        for i, msg in enumerate(results):
            pass
    # Run all disconnects concurrently
    results = await asyncio.gather(*(ws.receive() for ws in ws_list))

# 4. Throughput Test Cases

@pytest.mark.asyncio
async def test_receive_throughput_small_load():
    """
    Throughput test: small load of 5 websocket connections, each with 5 messages.
    """
    scope = {"type": "websocket"}
    n = 5
    ws_list = []
    for i in range(n):
        messages = [{"type": "websocket.connect"}] + [
            {"type": "websocket.receive", "text": f"ws{i}-msg{j}"} for j in range(5)
        ] + [{"type": "websocket.disconnect"}]
        ws = WebSocket(scope, make_receive(list(messages)), dummy_send)
        ws_list.append(ws)
    # Connect all
    await asyncio.gather(*(ws.receive() for ws in ws_list))
    # Receive all messages
    for j in range(5):
        await asyncio.gather(*(ws.receive() for ws in ws_list))
    # Disconnect all
    await asyncio.gather(*(ws.receive() for ws in ws_list))

@pytest.mark.asyncio
async def test_receive_throughput_medium_load():
    """
    Throughput test: medium load of 20 websocket connections, each with 20 messages.
    """
    scope = {"type": "websocket"}
    n = 20
    ws_list = []
    for i in range(n):
        messages = [{"type": "websocket.connect"}] + [
            {"type": "websocket.receive", "text": f"ws{i}-msg{j}"} for j in range(20)
        ] + [{"type": "websocket.disconnect"}]
        ws = WebSocket(scope, make_receive(list(messages)), dummy_send)
        ws_list.append(ws)
    # Connect all
    await asyncio.gather(*(ws.receive() for ws in ws_list))
    # Receive all messages
    for j in range(20):
        await asyncio.gather(*(ws.receive() for ws in ws_list))
    # Disconnect all
    await asyncio.gather(*(ws.receive() for ws in ws_list))

@pytest.mark.asyncio

To edit these changes git checkout codeflash/optimize-WebSocket.receive-mhbh4b0p and push.