From c9c06a355f2c5925801caf56c7bfff9b64d610ed Mon Sep 17 00:00:00 2001 From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com> Date: Sun, 16 Aug 2026 13:57:55 -0500 Subject: [PATCH] release the shutdown barrier on worker crash and replay retained publishes the broker never acked --- frigate/comms/mqtt.py | 54 +++++++++++++++++++ frigate/test/test_mqtt_lifecycle.py | 82 ++++++++++++++++++++++++++++- 2 files changed, 134 insertions(+), 2 deletions(-) diff --git a/frigate/comms/mqtt.py b/frigate/comms/mqtt.py index 7ec4b49df1..c80f723c40 100644 --- a/frigate/comms/mqtt.py +++ b/frigate/comms/mqtt.py @@ -50,6 +50,7 @@ class MqttClient(Communicator): self._callback_queue: queue.Queue[tuple[Any, ...]] = queue.Queue() self._retained_lock = threading.Lock() self._pending_retained: dict[str, tuple[Any, bool]] = {} + self._inflight_retained: dict[int, tuple[str, Any]] = {} self._subscription_mid: int | None = None self._subscription_ready = False self._next_connect_time = 0.0 @@ -273,6 +274,7 @@ class MqttClient(Communicator): client.on_disconnect = self._on_disconnect client.on_message = self._on_message client.on_subscribe = self._on_subscribe + client.on_publish = self._on_publish client.will_set( self.mqtt_config.topic_prefix + "/available", payload="offline", @@ -329,6 +331,9 @@ class MqttClient(Communicator): self._publish_offline_availability() self.connected = False finally: + # nothing drains the queue once the loop is gone, so release any + # waiter here or stop() blocks for the full flush timeout + self._requeue_disconnected_publishes() self._cleanup_client() def _publish_offline_availability(self) -> None: @@ -412,6 +417,7 @@ class MqttClient(Communicator): self.connected = False self._subscription_ready = False self._subscription_mid = None + self._requeue_inflight_retained() client = self.client self.client = None @@ -437,6 +443,19 @@ class MqttClient(Communicator): logger.info("MQTT reconnect scheduled in %.1fs", MQTT_RECONNECT_INTERVAL) self._cleanup_client() + def _requeue_inflight_retained(self) -> None: + """Rebuffer retained publishes paho took but the broker never acked. + + Dropping the client drops paho's outbound queue with it, and the session + is clean, so the broker will not resume delivery on the new one. + """ + with self._retained_lock: + inflight = list(self._inflight_retained.values()) + self._inflight_retained.clear() + + for topic, payload in inflight: + self._queue_retained(topic, payload, True, overwrite=False) + def _buffer_undelivered( self, queued_publish: QueuedPublish, overwrite: bool = True ) -> None: @@ -486,6 +505,8 @@ class MqttClient(Communicator): self._handle_subscribe_event(event[1], event[2]) elif event_type == "message": self._handle_inbound_message(event[1], event[2]) + elif event_type == "published": + self._handle_publish_event(event[1]) def _drain_publish_queue(self) -> None: """Publish queued work only after the session is fully subscribed.""" @@ -548,10 +569,24 @@ class MqttClient(Communicator): self._schedule_reconnect() return + # a successful rc only means paho accepted the message; above qos 0 it + # is not durable until the broker acks, so keep a copy for replay + if queued_publish.retain and not message_info.is_published(): + with self._retained_lock: + self._inflight_retained[message_info.mid] = ( + queued_publish.topic, + queued_publish.payload, + ) + if queued_publish.done is not None: self._wait_for_publish(message_info) queued_publish.done.set() + def _handle_publish_event(self, mid: int) -> None: + """Drop the replay copy once the broker has acknowledged the message.""" + with self._retained_lock: + self._inflight_retained.pop(mid, None) + def _wait_for_publish(self, message_info: mqtt.MQTTMessageInfo) -> None: """Pump the loop until a shutdown-critical publish is acknowledged.""" deadline = time.monotonic() + MQTT_SHUTDOWN_FLUSH_TIMEOUT @@ -759,6 +794,25 @@ class MqttClient(Communicator): """Handle subscribe acknowledgements from paho.""" self._callback_queue.put(("subscribed", mid, reason_codes)) + def _on_publish( + self, + client: mqtt.Client, + userdata: Any, + mid: int, + reason_code: mqtt.ReasonCode, # type: ignore[name-defined] + properties: Any, + ) -> None: + """Handle publish acknowledgements from paho. + + Only tracked retained messages need an event. At the default qos 0 + nothing is tracked, so this stays off the hot publish path. + """ + with self._retained_lock: + if mid not in self._inflight_retained: + return + + self._callback_queue.put(("published", mid)) + def _on_message( self, client: mqtt.Client, diff --git a/frigate/test/test_mqtt_lifecycle.py b/frigate/test/test_mqtt_lifecycle.py index b690e53b88..aa37870f47 100644 --- a/frigate/test/test_mqtt_lifecycle.py +++ b/frigate/test/test_mqtt_lifecycle.py @@ -399,8 +399,9 @@ class TestMqttClientLifecycle(unittest.TestCase): mock_client = MagicMock() mock_client.loop.return_value = mqtt.MQTT_ERR_SUCCESS self.client.client = mock_client - message_info = MagicMock(rc=mqtt.MQTT_ERR_SUCCESS) - message_info.is_published.side_effect = [False, True] + message_info = MagicMock(rc=mqtt.MQTT_ERR_SUCCESS, mid=1) + # inflight tracking checks once, then _wait_for_publish polls + message_info.is_published.side_effect = [False, False, True] mock_client.publish.return_value = message_info barrier = MagicMock() @@ -453,6 +454,83 @@ class TestMqttClientLifecycle(unittest.TestCase): self.client._pending_retained["frigate/available"], ("stopped", True) ) + def test_shutdown_barrier_releases_when_worker_crashes(self) -> None: + """stop() can queue the final publish just as the worker dies, and + nothing drains the queue after that.""" + barrier = threading.Event() + self.client._publish_queue.put( + QueuedPublish("frigate/available", "stopped", True, barrier) + ) + + with patch.object( + self.client, + "_mqtt_loop_worker", + side_effect=RuntimeError("unexpected bug"), + ): + self.client._worker_main() + + self.assertTrue(barrier.is_set()) + + def test_unacked_retained_publish_survives_reconnect(self) -> None: + """Above qos 0 a successful rc only means paho queued the message, and + dropping the client drops its outbound queue with it.""" + mock_client = MagicMock() + message_info = MagicMock(rc=mqtt.MQTT_ERR_SUCCESS, mid=12) + message_info.is_published.return_value = False + mock_client.publish.return_value = message_info + self.client.client = mock_client + self.client.connected = True + + self.client._publish_direct( + QueuedPublish("frigate/profile/state", "armed", True) + ) + self.assertEqual( + self.client._inflight_retained[12], ("frigate/profile/state", "armed") + ) + + self.client._cleanup_client() + + self.assertEqual(self.client._inflight_retained, {}) + self.assertEqual( + self.client._pending_retained["frigate/profile/state"], ("armed", True) + ) + + def test_acked_retained_publish_is_not_replayed(self) -> None: + mock_client = MagicMock() + message_info = MagicMock(rc=mqtt.MQTT_ERR_SUCCESS, mid=12) + message_info.is_published.return_value = False + mock_client.publish.return_value = message_info + self.client.client = mock_client + self.client.connected = True + + self.client._publish_direct( + QueuedPublish("frigate/profile/state", "armed", True) + ) + self.client._on_publish(mock_client, None, 12, MagicMock(), None) + self.client._drain_callback_queue() + + self.assertEqual(self.client._inflight_retained, {}) + + self.client._cleanup_client() + + self.assertEqual(self.client._pending_retained, {}) + + def test_already_published_retained_is_not_tracked(self) -> None: + """At the default qos 0 paho reports the message as published inline, + so there is nothing to wait on.""" + mock_client = MagicMock() + message_info = MagicMock(rc=mqtt.MQTT_ERR_SUCCESS, mid=12) + message_info.is_published.return_value = True + mock_client.publish.return_value = message_info + self.client.client = mock_client + self.client.connected = True + + self.client._publish_direct( + QueuedPublish("frigate/profile/state", "armed", True) + ) + + self.assertEqual(self.client._inflight_retained, {}) + def test_wait_for_publish_survives_disconnect_during_wait(self) -> None: mock_client = MagicMock() self.client.client = mock_client