Skip to content

Commit

Permalink
Ensure we hold strong references to tasks
Browse files Browse the repository at this point in the history
  • Loading branch information
bdraco committed Feb 13, 2023
1 parent f99db35 commit 35276b6
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 3 deletions.
20 changes: 18 additions & 2 deletions aioesphomeapi/connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ def __init__(
) -> None:
self._params = params
self.on_stop: Optional[Callable[[], Coroutine[Any, Any, None]]] = on_stop
self._on_stop_task: Optional[asyncio.Task[None]] = None
self._socket: Optional[socket.socket] = None
self._frame_helper: Optional[APIFrameHelper] = None
self._api_version: Optional[APIVersion] = None
Expand Down Expand Up @@ -142,6 +143,10 @@ def _cleanup(self) -> None:
self._connect_task.cancel()
self._connect_task = None

if self._keep_alive_task is not None:
self._keep_alive_task.cancel()
self._keep_alive_task = None

if self._frame_helper is not None:
self._frame_helper.close()
self._frame_helper = None
Expand All @@ -151,8 +156,19 @@ def _cleanup(self) -> None:
self._socket = None

if self.on_stop and self._connect_complete:

def _remove_on_stop_task():
"""Remove the stop task from the reconnect loop.
We need to do this because the asyncio does not hold
a strong reference to the task, so it can be garbage
collected unexpectedly.
"""
self._on_stop_task = None

# Ensure on_stop is called only once
asyncio.create_task(self.on_stop())
self._on_stop_task = asyncio.create_task(self.on_stop())
self._on_stop_task.add_done_callback(_remove_on_stop_task)
self.on_stop = None

# Note: we don't explicitly cancel the ping/read task here
Expand Down Expand Up @@ -318,7 +334,7 @@ async def _keep_alive_loop() -> None:
self._report_fatal_error(err)
return

asyncio.create_task(_keep_alive_loop())
self._keep_alive_task = asyncio.create_task(_keep_alive_loop())

async def connect(self, *, login: bool) -> None:
if self._connection_state != ConnectionState.INITIALIZED:
Expand Down
12 changes: 11 additions & 1 deletion aioesphomeapi/reconnect_logic.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,17 @@ async def stop(self) -> None:
await self._stop_zc_listen()

def stop_callback(self) -> None:
asyncio.create_task(self.stop())
def _remove_stop_task() -> None:
"""Remove the stop task from the reconnect loop.
We need to do this because the asyncio does not hold
a strong reference to the task, so it can be garbage
collected unexpectedly.
"""
self._stop_task = None

self._stop_task = asyncio.create_task(self.stop())
self._stop_task.add_done_callback(_remove_stop_task)

async def _start_zc_listen(self) -> None:
"""Listen for mDNS records.
Expand Down

0 comments on commit 35276b6

Please sign in to comment.