From 7ab8297e1356207c03f662c202b1356cf99339de Mon Sep 17 00:00:00 2001 From: Lurker00 Date: Fri, 18 Oct 2024 14:59:40 +0300 Subject: [PATCH] perf: close connection improvements (#369) * Eliminated unneeded self.heartbeater.cancel() * Logical order of actions * Direct call instead of await * Wait for async tasks to finish on close --- .../localtuya/core/pytuya/__init__.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/custom_components/localtuya/core/pytuya/__init__.py b/custom_components/localtuya/core/pytuya/__init__.py index 833cfc844..38a2aadec 100644 --- a/custom_components/localtuya/core/pytuya/__init__.py +++ b/custom_components/localtuya/core/pytuya/__init__.py @@ -965,6 +965,7 @@ async def keep_alive_loop(action): self.exception("Heartbeat failed (%s), disconnecting", ex) break + self.heartbeater = None if self.transport is not None: self.clean_up_session() @@ -1016,6 +1017,12 @@ async def close(self): self.debug("Closing connection") self.clean_up_session() + if self.heartbeater: + await self.heartbeater + + if self._sub_devs_query_task: + await self._sub_devs_query_task + def clean_up_session(self): """Clean up session.""" self.debug(f"Cleaning up session.") @@ -1027,12 +1034,12 @@ def clean_up_session(self): if self._sub_devs_query_task: self._sub_devs_query_task.cancel() - if self.dispatcher: - self.dispatcher.abort() - if self.is_connected: self.transport.close() + if self.dispatcher: + self.dispatcher.abort() + async def exchange_quick(self, payload, recv_retries): """Similar to exchange() but never retries sending and does not decode the response.""" if not self.is_connected: @@ -1048,8 +1055,8 @@ async def exchange_quick(self, payload, recv_retries): try: await self.transport_write(enc_payload) except Exception: # pylint: disable=broad-except - await self.close() - return None + return self.clean_up_session() + while recv_retries: try: seqno = MessageDispatcher.SESS_KEY_SEQNO