fix(mqtt): retire session workers on reconnect

mqtt_demo/__main__.py

@@ -166,7 +166,7 @@ def shutdown(*_):
         stopping.set()
         logger.info("shutting down…")
         for b in bridges:
-            b.stop.set()
+            b.request_stop()
             try: b.set_availability(False)
             except Exception: pass
 

mqtt_demo/bridge.py

@@ -87,6 +87,7 @@ def _href_to_segs(href: str) -> list[str]:
 # fixed-source-port reconnect invariant is untouched (see session_once).
 _GATE_RETRIES = 1
 _GATE_TIMEOUT_S = 4.0
+_WORKER_JOIN_TIMEOUT_S = 2.0
 
 
 class PushBridge:
@@ -123,6 +124,8 @@ def __init__(self,
         self.last_cycle_pub = None
         self.last_avail_pub: str | None = None
         self.stop = threading.Event()
+        self._session_stop_lock = threading.Lock()
+        self._session_stop: threading.Event | None = None
         self.started_ts = time.time()
         self.session_started_ts = None
         self.last_change_ts = None
@@ -175,6 +178,13 @@ def __init__(self,
                 app.topic_prefix, shared.HA_DISCOVERY_PREFIX, app.device_name,
                 model=descriptor.name.title()))
 
+    def request_stop(self) -> None:
+        """Stop the bridge and wake workers belonging to its current session."""
+        self.stop.set()
+        with self._session_stop_lock:
+            if self._session_stop is not None:
+                self._session_stop.set()
+
     # ---- cache plumbing ---------------------------------------------
 
     def _on_cache_change(self, changed: bool, source: str) -> None:
@@ -423,25 +433,54 @@ def _run_session_inner(self, sess):
         self.keepalive = keepalive
         self.observe_refresh = observe_refresh
 
+        # These workers belong to this DTLS session, not to the bridge
+        # process. A reconnect must retire them before the replacement
+        # session starts or they continue operating on the closed session.
+        session_stop = threading.Event()
+        with self._session_stop_lock:
+            self._session_stop = session_stop
+            # ``request_stop()`` sets the bridge event before taking this
+            # lock. Checking it while publishing the handle prevents a lost
+            # wakeup if shutdown races this session handoff.
+            if self.stop.is_set():
+                session_stop.set()
+
         sched_t = threading.Thread(
-            target=scheduler.run_forever, args=(self.stop,),
+            target=scheduler.run_forever, args=(session_stop,),
             daemon=True, name=f'{self.app.klass}-poll')
         ka_t = threading.Thread(
-            target=keepalive.run_forever, args=(self.stop,),
+            target=keepalive.run_forever, args=(session_stop,),
             daemon=True, name=f'{self.app.klass}-ping')
         ref_t = threading.Thread(
-            target=observe_refresh.run_forever, args=(self.stop,),
+            target=observe_refresh.run_forever, args=(session_stop,),
             daemon=True, name=f'{self.app.klass}-obsref')
-        sched_t.start()
-        ka_t.start()
-        ref_t.start()
+        workers = (sched_t, ka_t, ref_t)
+        started_workers = []
 
         try:
+            for worker in workers:
+                worker.start()
+                started_workers.append(worker)
             sess.join()
         finally:
+            # A worker already inside a tick can finish after the reader
+            # exits. Disable old-session reachability callbacks first so it
+            # cannot change availability after a replacement takes over.
+            keepalive.on_reachable = None
+            keepalive.on_unreachable = None
+            session_stop.set()
+            join_deadline = time.monotonic() + _WORKER_JOIN_TIMEOUT_S
+            for worker in started_workers:
+                worker.join(max(0.0, join_deadline - time.monotonic()))
+                if worker.is_alive():
+                    self.log.warning(
+                        "session worker did not stop: %s", worker.name)
             self.scheduler = None
             self.keepalive = None
             self.observe_refresh = None
+            with self._session_stop_lock:
+                if self._session_stop is session_stop:
+                    self._session_stop = None
 
     def _seed_from_device0(self, sess):
         code, pl = sess.get(self.descriptor.seed_path, timeout=15.0)

tests/test_worker_cleanup.py

@@ -3,7 +3,12 @@
 from __future__ import annotations
 
 import threading
+from types import SimpleNamespace
 
+import pytest
+
+from mqtt_demo import bridge as bridge_module
+from mqtt_demo.bridge import PushBridge
 from smartthings_local.ocf.keepalive import KeepaliveTask
 from smartthings_local.ocf.observe_refresh import ObserveRefreshTask
 from smartthings_local.ocf.poll_scheduler import PollScheduler, PollTier
@@ -73,3 +78,238 @@ def test_poll_scheduler_worker_stops_without_leaking_thread():
         [PollTier("idle", interval_s=3600.0, paths=())],
     )
     _assert_worker_stops(scheduler.run_forever, "test-poll-scheduler")
+
+
+class _SessionWorker:
+    def __init__(self, *args, **kwargs):
+        self.stop = None
+        self.started = threading.Event()
+        self.exited = threading.Event()
+        self.on_reachable = kwargs.get("on_reachable")
+        self.on_unreachable = kwargs.get("on_unreachable")
+        self.last_success_ts = 0.0
+
+    def run_forever(self, stop):
+        self.stop = stop
+        self.started.set()
+        stop.wait()
+        self.exited.set()
+
+
+class _JoinedSession:
+    def __init__(self, workers, start_index, error=None):
+        self.workers = workers
+        self.start_index = start_index
+        self.error = error
+
+    def join(self):
+        assert all(
+            worker.started.wait(_THREAD_DEADLINE_S)
+            for worker in self.workers[self.start_index:]
+        )
+        if self.error is not None:
+            raise self.error
+
+
+class _BlockingJoinedSession(_JoinedSession):
+    def __init__(self, workers):
+        super().__init__(workers, 0)
+        self.joined = threading.Event()
+        self.release = threading.Event()
+
+    def join(self):
+        super().join()
+        self.joined.set()
+        assert self.release.wait(_THREAD_DEADLINE_S)
+
+
+def _bridge():
+    bridge = object.__new__(PushBridge)
+    bridge.descriptor = SimpleNamespace(
+        observe_paths=(),
+        poll_tiers=[],
+        is_active=lambda _state: False,
+    )
+    bridge.shared = SimpleNamespace(PING_INTERVAL_S=3600.0)
+    bridge.app = SimpleNamespace(klass="test")
+    bridge.log = SimpleNamespace(info=lambda *args: None, warning=lambda *args: None)
+    bridge.cache = SimpleNamespace(links={})
+    bridge.stop = threading.Event()
+    bridge._session_stop_lock = threading.Lock()
+    bridge._session_stop = None
+    bridge.scheduler = None
+    bridge.keepalive = None
+    bridge.observe_refresh = None
+    bridge._seed_from_device0 = lambda _session: None
+    bridge._retag_logger_with_serial = lambda: None
+    bridge.maybe_publish_state = lambda **kwargs: None
+    bridge.set_availability = lambda _online: None
+    return bridge
+
+
+@pytest.mark.parametrize(
+    ("session_count", "join_error"),
+    [
+        pytest.param(2, None, id="reconnect"),
+        pytest.param(1, RuntimeError("reader failed"), id="reader-error"),
+    ],
+)
+def test_bridge_retires_session_workers_before_returning(
+    monkeypatch, session_count, join_error
+):
+    workers = []
+
+    def worker_factory(*args, **kwargs):
+        worker = _SessionWorker(*args, **kwargs)
+        workers.append(worker)
+        return worker
+
+    monkeypatch.setattr(bridge_module, "PollScheduler", worker_factory)
+    monkeypatch.setattr(bridge_module, "KeepaliveTask", worker_factory)
+    monkeypatch.setattr(bridge_module, "ObserveRefreshTask", worker_factory)
+
+    bridge = _bridge()
+
+    session_stops = []
+    for session_index in range(session_count):
+        start_index = len(workers)
+        session = _JoinedSession(
+            workers,
+            start_index,
+            error=join_error if session_index == session_count - 1 else None,
+        )
+        if session.error is None:
+            bridge._run_session_inner(session)
+        else:
+            with pytest.raises(RuntimeError, match="reader failed"):
+                bridge._run_session_inner(session)
+
+        session_workers = workers[start_index:]
+        assert len(session_workers) == 3
+        assert len({id(worker.stop) for worker in session_workers}) == 1
+        session_stops.append(session_workers[0].stop)
+
+    assert len(workers) == session_count * 3
+    assert len({id(stop) for stop in session_stops}) == session_count
+    assert all(stop is not bridge.stop for stop in session_stops)
+    assert all(stop.is_set() for stop in session_stops)
+    assert not bridge.stop.is_set()
+    assert all(worker.exited.is_set() for worker in workers)
+    keepalive_workers = tuple(
+        workers[index] for index in range(1, len(workers), 3)
+    )
+    assert all(worker.on_reachable is None for worker in keepalive_workers)
+    assert all(worker.on_unreachable is None for worker in keepalive_workers)
+    assert bridge.scheduler is None
+    assert bridge.keepalive is None
+    assert bridge.observe_refresh is None
+    assert bridge._session_stop is None
+
+
+def test_bridge_retires_started_worker_when_later_thread_fails_to_start(
+    monkeypatch,
+):
+    workers = []
+
+    def worker_factory(*args, **kwargs):
+        worker = _SessionWorker(*args, **kwargs)
+        workers.append(worker)
+        return worker
+
+    monkeypatch.setattr(bridge_module, "PollScheduler", worker_factory)
+    monkeypatch.setattr(bridge_module, "KeepaliveTask", worker_factory)
+    monkeypatch.setattr(bridge_module, "ObserveRefreshTask", worker_factory)
+
+    bridge = _bridge()
+
+    original_start = threading.Thread.start
+    start_count = 0
+
+    def fail_second_start(thread):
+        nonlocal start_count
+        start_count += 1
+        if start_count == 2:
+            raise RuntimeError("synthetic thread start failure")
+        original_start(thread)
+
+    monkeypatch.setattr(threading.Thread, "start", fail_second_start)
+
+    with pytest.raises(RuntimeError, match="synthetic thread start failure"):
+        bridge._run_session_inner(SimpleNamespace(join=lambda: None))
+
+    assert len(workers) == 3
+    assert workers[0].started.wait(_THREAD_DEADLINE_S)
+    assert workers[0].exited.wait(_THREAD_DEADLINE_S)
+    assert workers[0].stop is not bridge.stop
+    assert workers[0].stop.is_set()
+    assert workers[1].stop is None
+    assert workers[2].stop is None
+    assert bridge.scheduler is None
+    assert bridge.keepalive is None
+    assert bridge.observe_refresh is None
+    assert bridge._session_stop is None
+
+
+def test_request_stop_wakes_current_session_workers(monkeypatch):
+    workers = []
+
+    def worker_factory(*args, **kwargs):
+        worker = _SessionWorker(*args, **kwargs)
+        workers.append(worker)
+        return worker
+
+    monkeypatch.setattr(bridge_module, "PollScheduler", worker_factory)
+    monkeypatch.setattr(bridge_module, "KeepaliveTask", worker_factory)
+    monkeypatch.setattr(bridge_module, "ObserveRefreshTask", worker_factory)
+
+    bridge = _bridge()
+    session = _BlockingJoinedSession(workers)
+    session_thread = threading.Thread(
+        target=bridge._run_session_inner,
+        args=(session,),
+        daemon=True,
+    )
+    session_thread.start()
+
+    try:
+        assert session.joined.wait(_THREAD_DEADLINE_S)
+        session_stop = bridge._session_stop
+        assert session_stop is not None
+        assert not session_stop.is_set()
+
+        bridge.request_stop()
+
+        assert bridge.stop.is_set()
+        assert session_stop.is_set()
+        assert all(
+            worker.exited.wait(_THREAD_DEADLINE_S) for worker in workers
+        )
+        assert session_thread.is_alive()
+    finally:
+        session.release.set()
+        session_thread.join(_THREAD_DEADLINE_S)
+
+    assert not session_thread.is_alive()
+    assert bridge._session_stop is None
+
+
+def test_session_workers_observe_stop_requested_before_handoff(monkeypatch):
+    workers = []
+
+    def worker_factory(*args, **kwargs):
+        worker = _SessionWorker(*args, **kwargs)
+        workers.append(worker)
+        return worker
+
+    monkeypatch.setattr(bridge_module, "PollScheduler", worker_factory)
+    monkeypatch.setattr(bridge_module, "KeepaliveTask", worker_factory)
+    monkeypatch.setattr(bridge_module, "ObserveRefreshTask", worker_factory)
+
+    bridge = _bridge()
+    bridge.request_stop()
+    bridge._run_session_inner(_JoinedSession(workers, 0))
+
+    assert len(workers) == 3
+    assert all(worker.stop.is_set() for worker in workers)
+    assert all(worker.exited.wait(_THREAD_DEADLINE_S) for worker in workers)
+    assert bridge._session_stop is None