release the shutdown barrier on worker crash and replay retained publishes the broker never acked

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,

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