diff --git a/mqtt_demo/__main__.py b/mqtt_demo/__main__.py index b1a0ee5..523412f 100644 --- a/mqtt_demo/__main__.py +++ b/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 diff --git a/mqtt_demo/bridge.py b/mqtt_demo/bridge.py index d1f66d2..8a8d8a9 100644 --- a/mqtt_demo/bridge.py +++ b/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) diff --git a/tests/test_worker_cleanup.py b/tests/test_worker_cleanup.py index 11694b6..472f361 100644 --- a/tests/test_worker_cleanup.py +++ b/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