Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

first step to fix dvd-dev/hilo#486 by making room to a second websock… #505

Draft
wants to merge 4 commits into
base: main
Choose a base branch
from
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 54 additions & 28 deletions custom_components/hilo/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,9 +232,12 @@ def __init__(self, hass: HomeAssistant, entry: ConfigEntry, api: API) -> None:
self.find_meter(self._hass)
self.entry = entry
self.devices: Devices = Devices(api)
self._websocket_reconnect_task: asyncio.Task | None = None
self._update_task: asyncio.Task | None = None
self.invocations = {0: self.subscribe_to_location}
self._websocket_reconnect_tasks: list[asyncio.Task | None] = [None, None]
self._update_task: list[asyncio.Task | None] = [None, None]
Comment on lines +235 to +236
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switching from a single asyncio.Task to a list for both _websocket_reconnect_tasks and _update_task introduces complexity, which could require additional management logic to handle these lists properly and ensure synchronization. Consider commenting the rationale behind using lists here or ensuring robust handling to avoid potential issues.

self.invocations = {
0: self.subscribe_to_location,
1: self.subscribe_to_challenge,
}
self.hq_plan_name = entry.options.get(CONF_HQ_PLAN_NAME, DEFAULT_HQ_PLAN_NAME)
self.appreciation = entry.options.get(
CONF_APPRECIATION_PHASE, DEFAULT_APPRECIATION_PHASE
Expand Down Expand Up @@ -344,13 +347,29 @@ async def on_websocket_event(self, event: WebsocketEvent) -> None:
async def subscribe_to_location(self, inv_id: int) -> None:
"""Sends the json payload to receive updates from the location."""
LOG.debug(f"Subscribing to location {self.devices.location_id}")
await self._api.websocket.async_invoke(
await self._api.websocket_devices.async_invoke(
[self.devices.location_id], "SubscribeToLocation", inv_id
)

@callback
async def subscribe_to_challenge(self, inv_id: int, event_id: int = 0) -> None:
"""Sends the json payload to receive updates from the challenge."""
LOG.debug(
f"Subscribing to challenge {event_id} at location {self.devices.location_id}"
)
await self._api.websocket_challenges.async_invoke(
[event_id, self.devices.location_id], "SubscribeToChallenge", inv_id
)

@callback
async def request_status_update(self) -> None:
await self._api.websocket.send_status()
await self._api.websocket_devices.send_status()
for inv_id, inv_cb in self.invocations.items():
await inv_cb(inv_id)

@callback
async def request_status_update_challenge(self) -> None:
await self._api.websocket_challenges.send_status()
for inv_id, inv_cb in self.invocations.items():
await inv_cb(inv_id)

Expand Down Expand Up @@ -424,20 +443,27 @@ async def async_init(self, scan_interval: int) -> None:
self._hass, self.entry, self.unknown_tracker_device
)

self._api.websocket.add_connect_callback(self.request_status_update)
self._api.websocket.add_event_callback(self.on_websocket_event)
self._websocket_reconnect_task = asyncio.create_task(
self.start_websocket_loop()
self._api.websocket_devices.add_connect_callback(self.request_status_update)
self._api.websocket_devices.add_event_callback(self.on_websocket_event)
self._api.websocket_challenges.add_connect_callback(
self.request_status_update_challenge
)
self._api.websocket_challenges.add_event_callback(self.on_websocket_event)
self._websocket_reconnect_tasks[0] = asyncio.create_task(
self.start_websocket_loop(self._api.websocket_devices, 0)
)
# asyncio.create_task(self._api.websocket.async_connect())
self._websocket_reconnect_tasks[1] = asyncio.create_task(
self.start_websocket_loop(self._api.websocket_challenges, 1)
)
# asyncio.create_task(self._api.websocket_devices.async_connect())

async def websocket_disconnect_listener(_: Event) -> None:
"""Define an event handler to disconnect from the websocket."""
if TYPE_CHECKING:
assert self._api.websocket
assert self._api.websocket_devices

if self._api.websocket.connected:
await self._api.websocket.async_disconnect()
if self._api.websocket_devices.connected:
await self._api.websocket_devices.async_disconnect()

self.entry.async_on_unload(
self._hass.bus.async_listen_once(
Expand All @@ -452,37 +478,37 @@ async def websocket_disconnect_listener(_: Event) -> None:
update_method=self.async_update,
)

async def start_websocket_loop(self) -> None:
async def start_websocket_loop(self, websocket, id) -> None:
"""Start a websocket reconnection loop."""
if TYPE_CHECKING:
assert self._api.websocket
assert websocket

should_reconnect = True

try:
await self._api.websocket.async_connect()
await self._api.websocket.async_listen()
await websocket.async_connect()
await websocket.async_listen()
except asyncio.CancelledError:
LOG.debug("Request to cancel websocket loop received")
raise
except WebsocketError as err:
LOG.error(f"Failed to connect to websocket: {err}", exc_info=err)
await self.cancel_websocket_loop()
await self.cancel_websocket_loop(websocket, id)
except InvalidCredentialsError:
LOG.warning("Invalid credentials? Refreshing websocket infos")
await self.cancel_websocket_loop()
await self.cancel_websocket_loop(websocket, id)
await self._api.refresh_ws_token()
except Exception as err: # pylint: disable=broad-except
LOG.error(
f"Unknown exception while connecting to websocket: {err}", exc_info=err
)
await self.cancel_websocket_loop()
await self.cancel_websocket_loop(websocket, id)

if should_reconnect:
LOG.info("Disconnected from websocket; reconnecting in 5 seconds.")
await asyncio.sleep(5)
self._websocket_reconnect_task = self._hass.async_create_task(
self.start_websocket_loop()
self._websocket_reconnect_tasks[id] = self._hass.async_create_task(
self.start_websocket_loop(websocket, id)
)

async def cancel_task(self, task) -> None:
Expand All @@ -496,15 +522,15 @@ async def cancel_task(self, task) -> None:
task = None
return task

async def cancel_websocket_loop(self) -> None:
async def cancel_websocket_loop(self, websocket, id) -> None:
"""Stop any existing websocket reconnection loop."""
self._websocket_reconnect_task = await self.cancel_task(
self._websocket_reconnect_task
self._websocket_reconnect_tasks[id] = await self.cancel_task(
self._websocket_reconnect_tasks[id]
)
self._update_task = await self.cancel_task(self._update_task)
self._update_task[id] = await self.cancel_task(self._update_task[id])
if TYPE_CHECKING:
assert self._api.websocket
await self._api.websocket.async_disconnect()
assert websocket
await websocket.async_disconnect()

async def async_update(self) -> None:
"""Updates tarif periodically."""
Expand Down
Loading