Split connection process to enable earlier zeroconf listening (#1205)

Author: bdraco

aioesphomeapi/client.py

@@ -206,7 +206,8 @@ async def connect(
         login: bool = False,
     ) -> None:
         """Connect to the device."""
-        await self.start_connection(on_stop)
+        await self.start_resolve_host(on_stop)
+        await self.start_connection()
         await self.finish_connection(login)
 
     def _on_stop(
@@ -219,11 +220,11 @@ def _on_stop(
         if on_stop:
             self._create_background_task(on_stop(expected_disconnect))
 
-    async def start_connection(
+    async def start_resolve_host(
         self,
         on_stop: Callable[[bool], Coroutine[Any, Any, None]] | None = None,
     ) -> None:
-        """Start connecting to the device."""
+        """Start resolving the host."""
         if self._connection is not None:
             raise APIConnectionError(f"Already connected to {self.log_name}!")
         self._connection = APIConnection(
@@ -232,6 +233,12 @@ async def start_connection(
             self._debug_enabled,
             self.log_name,
         )
+        await self._execute_connection_coro(self._connection.start_resolve_host())
+
+    async def start_connection(self) -> None:
+        """Start connecting to the device."""
+        if TYPE_CHECKING:
+            assert self._connection is not None
         await self._execute_connection_coro(self._connection.start_connection())
         # If we connected, we should set the log name now
         if self._connection.connected_address:

aioesphomeapi/connection.pxd

@@ -53,6 +53,7 @@ cdef object SocketClosedAPIError
 cdef object astuple
 
 cdef object CONNECTION_STATE_INITIALIZED
+cdef object CONNECTION_STATE_HOST_RESOLVED
 cdef object CONNECTION_STATE_SOCKET_OPENED
 cdef object CONNECTION_STATE_HANDSHAKE_COMPLETE
 cdef object CONNECTION_STATE_CONNECTED
@@ -104,6 +105,7 @@ cdef class APIConnection:
     cdef object _pong_timer
     cdef float _keep_alive_interval
     cdef float _keep_alive_timeout
+    cdef object _resolve_host_future
     cdef object _start_connect_future
     cdef object _finish_connect_future
     cdef public Exception _fatal_exception
@@ -115,6 +117,7 @@ cdef class APIConnection:
     cdef bint _debug_enabled
     cdef public str received_name
     cdef public str connected_address
+    cdef list _addrs_info
 
     cpdef void send_message(self, object msg) except *
 

aioesphomeapi/connection.py

@@ -128,16 +128,19 @@ class ConnectionParams:
 class ConnectionState(enum.Enum):
     # The connection is initialized, but connect() wasn't called yet
     INITIALIZED = 0
+    # The host has been resolved, but the socket hasn't been opened yet
+    HOST_RESOLVED = 1
     # The socket has been opened, but the handshake and login haven't been completed
-    SOCKET_OPENED = 1
+    SOCKET_OPENED = 2
     # The handshake has been completed, messages can be exchanged
-    HANDSHAKE_COMPLETE = 2
+    HANDSHAKE_COMPLETE = 3
     # The connection has been established, authenticated data can be exchanged
-    CONNECTED = 3
-    CLOSED = 4
+    CONNECTED = 4
+    CLOSED = 5
 
 
 CONNECTION_STATE_INITIALIZED = ConnectionState.INITIALIZED
+CONNECTION_STATE_HOST_RESOLVED = ConnectionState.HOST_RESOLVED
 CONNECTION_STATE_SOCKET_OPENED = ConnectionState.SOCKET_OPENED
 CONNECTION_STATE_HANDSHAKE_COMPLETE = ConnectionState.HANDSHAKE_COMPLETE
 CONNECTION_STATE_CONNECTED = ConnectionState.CONNECTED
@@ -192,6 +195,7 @@ class APIConnection:
     """
 
     __slots__ = (
+        "_addrs_info",
         "_debug_enabled",
         "_expected_disconnect",
         "_fatal_exception",
@@ -206,6 +210,7 @@ class APIConnection:
         "_ping_timer",
         "_pong_timer",
         "_read_exception_futures",
+        "_resolve_host_future",
         "_send_pending_ping",
         "_socket",
         "_start_connect_future",
@@ -247,6 +252,7 @@ def __init__(
         self._keep_alive_interval = keepalive
         self._keep_alive_timeout = keepalive * KEEP_ALIVE_TIMEOUT_RATIO
 
+        self._resolve_host_future: asyncio.Future[None] | None = None
         self._start_connect_future: asyncio.Future[None] | None = None
         self._finish_connect_future: asyncio.Future[None] | None = None
         self._fatal_exception: Exception | None = None
@@ -258,6 +264,7 @@ def __init__(
         self._debug_enabled = debug_enabled
         self.received_name: str = ""
         self.connected_address: str | None = None
+        self._addrs_info: list[hr.AddrInfo] = []
 
     def set_log_name(self, name: str) -> None:
         """Set the friendly log name for this connection."""
@@ -286,6 +293,7 @@ def _cleanup(self) -> None:
                 fut.set_exception(new_exc)
         self._read_exception_futures.clear()
 
+        self._set_resolve_host_future()
         self._set_start_connect_future()
         self._set_finish_connect_future()
 
@@ -572,32 +580,61 @@ def _async_pong_not_received(self) -> None:
             )
         )
 
-    async def _do_connect(self) -> None:
-        """Do the actual connect process."""
-        addrs_info = await hr.async_resolve_host(
-            self._params.addresses,
-            self._params.port,
-            self._params.zeroconf_manager,
-        )
-        await self._connect_socket_connect(addrs_info)
+    async def start_resolve_host(self) -> None:
+        """Start the host resolution process.
+
+        This part of the process resolves the hostnames to IP addresses
+        and prepares the connection for the next step.
+        """
+        if self.connection_state is not CONNECTION_STATE_INITIALIZED:
+            raise RuntimeError(
+                "Connection can only be used once, connection is not in init state"
+            )
+
+        self._resolve_host_future = self._loop.create_future()
+        try:
+            async with interrupt(
+                self._resolve_host_future, ConnectionInterruptedError, None
+            ):
+                self._addrs_info = await hr.async_resolve_host(
+                    self._params.addresses,
+                    self._params.port,
+                    self._params.zeroconf_manager,
+                )
+        except (Exception, CancelledError) as ex:
+            # If the task was cancelled, we need to clean up the connection
+            # and raise the CancelledError as APIConnectionError
+            self._cleanup()
+            raise self._wrap_fatal_connection_exception("resolving", ex)
+        finally:
+            self._set_resolve_host_future()
+        self._set_connection_state(CONNECTION_STATE_HOST_RESOLVED)
+
+    def _set_resolve_host_future(self) -> None:
+        if (
+            self._resolve_host_future is not None
+            and not self._resolve_host_future.done()
+        ):
+            self._resolve_host_future.set_result(None)
+            self._resolve_host_future = None
 
     async def start_connection(self) -> None:
         """Start the connection process.
 
         This part of the process establishes the socket connection but
         does not initialize the frame helper or send the hello message.
         """
-        if self.connection_state is not CONNECTION_STATE_INITIALIZED:
+        if self.connection_state is not CONNECTION_STATE_HOST_RESOLVED:
             raise RuntimeError(
-                "Connection can only be used once, connection is not in init state"
+                "Connection must be in HOST_RESOLVED state to start connection"
             )
 
         self._start_connect_future = self._loop.create_future()
         try:
             async with interrupt(
                 self._start_connect_future, ConnectionInterruptedError, None
             ):
-                await self._do_connect()
+                await self._connect_socket_connect(self._addrs_info)
         except (Exception, CancelledError) as ex:
             # If the task was cancelled, we need to clean up the connection
             # and raise the CancelledError as APIConnectionError

aioesphomeapi/reconnect_logic.py

@@ -34,15 +34,17 @@
 
 
 class ReconnectLogicState(Enum):
-    CONNECTING = 0
-    HANDSHAKING = 1
-    READY = 2
-    DISCONNECTED = 3
+    RESOLVING = 0
+    CONNECTING = 1
+    HANDSHAKING = 2
+    READY = 3
+    DISCONNECTED = 4
 
 
 NOT_YET_CONNECTED_STATES = {
     ReconnectLogicState.DISCONNECTED,
     ReconnectLogicState.CONNECTING,
+    ReconnectLogicState.RESOLVING,
 }
 
 
@@ -187,10 +189,21 @@ def _async_log_connection_error(self, err: Exception) -> None:
 
     async def _try_connect(self) -> bool:
         """Try connecting to the API client."""
+        self._async_set_connection_state_while_locked(ReconnectLogicState.RESOLVING)
+        start_resolve_time = time.perf_counter()
+        try:
+            await self._cli.start_resolve_host(on_stop=self._on_disconnect)
+        except Exception as err:  # pylint: disable=broad-except
+            await self._handle_connection_failure(err)
+            return False
         self._async_set_connection_state_while_locked(ReconnectLogicState.CONNECTING)
         start_connect_time = time.perf_counter()
+        resolve_time = start_connect_time - start_resolve_time
+        _LOGGER.info(
+            "Successfully resolved %s in %0.3fs", self._cli.log_name, resolve_time
+        )
         try:
-            await self._cli.start_connection(on_stop=self._on_disconnect)
+            await self._cli.start_connection()
         except Exception as err:  # pylint: disable=broad-except
             await self._handle_connection_failure(err)
             return False
@@ -248,7 +261,13 @@ def _call_connect_once(self) -> None:
         if self._connect_task and not self._connect_task.done():
             if self._connection_state != ReconnectLogicState.CONNECTING:
                 # Connection state is far enough along that we should
-                # not restart the connect task
+                # not restart the connect task.
+                #
+                # Zeroconf triggering scenarios:
+                # - RESOLVING state: Don't cancel, the resolve task will complete immediately
+                #   since it's waiting for the same records zeroconf is delivering
+                # - CONNECTING state: Cancel and restart to use potentially updated connection info
+                # - HANDSHAKING state or later: Don't cancel, too far along in the process
                 _LOGGER.debug(
                     "%s: Not cancelling existing connect task as its already %s!",
                     self._cli.log_name,
@@ -300,6 +319,7 @@ async def _connect_once_or_reschedule(self) -> None:
                 or self._is_stopped
             ):
                 return
+            self._start_zc_listen()
             if await self._try_connect():
                 return
             tries = min(self._tries, 10)  # prevent OverflowError
@@ -309,9 +329,6 @@ async def _connect_once_or_reschedule(self) -> None:
                     "Trying to connect to %s in the background", self._cli.log_name
                 )
             _LOGGER.debug("Retrying %s in %.2f seconds", self._cli.log_name, wait_time)
-            if wait_time:
-                # If we are waiting, start listening for mDNS records
-                self._start_zc_listen()
             self._schedule_connect(wait_time)
 
     def _remove_stop_task(self, _fut: asyncio.Future[None]) -> None:

tests/common.py

@@ -146,6 +146,7 @@ def async_fire_time_changed(
 
 async def connect(conn: APIConnection, login: bool = True):
     """Wrapper for connection logic to do both parts."""
+    await conn.start_resolve_host()
     await conn.start_connection()
     await conn.finish_connection(login=login)
 
@@ -156,7 +157,8 @@ async def connect_client(
     on_stop: Callable[[bool], Awaitable[None]] | None = None,
 ) -> None:
     """Wrapper for connection logic to do both parts."""
-    await client.start_connection(on_stop=on_stop)
+    await client.start_resolve_host(on_stop=on_stop)
+    await client.start_connection()
     await client.finish_connection(login=login)
 
 

tests/test_client.py

@@ -196,16 +196,18 @@ async def test_expected_name(auth_client: APIClient) -> None:
 
 
 async def test_connect_backwards_compat() -> None:
-    """Verify connect is a thin wrapper around start_connection and finish_connection."""
+    """Verify connect is a thin wrapper around start_resolve_host, start_connection and finish_connection."""
 
     cli = PatchableAPIClient("host", 1234, None)
     with (
+        patch.object(cli, "start_resolve_host") as mock_start_resolve_host,
         patch.object(cli, "start_connection") as mock_start_connection,
         patch.object(cli, "finish_connection") as mock_finish_connection,
     ):
         await cli.connect()
 
-    assert mock_start_connection.mock_calls == [call(None)]
+    assert mock_start_resolve_host.mock_calls == [call(None)]
+    assert mock_start_connection.mock_calls == [call()]
     assert mock_finish_connection.mock_calls == [call(False)]
 
 
@@ -216,6 +218,7 @@ async def test_finish_connection_wraps_exceptions_as_unhandled_api_error(
 
     cli = APIClient("127.0.0.1", 1234, None)
     with patch("aioesphomeapi.client.APIConnection", PatchableAPIConnection):
+        await cli.start_resolve_host()
         await cli.start_connection()
 
     with (
@@ -244,6 +247,7 @@ async def _start_connection_with_delay(*args, **kwargs):
         "aioesphomeapi.connection.aiohappyeyeballs.start_connection",
         _start_connection_with_delay,
     ):
+        await cli.start_resolve_host()
         start_task = asyncio.create_task(cli.start_connection())
         await asyncio.sleep(0)
         assert cli._connection is not None
@@ -265,6 +269,7 @@ async def _start_connection_without_delay(*args, **kwargs):
             _start_connection_without_delay,
         ),
     ):
+        await cli.start_resolve_host()
         await cli.start_connection()
         await asyncio.sleep(0)
 
@@ -303,7 +308,7 @@ async def test_request_while_handshaking() -> None:
 async def test_connect_while_already_connected(auth_client: APIClient) -> None:
     """Test connecting while already connected raises."""
     with pytest.raises(APIConnectionError):
-        await auth_client.start_connection()
+        await auth_client.start_resolve_host()
 
 
 @pytest.mark.parametrize(
@@ -1066,7 +1071,11 @@ async def test_noise_psk_handles_subclassed_string():
     )
     assert rl._connection_state is ReconnectLogicState.DISCONNECTED
 
-    with patch.object(cli, "start_connection"), patch.object(cli, "finish_connection"):
+    with (
+        patch.object(cli, "start_resolve_host"),
+        patch.object(cli, "start_connection"),
+        patch.object(cli, "finish_connection"),
+    ):
         await rl.start()
         for _ in range(3):
             await asyncio.sleep(0)

tests/test_connection.py

@@ -1037,7 +1037,7 @@ async def test_connection_cannot_be_reused(
     send_plaintext_connect_response(protocol, False)
     await connect_task
     with pytest.raises(RuntimeError):
-        await conn.start_connection()
+        await conn.start_resolve_host()
 
 
 async def test_attempting_to_finish_unstarted_connection(
@@ -1048,6 +1048,17 @@ async def test_attempting_to_finish_unstarted_connection(
         await conn.finish_connection(login=False)
 
 
+async def test_start_connection_wrong_state(
+    conn: APIConnection,
+) -> None:
+    """Test that we raise when trying to start connection in wrong state."""
+    with pytest.raises(
+        RuntimeError,
+        match="Connection must be in HOST_RESOLVED state to start connection",
+    ):
+        await conn.start_connection()
+
+
 async def test_internal_message_received_immediately_after_connection(
     conn: APIConnection,
     resolve_host: AsyncMock,