Cleanly handle connection loss

Author: puddly

bellows/ash.py

@@ -367,7 +367,7 @@ def connection_made(self, transport):
         self._transport = transport
         self._ezsp_protocol.connection_made(self)
 
-    def connection_lost(self, exc):
+    def connection_lost(self, exc: Exception | None) -> None:
         self._ezsp_protocol.connection_lost(exc)
 
     def eof_received(self):

bellows/ezsp/__init__.py

@@ -58,13 +58,14 @@ class EZSP:
         v13.EZSPv13.VERSION: v13.EZSPv13,
     }
 
-    def __init__(self, device_config: dict):
+    def __init__(self, device_config: dict, application: Any | None = None):
         self._config = device_config
         self._callbacks = {}
         self._ezsp_event = asyncio.Event()
         self._ezsp_version = v4.EZSPv4.VERSION
         self._gw = None
         self._protocol = None
+        self._application = application
         self._send_sem = PriorityDynamicBoundedSemaphore(value=MAX_COMMAND_CONCURRENCY)
 
         self._stack_status_listeners: collections.defaultdict[
@@ -126,24 +127,18 @@ async def startup_reset(self) -> None:
 
         await self.version()
 
-    @classmethod
-    async def initialize(cls, zigpy_config: dict) -> EZSP:
-        """Return initialized EZSP instance."""
-        ezsp = cls(zigpy_config[conf.CONF_DEVICE])
-        await ezsp.connect(use_thread=zigpy_config[conf.CONF_USE_THREAD])
-
-        try:
-            await ezsp.startup_reset()
-        except Exception:
-            ezsp.close()
-            raise
-
-        return ezsp
-
     async def connect(self, *, use_thread: bool = True) -> None:
         assert self._gw is None
         self._gw = await bellows.uart.connect(self._config, self, use_thread=use_thread)
         self._protocol = v4.EZSPv4(self.handle_callback, self._gw)
+        await self.startup_reset()
+
+    async def disconnect(self) -> None:
+        if self._gw is not None:
+            await self._gw.disconnect()
+            self._gw = None
+        elif self._application is not None:
+            self._application.connection_lost(None)
 
     async def reset(self):
         LOGGER.debug("Resetting EZSP")
@@ -282,18 +277,13 @@ def connection_lost(self, exc):
             self._config[conf.CONF_DEVICE_PATH],
             exc,
         )
-        self.enter_failed_state(f"Serial connection loss: {exc!r}")
+        if self._application is not None:
+            self._application.connection_lost(exc)
 
     def enter_failed_state(self, error):
         """UART received error frame."""
-        if len(self._callbacks) > 1:
-            LOGGER.error("NCP entered failed state. Requesting APP controller restart")
-            self.close()
-            self.handle_callback("_reset_controller_application", (error,))
-        else:
-            LOGGER.info(
-                "NCP entered failed state. No application handler registered, ignoring..."
-            )
+        if self._application is not None:
+            self._application.connection_lost(error)
 
     def __getattr__(self, name: str) -> Callable:
         if name not in self._protocol.COMMANDS:

bellows/uart.py

@@ -18,41 +18,15 @@
 RESET_TIMEOUT = 5
 
 
-class Gateway(asyncio.Protocol):
-    FLAG = b"\x7E"  # Marks end of frame
-    ESCAPE = b"\x7D"
-    XON = b"\x11"  # Resume transmission
-    XOFF = b"\x13"  # Stop transmission
-    SUBSTITUTE = b"\x18"
-    CANCEL = b"\x1A"  # Terminates a frame in progress
-    STUFF = 0x20
-    RANDOMIZE_START = 0x42
-    RANDOMIZE_SEQ = 0xB8
-
-    RESERVED = FLAG + ESCAPE + XON + XOFF + SUBSTITUTE + CANCEL
-
-    class Terminator:
-        pass
-
-    def __init__(self, application, connected_future=None, connection_done_future=None):
+class Gateway(zigpy.serial.SerialProtocol):
+    def __init__(self, application, connection_done_future=None):
+        super().__init__()
         self._application = application
 
         self._reset_future = None
         self._startup_reset_future = None
-        self._connected_future = connected_future
         self._connection_done_future = connection_done_future
 
-        self._transport = None
-
-    def close(self):
-        self._transport.close()
-
-    def connection_made(self, transport):
-        """Callback when the uart is connected"""
-        self._transport = transport
-        if self._connected_future is not None:
-            self._connected_future.set_result(True)
-
     async def send_data(self, data: bytes) -> None:
         await self._transport.send_data(data)
 
@@ -92,12 +66,9 @@ def _reset_cleanup(self, future):
         """Delete reset future."""
         self._reset_future = None
 
-    def eof_received(self):
-        """Server gracefully closed its side of the connection."""
-        self.connection_lost(ConnectionResetError("Remote server closed connection"))
-
     def connection_lost(self, exc):
         """Port was closed unexpectedly."""
+        super().connection_lost(exc)
 
         LOGGER.debug("Connection lost: %r", exc)
         reason = exc or ConnectionResetError("Remote server closed connection")
@@ -117,11 +88,6 @@ def connection_lost(self, exc):
             self._reset_future.set_exception(reason)
             self._reset_future = None
 
-        if exc is None:
-            LOGGER.debug("Closed serial connection")
-            return
-
-        LOGGER.error("Lost serial connection: %r", exc)
         self._application.connection_lost(exc)
 
     async def reset(self):
@@ -144,10 +110,9 @@ async def reset(self):
 async def _connect(config, application):
     loop = asyncio.get_event_loop()
 
-    connection_future = loop.create_future()
     connection_done_future = loop.create_future()
 
-    gateway = Gateway(application, connection_future, connection_done_future)
+    gateway = Gateway(application, connection_done_future)
     protocol = AshProtocol(gateway)
 
     if config[zigpy.config.CONF_DEVICE_FLOW_CONTROL] is None:
@@ -164,7 +129,7 @@ async def _connect(config, application):
         rtscts=rtscts,
     )
 
-    await connection_future
+    await gateway.wait_until_connected()
 
     thread_safe_protocol = ThreadsafeProxy(gateway, loop)
     return thread_safe_protocol, connection_done_future

bellows/zigbee/application.py

@@ -142,22 +142,22 @@ async def _get_board_info(self) -> tuple[str, str, str] | tuple[None, None, None
         return None, None, None
 
     async def connect(self) -> None:
-        ezsp = bellows.ezsp.EZSP(self.config[zigpy.config.CONF_DEVICE])
-        await ezsp.connect(use_thread=self.config[CONF_USE_THREAD])
+        self._ezsp = bellows.ezsp.EZSP(self.config[zigpy.config.CONF_DEVICE], self)
 
         try:
-            await ezsp.startup_reset()
+            await self._ezsp.connect(use_thread=self.config[CONF_USE_THREAD])
+            await self._ezsp.startup_reset()
 
             # Writing config is required here because network info can't be loaded
-            await ezsp.write_config(self.config[CONF_EZSP_CONFIG])
-        except Exception:
-            ezsp.close()
-            raise
-
-        self._ezsp = ezsp
+            await self._ezsp.write_config(self.config[CONF_EZSP_CONFIG])
 
-        self._created_device_endpoints.clear()
-        await self.register_endpoints()
+            self._created_device_endpoints.clear()
+            await self.register_endpoints()
+        except Exception as exc:
+            await self._ezsp.disconnect()
+            self._ezsp = None
+            self.connection_lost(exc)
+            raise
 
     async def _ensure_network_running(self) -> bool:
         """Ensures the network is currently running and returns whether or not the network
@@ -603,8 +603,9 @@ async def disconnect(self):
         # TODO: how do you shut down the stack?
         self.controller_event.clear()
         if self._ezsp is not None:
-            self._ezsp.close()
+            await self._ezsp.disconnect()
             self._ezsp = None
+            self.connection_lost(None)
 
     async def force_remove(self, dev):
         # This should probably be delivered to the parent device instead
@@ -623,8 +624,6 @@ def ezsp_callback_handler(self, frame_name, args):
             self.handle_route_record(*args)
         elif frame_name == "incomingRouteErrorHandler":
             self.handle_route_error(*args)
-        elif frame_name == "_reset_controller_application":
-            self.connection_lost(args[0])
         elif frame_name == "idConflictHandler":
             self._handle_id_conflict(*args)