Fix permanent disconnect issue (#273)

This commit is contained in:
Alex Yao 2022-10-24 07:11:16 -05:00 committed by GitHub
parent efeab0dcf3
commit b550b294c7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -109,11 +109,17 @@ class APIConnection:
self._process_task: Optional[asyncio.Task[None]] = None self._process_task: Optional[asyncio.Task[None]] = None
self._connect_lock: asyncio.Lock = asyncio.Lock()
self._cleanup_task: Optional[asyncio.Task[None]] = None
async def _cleanup(self) -> None: async def _cleanup(self) -> None:
"""Clean up all resources that have been allocated. """Clean up all resources that have been allocated.
Safe to call multiple times. Safe to call multiple times.
""" """
async def _do_cleanup() -> None:
async with self._connect_lock:
if self._frame_helper is not None: if self._frame_helper is not None:
await self._frame_helper.close() await self._frame_helper.close()
self._frame_helper = None self._frame_helper = None
@ -138,6 +144,9 @@ class APIConnection:
# themself, effectively ending execution after _cleanup which may be unexpected # themself, effectively ending execution after _cleanup which may be unexpected
self._ping_stop_event.set() self._ping_stop_event.set()
if not self._cleanup_task or not self._cleanup_task.done():
self._cleanup_task = asyncio.create_task(_do_cleanup())
async def _connect_resolve_host(self) -> hr.AddrInfo: async def _connect_resolve_host(self) -> hr.AddrInfo:
"""Step 1 in connect process: resolve the address.""" """Step 1 in connect process: resolve the address."""
try: try:
@ -305,6 +314,10 @@ class APIConnection:
if login: if login:
await self.login() await self.login()
# A connection lock must be created to avoid potential issues where
# connect has succeeded but not yet returned, followed by a disconnect.
# See esphome/aioesphomeapi#258 for more information
async with self._connect_lock:
try: try:
# Allow 2 minutes for connect; this is only as a last measure # Allow 2 minutes for connect; this is only as a last measure
# to protect from issues if some part of the connect process mistakenly # to protect from issues if some part of the connect process mistakenly