Add disconnect() to stop the reader thread and free its resources

pylutron/__init__.py

@@ -104,6 +104,7 @@ def __init__(self, host: str, user: str, password: str, recv_callback: Callable[
     self._connection_factory = connection_factory
     self._done = False
     self._loop = asyncio.new_event_loop()
+    self._main_task: Optional[asyncio.Task[None]] = None
     self._exception: Optional[LutronException] = None
 
     self.daemon = True
@@ -123,6 +124,54 @@ def connect(self) -> None:
           raise self._exception
         raise LutronConnectionError("Failed to connect to Lutron controller")
 
+  def disconnect(self, timeout: float = 5.0) -> None:
+    """Closes the connection and stops the reader thread.
+
+    Idempotent, and safe to call from any thread. Blocks until the thread has
+    exited or timeout seconds elapse, except when called from the reader thread
+    itself (e.g. inside a receive callback), where shutdown is scheduled and
+    this returns immediately rather than joining itself.
+
+    Terminal: a LutronConnection that has been disconnected cannot be
+    reconnected, because the underlying thread cannot be restarted. Callers
+    that need a new session should construct a new Lutron.
+    """
+    with self._lock:
+      self._done = True
+      self._connect_cond.notify_all()
+    if not self.is_alive():
+      # Either never started, or already finished. Either way the loop is ours
+      # to close, and join() would raise if the thread was never started.
+      if not self._loop.is_closed():
+        self._loop.close()
+      return
+    try:
+      self._loop.call_soon_threadsafe(self._shutdown_on_loop)
+    except RuntimeError:
+      # The reader finished and run() closed the loop between the is_alive()
+      # check above and this call. Nothing left to wake; checking is_closed()
+      # first would only narrow that window, not remove it.
+      pass
+    if threading.current_thread() is self:
+      # A receive callback asked us to disconnect. Shutdown is queued on our own
+      # loop and will run as the task unwinds; joining ourselves would raise.
+      return
+    self.join(timeout)
+    if self.is_alive():
+      _LOGGER.warning("Reader thread did not exit within %.1fs", timeout)
+
+  def _shutdown_on_loop(self) -> None:
+    """Wakes the reader for shutdown. Runs on the reader thread's loop.
+
+    Cancelling the task is what does the work: it unblocks a pending readline()
+    and cuts short the reconnect backoff alike. Closing the writer first means
+    the socket goes away even if the task is between awaits.
+    """
+    if self._writer:
+      self._writer.close()
+    if self._main_task is not None:
+      self._main_task.cancel()
+
   def send(self, cmd: str) -> None:
     """Sends the specified command to the lutron controller.
 
@@ -272,9 +321,18 @@ async def _main_loop(self) -> None:
         self._disconnect_locked()
 
       if not self._done:
-        # don't spam reconnect
+        # don't spam reconnect. disconnect() cancels the task, so this does not
+        # delay shutdown.
         await asyncio.sleep(5)
 
+  async def _runner(self) -> None:
+    """Wraps _main_loop so that disconnect() has a task it can cancel."""
+    self._main_task = asyncio.current_task()
+    try:
+      await self._main_loop()
+    except asyncio.CancelledError:
+      _LOGGER.debug("Reader task cancelled during shutdown")
+
   def run(self) -> None:
     """Main entry point into our receive thread.
 
@@ -283,10 +341,21 @@ def run(self) -> None:
     _LOGGER.info("Started")
     asyncio.set_event_loop(self._loop)
     try:
-      self._loop.run_until_complete(self._main_loop())
+      self._loop.run_until_complete(self._runner())
     except Exception:
       _LOGGER.exception("Uncaught exception in run")
       raise
+    finally:
+      # Anyone blocked in connect() must not wait on a thread that is gone.
+      with self._lock:
+        self._done = True
+        self._connected = False
+        self._connect_cond.notify_all()
+      try:
+        self._loop.run_until_complete(self._loop.shutdown_asyncgens())
+      finally:
+        self._loop.close()
+        _LOGGER.info("Stopped")
 
 
 class LutronXmlDbParser(object):
@@ -607,6 +676,15 @@ def connect(self) -> None:
     """Connects to the Lutron controller to send and receive commands and status"""
     self._conn.connect()
 
+  def disconnect(self, timeout: float = 5.0) -> None:
+    """Closes the connection and stops the background reader thread.
+
+    Idempotent. Blocks until the thread exits or timeout seconds elapse, except
+    when called from the reader thread itself, where it returns immediately. A
+    disconnected Lutron cannot be reconnected; construct a new one instead.
+    """
+    self._conn.disconnect(timeout)
+
   def send(self, op: str, cmd: str, integration_id: int, *args: Any) -> None:
     """Formats and sends the requested command to the Lutron controller."""
     out_cmd = ",".join(

tests/test_connection.py

@@ -309,5 +309,181 @@ async def do_login() -> None:
         self.assertIsNone(self.conn._exception)
 
 
+
+class TestLutronConnectionDisconnect(unittest.TestCase):
+    """disconnect() must stop the reader thread and free the event loop.
+
+    These drive the real thread rather than _main_loop directly, because the
+    whole point is the cross-thread shutdown handshake.
+    """
+
+    def setUp(self) -> None:
+        self.mock_reader = AsyncMock()
+        self.mock_writer = MagicMock()
+        self.mock_writer.drain = AsyncMock()
+
+        async def factory(host: str, port: int, connect_timeout: Optional[float] = None,
+                          encoding: Optional[str] = None) -> Tuple[AsyncMock, MagicMock]:
+            return self.mock_reader, self.mock_writer
+
+        self.conn = LutronConnection('127.0.0.1', 'user', 'pass', lambda line: None,
+                                     connection_factory=factory)
+
+    def tearDown(self) -> None:
+        # Never leave a reader thread behind, even if an assertion failed.
+        self.conn.disconnect(timeout=2.0)
+
+    def _connect_and_idle(self) -> None:
+        """Connects, then parks the reader in readline() like a live session."""
+        async def do_login() -> None:
+            self.conn._reader = self.mock_reader
+            self.conn._writer = self.mock_writer
+
+        async def park() -> bytes:
+            await asyncio.Event().wait()   # a connection with no traffic
+            return b""                     # pragma: no cover
+
+        self.mock_reader.readline = park
+        with patch.object(LutronConnection, '_do_login', side_effect=do_login):
+            self.conn.connect()
+
+    def test_disconnect_stops_thread_and_closes_loop(self) -> None:
+        """The reader is blocked in readline(); disconnect() must still stop it.
+
+        Setting _done alone cannot do this -- readline() only returns on data or
+        a closed socket -- so this fails without the cancellation in
+        _shutdown_on_loop.
+        """
+        self._connect_and_idle()
+        self.assertTrue(self.conn.is_alive())
+
+        self.conn.disconnect(timeout=5.0)
+
+        self.assertFalse(self.conn.is_alive(), "reader thread should have exited")
+        self.assertTrue(self.conn._loop.is_closed(), "event loop should be closed")
+        # the socket goes away, not just the thread
+        self.mock_writer.close.assert_called()
+
+    def test_disconnect_is_idempotent(self) -> None:
+        """Repeat calls must not raise, including after the thread is gone."""
+        self._connect_and_idle()
+        self.conn.disconnect(timeout=5.0)
+        self.conn.disconnect(timeout=5.0)
+        self.conn.disconnect(timeout=5.0)
+        self.assertFalse(self.conn.is_alive())
+        self.assertTrue(self.conn._loop.is_closed())
+
+    def test_disconnect_without_connect(self) -> None:
+        """Never started: must not raise, and must still close the loop.
+
+        Thread.join() raises RuntimeError on a thread that was never started, so
+        disconnect() has to check is_alive() before joining.
+        """
+        self.conn.disconnect(timeout=5.0)
+        self.assertFalse(self.conn.is_alive())
+        self.assertTrue(self.conn._loop.is_closed())
+
+    def test_disconnect_interrupts_reconnect_backoff(self) -> None:
+        """A disconnect during the 5s backoff must not wait the backoff out.
+
+        The first session parks in readline() so connect() returns against a
+        connection that is genuinely up; the test then releases that read to
+        drop it, and every reconnect fails, leaving the loop in the backoff.
+        With a plain asyncio.sleep(5) there, disconnect() blocks until the sleep
+        expires; with the interruptible wait it returns at once.
+        """
+        calls = {'n': 0}
+        in_backoff = threading.Event()
+        # Built on the reader's loop, not here: on Python 3.9 asyncio.Event()
+        # binds a loop via get_event_loop() at construction, and this thread has
+        # none. By the time connect() returns, _do_login has populated it.
+        loop_state: dict[str, asyncio.Event] = {}
+
+        async def park_then_drop() -> bytes:
+            await loop_state['release'].wait()
+            return b""
+
+        async def do_login() -> None:
+            calls['n'] += 1
+            if calls['n'] == 1:
+                loop_state['release'] = asyncio.Event()
+                self.conn._reader = self.mock_reader
+                self.conn._writer = self.mock_writer
+                self.mock_reader.readline = park_then_drop
+                return
+            in_backoff.set()
+            raise OSError(113, "Connect call failed")
+
+        with patch.object(LutronConnection, '_do_login', side_effect=do_login):
+            self.conn.connect()
+            # Drop the live session from the loop thread, which owns the event.
+            self.conn._loop.call_soon_threadsafe(loop_state['release'].set)
+            self.assertTrue(in_backoff.wait(timeout=10.0), "never reached the backoff")
+
+            start = time.monotonic()
+            self.conn.disconnect(timeout=10.0)
+            elapsed = time.monotonic() - start
+
+        self.assertFalse(self.conn.is_alive())
+        self.assertLess(elapsed, 4.0,
+                        f"disconnect waited out the backoff ({elapsed:.2f}s)")
+
+
+    def test_disconnect_survives_loop_closing_underneath_it(self) -> None:
+        """The reader may finish and close the loop mid-disconnect().
+
+        is_alive() can be true while run()'s finally is already closing the
+        loop, so call_soon_threadsafe races loop closure and raises
+        RuntimeError: Event loop is closed. Checking is_closed() first only
+        narrows that window. disconnect() must swallow it rather than propagate
+        into the caller -- in Home Assistant that is async_unload_entry.
+        """
+        self._connect_and_idle()
+
+        with patch.object(self.conn._loop, 'call_soon_threadsafe',
+                          side_effect=RuntimeError("Event loop is closed")):
+            self.conn.disconnect(timeout=1.0)   # must not raise
+
+        # The thread is still up (nothing woke it); a real disconnect still works.
+        self.conn.disconnect(timeout=5.0)
+        self.assertFalse(self.conn.is_alive())
+
+    def test_disconnect_from_the_reader_thread(self) -> None:
+        """A receive callback may call disconnect(); joining self would raise.
+
+        Thread.join() on the current thread raises "cannot join current thread",
+        so disconnect() schedules the shutdown and returns instead of joining.
+        The thread then unwinds on its own.
+        """
+        errors: List[BaseException] = []
+        # Hold the callback until connect() has returned on the main thread.
+        # Disconnecting sooner races connect()'s own predicate, which is a
+        # separate pre-existing issue and not what this test is about.
+        proceed = threading.Event()
+
+        def recv_cb(line: str) -> None:
+            if not proceed.wait(5.0):
+                return
+            try:
+                self.conn.disconnect(timeout=5.0)
+            except BaseException as exc:      # noqa: BLE001 - recorded, re-checked below
+                errors.append(exc)
+
+        self.conn._recv_cb = recv_cb
+        self.mock_reader.readline = AsyncMock(return_value=b"~OUTPUT,1,1,100.0\r\n")
+
+        async def do_login() -> None:
+            self.conn._reader = self.mock_reader
+            self.conn._writer = self.mock_writer
+
+        with patch.object(LutronConnection, '_do_login', side_effect=do_login):
+            self.conn.connect()
+            proceed.set()
+            self.conn.join(5.0)
+
+        self.assertEqual(errors, [], f"disconnect() raised on the reader thread: {errors}")
+        self.assertFalse(self.conn.is_alive(), "reader thread should have unwound")
+
+
 if __name__ == '__main__':
     unittest.main()