From ca0f1b76cfa16aa5959d7b38df769fcfbf64ca07 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 9 Jul 2023 17:02:39 -1000
Subject: [PATCH 1/2] Avoid sending pings when we have traffic

This is the same optimization that the ESPHome side has

https://github.com/esphome/esphome/blob/ddde1ee31ef96476c7a7d1c8b8fe306aa4d2e503/esphome/components/api/api_connection.cpp#L119
---
 aioesphomeapi/connection.py | 34 ++++++++++++++++++++--------------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py
index 365de5a3..f4b5c637 100644
--- a/aioesphomeapi/connection.py
+++ b/aioesphomeapi/connection.py
@@ -135,6 +135,8 @@ class APIConnection:
         "_connect_task",
         "_fatal_exception",
         "_expected_disconnect",
+        "_loop",
+        "_send_pending_ping",
     )
 
     def __init__(
@@ -172,6 +174,8 @@ def __init__(
         self._connect_task: Optional[asyncio.Task[None]] = None
         self._fatal_exception: Optional[Exception] = None
         self._expected_disconnect = False
+        self._send_pending_ping = False
+        self._loop = asyncio.get_running_loop()
 
     @property
     def connection_state(self) -> ConnectionState:
@@ -271,7 +275,7 @@ async def _connect_socket_connect(self, addr: hr.AddrInfo) -> None:
         sockaddr = astuple(addr.sockaddr)
 
         try:
-            coro = asyncio.get_event_loop().sock_connect(self._socket, sockaddr)
+            coro = self._loop.sock_connect(self._socket, sockaddr)
             async with async_timeout.timeout(TCP_CONNECT_TIMEOUT):
                 await coro
         except OSError as err:
@@ -290,7 +294,7 @@ async def _connect_socket_connect(self, addr: hr.AddrInfo) -> None:
     async def _connect_init_frame_helper(self) -> None:
         """Step 3 in connect process: initialize the frame helper and init read loop."""
         fh: Union[APIPlaintextFrameHelper, APINoiseFrameHelper]
-        loop = asyncio.get_event_loop()
+        loop = self._loop
 
         if self._params.noise_psk is None:
             _, fh = await loop.create_connection(
@@ -357,13 +361,10 @@ async def _connect_hello(self) -> None:
                 f"Server sent a different name '{resp.name}'", resp.name
             )
 
-    async def _connect_start_ping(self) -> None:
-        """Step 5 in connect process: start the ping loop."""
-        self._async_schedule_keep_alive(asyncio.get_running_loop())
-
-    def _async_schedule_keep_alive(self, loop: asyncio.AbstractEventLoop) -> None:
+    def _async_schedule_keep_alive(self) -> None:
         """Start the keep alive task."""
-        self._ping_timer = loop.call_later(
+        self._send_pending_ping = True
+        self._ping_timer = self._loop.call_later(
             self._keep_alive_interval, self._async_send_keep_alive
         )
 
@@ -372,14 +373,14 @@ def _async_send_keep_alive(self) -> None:
         if not self._is_socket_open:
             return
 
-        loop = asyncio.get_running_loop()
-        self.send_message(PING_REQUEST_MESSAGE)
+        if self._send_pending_ping:
+            self.send_message(PING_REQUEST_MESSAGE)
 
         if self._pong_timer is None:
             # Do not reset the timer if it's already set
             # since the only thing we want to reset the timer
             # is if we receive a pong.
-            self._pong_timer = loop.call_later(
+            self._pong_timer = self._loop.call_later(
                 self._keep_alive_timeout, self._async_pong_not_received
             )
         else:
@@ -399,7 +400,7 @@ def _async_send_keep_alive(self) -> None:
                 self._keep_alive_interval,
             )
 
-        self._async_schedule_keep_alive(loop)
+        self._async_schedule_keep_alive()
 
     def _async_cancel_pong_timer(self) -> None:
         """Cancel the pong timer."""
@@ -434,7 +435,7 @@ async def _do_connect() -> None:
             await self._connect_socket_connect(addr)
             await self._connect_init_frame_helper()
             await self._connect_hello()
-            await self._connect_start_ping()
+            self._async_schedule_keep_alive()
             if login:
                 await self.login(check_connected=False)
 
@@ -599,7 +600,7 @@ async def send_message_await_response_complex(
 
         :raises TimeoutAPIError: if a timeout occured
         """
-        fut = asyncio.get_event_loop().create_future()
+        fut = self._loop.create_future()
         responses = []
 
         def on_message(resp: message.Message) -> None:
@@ -725,6 +726,11 @@ def _process_packet(self, msg_type_proto: int, data: bytes) -> None:
             # as we know the connection is still alive
             self._async_cancel_pong_timer()
 
+        if self._send_pending_ping:
+            # Any valid message from the remove cancels the pending ping
+            # since we know the connection is still alive
+            self._send_pending_ping = False
+
         for handler in self._message_handlers.get(msg_type, [])[:]:
             handler(msg)
 

From 7f9ff19775f0dcbb4f8a6bb7d8d310e35360fb00 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Sun, 9 Jul 2023 17:06:17 -1000
Subject: [PATCH 2/2] legacy python

---
 aioesphomeapi/connection.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aioesphomeapi/connection.py b/aioesphomeapi/connection.py
index f4b5c637..720dafa9 100644
--- a/aioesphomeapi/connection.py
+++ b/aioesphomeapi/connection.py
@@ -175,7 +175,7 @@ def __init__(
         self._fatal_exception: Optional[Exception] = None
         self._expected_disconnect = False
         self._send_pending_ping = False
-        self._loop = asyncio.get_running_loop()
+        self._loop = asyncio.get_event_loop()
 
     @property
     def connection_state(self) -> ConnectionState: