From 19a28735839e4512bfc5618da746ab2569a92f70 Mon Sep 17 00:00:00 2001
From: Wentao Wu <wuwentao@canaan-creative.com>
Date: Wed, 25 Sep 2024 18:59:39 +0800
Subject: [PATCH] chore: update socket close and loop break

---
 midealocal/device.py | 41 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

diff --git a/midealocal/device.py b/midealocal/device.py
index 38288151..23560850 100644
--- a/midealocal/device.py
+++ b/midealocal/device.py
@@ -294,6 +294,29 @@ def send_message_v2(self, data: bytes, query: bool = False) -> None:
             if query:
                 self._socket.settimeout(QUERY_TIMEOUT)
             self._socket.send(data)
+        except TimeoutError:
+            _LOGGER.debug(
+                "[%s] send_message_v2 timed out",
+                self._device_id,
+            )
+            # raise exception to main loop
+            raise
+        except ConnectionResetError as e:
+            _LOGGER.debug(
+                "[%s] send_message_v2 ConnectionResetError: %s",
+                self._device_id,
+                e,
+            )
+            # raise exception to main loop
+            raise
+        except OSError as e:
+            _LOGGER.debug(
+                "[%s] send_message_v2 OSError: %s",
+                self._device_id,
+                e,
+            )
+            # raise exception to main loop
+            raise
         except Exception as e:
             _LOGGER.exception(
                 "[%s] send_message_v2 Unexpected socket error",
@@ -350,12 +373,13 @@ def refresh_status(self, check_protocol: bool = False) -> None:
                         if len(msg) == 0:
                             raise OSError("Empty message received.")
                         result = self.parse_message(msg)
+                        # Prevent infinite loop
                         if result == MessageResult.SUCCESS:
                             break
-                        if result == MessageResult.PADDING:
+                        elif result == MessageResult.PADDING:  # noqa: RET508
                             continue
-                        # parse msg error
-                        error_count += 1
+                        else:
+                            raise ResponseException  # noqa: TRY301
                     # recovery SOCKET_TIMEOUT after recv msg
                     self._socket.settimeout(SOCKET_TIMEOUT)
                 # only catch TimoutError for check_protocol
@@ -380,6 +404,9 @@ def refresh_status(self, check_protocol: bool = False) -> None:
                     # refresh_status, raise timeout exception to main loop
                     else:
                         raise
+                except ResponseException:
+                    # parse msg error
+                    error_count += 1
             else:
                 error_count += 1
             # init check_protocol and all the query failed
@@ -557,6 +584,7 @@ def close_socket(self, init: bool = False) -> None:
         self._buffer = b""
         if self._socket:
             try:
+                self._socket.shutdown(socket.SHUT_RDWR)
                 self._socket.close()
                 _LOGGER.debug("[%s] Socket closed", self._device_id)
             except OSError as e:
@@ -620,8 +648,8 @@ def run(self) -> None:
                 if self.connect(init=True) is False:
                     self.close_socket(init=True)
                     connection_retries += 1
-                    # sleep time increase, maximum is 600 seconds
-                    sleep_time = min(5 * connection_retries, 600)
+                    # Sleep time with exponential backoff, maximum 600 seconds
+                    sleep_time = min(5 * (2 ** (connection_retries - 1)), 600)
                     _LOGGER.warning(
                         "[%s] Unable to connect, sleep %s seconds and retry",
                         self._device_id,
@@ -649,6 +677,9 @@ def run(self) -> None:
                     _LOGGER.debug("[%s] No Supported protocol", self._device_id)
                     # ignore and continue loop
                     continue
+                except ConnectionResetError:  # refresh_status -> build_send exception
+                    _LOGGER.debug("[%s] Connection reset by peer", self._device_id)
+                    reconnect = True
                 except OSError:  # refresh_status
                     _LOGGER.debug("[%s] OS error", self._device_id)
                     reconnect = True