Skip to content

Commit 2dae682

Browse files
committed
Fix async websocket reader starvation and dispatch races
Replace the racing-recv() polling loop with a single background reader task that is the only thing allowed to touch the socket. recv() calls now wait on a shared Condition and get notified per dispatched message, instead of each concurrently calling _async_recv() and starving each other. Errors are routed to the specific msg_id they belong to when attributable, and broadcast to all waiters only when the connection itself is dead. Bump nimax to 1.1.1 and enable id-aware WebSocket replay gating so test cassettes correctly model concurrent, out-of-order request/response correlation instead of relying on positional ordering.
1 parent 9cc48b9 commit 2dae682

2 files changed

Lines changed: 70 additions & 15 deletions

File tree

‎homeassistant_api/asyncwebsocket.py‎

Lines changed: 68 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
from __future__ import annotations
22

3+
import asyncio
34
import contextlib
45
import json
56
import logging
@@ -54,6 +55,7 @@
5455
class AsyncWebsocketClient(BaseWebsocketClient):
5556
_session: AsyncSession
5657
_ws: AsyncExtensionFromHTTP
58+
_reader_task: asyncio.Task[None]
5759

5860
def __init__(
5961
self,
@@ -65,6 +67,9 @@ def __init__(
6567
) -> None:
6668
super().__init__(api_url, token, max_size=max_size)
6769
self._session = session if session is not None else AsyncSession()
70+
self._new_message = asyncio.Condition()
71+
self._fatal_error: BaseException | None = None
72+
self._id_errors: dict[int, BaseException] = {}
6873

6974
async def __aenter__(self) -> Self:
7075
await self._session.__aenter__()
@@ -74,8 +79,11 @@ async def __aenter__(self) -> Self:
7479
msg = "Server did not upgrade to WebSocket"
7580
raise ReceivingError(msg)
7681
self._ws = resp.extension
82+
# Authenticate before starting the reader task: these messages have no "id"
83+
# and aren't dispatched through it, so nothing else may read the socket yet.
7784
okay = await self.authentication_phase()
7885
logger.info("Authenticated with Home Assistant (%s)", okay.ha_version)
86+
self._reader_task = asyncio.create_task(self._reader_loop())
7987
await self.supported_features_phase()
8088
return self
8189

@@ -85,10 +93,50 @@ async def __aexit__(
8593
exc_value: BaseException | None,
8694
traceback: TracebackType | None,
8795
) -> None:
96+
self._reader_task.cancel()
97+
with contextlib.suppress(asyncio.CancelledError):
98+
await self._reader_task
8899
await self._ws.close()
89100
del self._ws
90101
await self._session.__aexit__(exc_type, exc_value, traceback)
91102

103+
async def _reader_loop(self) -> None:
104+
"""
105+
Continuously read messages off the socket and dispatch them by id.
106+
107+
This is the only task allowed to touch the socket directly, so reads can
108+
never interleave or starve a waiting `recv()` call.
109+
"""
110+
while True:
111+
try:
112+
data = await self._async_recv()
113+
except Exception as e: # noqa: BLE001
114+
# No message id is available yet, so there's no single owner to
115+
# blame this on: the connection is dead for every pending recv().
116+
async with self._new_message:
117+
self._fatal_error = e
118+
self._new_message.notify_all()
119+
return
120+
await self._dispatch(data)
121+
122+
async def _dispatch(self, data: dict[str, Any]) -> None:
123+
"""Hand a received message to `handle_recv`, routing any failure to its owner."""
124+
try:
125+
self.handle_recv(data)
126+
except Exception as e: # noqa: BLE001
127+
data_id = data.get("id")
128+
async with self._new_message:
129+
if isinstance(data_id, int):
130+
# The message carried its own id, so only the caller waiting
131+
# on that id needs to see the failure.
132+
self._id_errors[data_id] = e
133+
else:
134+
self._fatal_error = e
135+
self._new_message.notify_all()
136+
return
137+
async with self._new_message:
138+
self._new_message.notify_all()
139+
92140
async def _async_send(self, data: dict[str, Any]) -> None:
93141
"""Send a message to the websocket server."""
94142
if data.get("type") != "auth":
@@ -147,20 +195,26 @@ async def recv(
147195
msg_id: int,
148196
) -> EventResponse | ResultResponse | PingResponse | None:
149197
"""Receive a response to a message from the websocket server."""
150-
while True:
151-
## have we received a message with the id we're looking for?
152-
if self._result_responses.get(msg_id) is not None:
153-
return self._result_responses.pop(msg_id)
154-
if self._event_responses.get(msg_id, []):
155-
return self._event_responses[msg_id].pop(0)
156-
if (
157-
self._ping_responses.get(msg_id) is not None
158-
and self._ping_responses[msg_id].end is not None
159-
):
160-
return self._ping_responses.pop(msg_id)
161-
162-
## if not, keep receiving messages until we do
163-
self.handle_recv(await self._async_recv())
198+
async with self._new_message:
199+
while True:
200+
## have we received a message with the id we're looking for?
201+
if self._result_responses.get(msg_id) is not None:
202+
return self._result_responses.pop(msg_id)
203+
if self._event_responses.get(msg_id, []):
204+
return self._event_responses[msg_id].pop(0)
205+
if (
206+
self._ping_responses.get(msg_id) is not None
207+
and self._ping_responses[msg_id].end is not None
208+
):
209+
return self._ping_responses.pop(msg_id)
210+
## did our own request fail, or did the connection die outright?
211+
if msg_id in self._id_errors:
212+
raise self._id_errors.pop(msg_id)
213+
if self._fatal_error is not None:
214+
raise self._fatal_error
215+
216+
## if not, wait for the reader task to dispatch another message
217+
await self._new_message.wait()
164218

165219
async def recv_result(self, msg_id: int) -> ResultResponse:
166220
"""Receive a ResultResponse, raising TypeError if the response is not a ResultResponse."""

‎pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,13 @@ dev = [
4343
"aiosqlite>=0.22",
4444
"prek>=0.3.8",
4545
"pre-commit>=4.5.1",
46-
"nimax",
46+
"nimax>=1.1.1",
4747
]
4848

4949
[tool.nimax]
5050
cassette_library_dir = "tests/cassettes"
5151
match_on = ["method", "uri"]
52+
ws_id_extractor = "id"
5253

5354
[tool.pytest.ini_options]
5455
asyncio_mode = "auto"

0 commit comments

Comments
 (0)