From 4e7fb23a9ea5a2b7ad2c7daef9d8c5017f05bc86 Mon Sep 17 00:00:00 2001 From: Riaz Virani Date: Thu, 8 Oct 2026 17:07:23 -0700 Subject: [PATCH 1/5] ENT-1646: Add MQTT Writer v2 with fire_and_forget Add roboflow_enterprise/mqtt_writer_sink@v2 next to the unchanged v1. The new optional fire_and_forget input (default False) lets a run queue the message and return without waiting for the broker's acknowledgement or, after the first run, for a reconnect, so a broker outage no longer stalls every frame. While disconnected, QoS 0 messages are dropped and QoS 1/2 messages are queued in paho and sent after reconnect, both reported with error_status set. A first connect that fails keeps the client retrying in the background. v2 also hardens outages in both modes: the client holds at most 1000 QoS 1/2 messages and 1000 unsent packets, repeated failures and reconnect attempts are logged once per outage, the reconnect delay has a floor, and selector-supplied switches, port and qos are coerced the way the manifest validates them. The step declares the fire_and_forget_hides_persistence_failures and connection_and_state_rebuilt_per_request restrictions when the switch is on. Co-Authored-By: Claude Opus 5.5 --- .../test_workflow_with_mqtt_writer_v2.py | 680 +++++++ workflows/CHANGELOG.md | 11 + .../enterprise_blocks/loader.py | 4 + .../enterprise_blocks/sinks/mqtt_writer/v2.py | 1003 ++++++++++ ...e_and_forget_sink_workload_declarations.py | 37 +- .../core_steps/sinks/test_mqtt_writer_v2.py | 1635 +++++++++++++++++ 6 files changed, 3366 insertions(+), 4 deletions(-) create mode 100644 tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py create mode 100644 workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py create mode 100644 workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py new file mode 100644 index 0000000000..a9f79abc0b --- /dev/null +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py @@ -0,0 +1,680 @@ +import copy +import os +import socket +import threading +import time +from unittest.mock import patch + +import pytest +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v2 +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + MQTTWriterSinkBlockV2, +) + +from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.workflows.core_steps.sinks.noop import disabled_sink_message +from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.conftest import FakeMQTTBroker + + +def wait_until(predicate, timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return False + + +def closed_port() -> int: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.bind(("localhost", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +@pytest.mark.timeout(10) +def test_successful_connection_and_publishing(fake_mqtt_broker): + # given + block = MQTTWriterSinkBlockV2() + published_message = "Test message" + expected_message = "Message published successfully" + + fake_mqtt_broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=fake_mqtt_broker.start) + broker_thread.start() + + try: + # when + result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message=published_message, + timeout=2.0, + ) + + broker_thread.join(timeout=2) + + # then + assert result["error_status"] is False, "No error expected" + assert result["message"] == expected_message + + assert published_message.encode() in fake_mqtt_broker.messages[-1] + finally: + block.close() + + +@pytest.mark.timeout(10) +def test_rejected_connack_is_not_treated_as_connected(): + # given + broker = FakeMQTTBroker(connack_reason_code=5) + broker.messages_count_to_wait_for = 0 + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + block = MQTTWriterSinkBlockV2() + + try: + # when + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="should not be delivered", + timeout=1.0, + ) + broker_thread.join(timeout=2) + + # then + assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "code 5" in result["message"] + assert not block._connected.is_set() + assert broker.messages == [] + finally: + block.close() + broker.finish() + + +@pytest.mark.timeout(15) +def test_refused_connection_is_reported_once_and_not_retried(): + # given - a broker that refuses the credentials and would accept any retry + broker = FakeMQTTBroker(connack_reason_code=5, keep_serving=True) + broker_thread = threading.Thread(target=broker.serve, daemon=True) + broker_thread.start() + block = MQTTWriterSinkBlockV2() + kwargs = dict( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="should not be delivered", + timeout=0.5, + ) + + try: + # when - the first run hits the refusal, then the block lives on past + # several reconnect intervals (max delay is 2 * timeout) + started = time.monotonic() + first = block.run(**kwargs) + first_elapsed = time.monotonic() - started + time.sleep(2.5) + started = time.monotonic() + second = block.run(**kwargs) + second_elapsed = time.monotonic() - started + + # then - the refusal is named at once, the loop stopped, nothing published + assert first["error_status"] is True + assert "not authorised" in first["message"] + assert "Check username and password" in first["message"] + assert "Raise 'timeout'" not in first["message"] + assert first_elapsed < 0.4 + assert second == first + assert second_elapsed < 0.1 + assert broker.connections_accepted == 1 + assert broker.messages == [] + finally: + block.close() + broker.finish() + broker_thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_block_recovers_when_broker_becomes_available_after_first_run(): + # given - the port is bound but refuses connections, so the first TCP + # attempt fails; the failed run must leave no client or thread behind + broker = FakeMQTTBroker(listening=False) + broker.messages_count_to_wait_for = 1 + block = MQTTWriterSinkBlockV2() + + try: + # when - broker is not accepting connections yet + first_result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="first message", + timeout=0.3, + ) + + assert block.mqtt_client is None + + # broker becomes available, the next run connects from scratch + broker.listen() + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + + second_result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="second message", + timeout=2.0, + ) + broker_thread.join(timeout=2) + + # then + assert first_result["error_status"] is True + assert second_result["error_status"] is False + assert b"second message" in broker.messages[-1] + finally: + block.close() + broker.finish() + + +@pytest.mark.timeout(15) +def test_lost_connection_clears_readiness_and_run_reports_it(fake_mqtt_broker): + # given + block = MQTTWriterSinkBlockV2() + fake_mqtt_broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=fake_mqtt_broker.start) + broker_thread.start() + + try: + # when + first_result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message="delivered", + timeout=2.0, + ) + broker_thread.join(timeout=2) + fake_mqtt_broker.finish() + + assert wait_until(lambda: not block._connected.is_set()) + + second_result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message="lost", + timeout=2.0, + ) + + # then + assert first_result["error_status"] is False + assert second_result["error_status"] is True + assert "not connected" in second_result["message"].lower() + finally: + block.close() + + +@pytest.mark.timeout(10) +def test_missing_puback_reported_as_delivery_unknown_not_final_failure( + fake_mqtt_broker, +): + # given - broker accepts the PUBLISH but never sends PUBACK + block = MQTTWriterSinkBlockV2() + published_message = "qos1 message" + fake_mqtt_broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=fake_mqtt_broker.start) + broker_thread.start() + + try: + # when + result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message=published_message, + qos=1, + timeout=1.0, + ) + broker_thread.join(timeout=2) + + # then - the block reports delivery-unknown, yet the broker received it + assert result["error_status"] is True + assert "delivery status unknown" in result["message"].lower() + assert published_message.encode() in fake_mqtt_broker.messages[-1] + finally: + block.close() + + +@pytest.mark.timeout(10) +def test_fire_and_forget_returns_before_puback_and_message_still_arrives( + fake_mqtt_broker, +): + # given - broker accepts the PUBLISH but never sends PUBACK + block = MQTTWriterSinkBlockV2() + published_message = "fire and forget message" + fake_mqtt_broker.messages_count_to_wait_for = 2 + broker_thread = threading.Thread(target=fake_mqtt_broker.start) + broker_thread.start() + + try: + # when - a first run that publishes; the connect is the only wait + block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message="first", + qos=1, + timeout=2.0, + fire_and_forget=True, + ) + started = time.monotonic() + result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message=published_message, + qos=1, + timeout=2.0, + fire_and_forget=True, + ) + elapsed = time.monotonic() - started + + # then - no PUBACK wait, and the broker still receives the message + assert result["error_status"] is False + assert "queued" in result["message"].lower() + assert elapsed < 1.0 + assert wait_until( + lambda: any( + published_message.encode() in message + for message in fake_mqtt_broker.messages + ) + ) + finally: + block.close() + fake_mqtt_broker.finish() + broker_thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_fire_and_forget_cold_start_retries_in_background_without_blocking(): + # given - the port is bound but refuses connections when the pipeline starts + broker = FakeMQTTBroker(listening=False) + broker.messages_count_to_wait_for = 1 + block = MQTTWriterSinkBlockV2() + run_kwargs = dict( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + timeout=0.3, + fire_and_forget=True, + ) + broker_thread = None + + try: + # when - the first run fails, later runs while the broker is down + first_result = block.run(message="first", **run_kwargs) + client = block.mqtt_client + started = time.monotonic() + down_results = [block.run(message="down", **run_kwargs) for _ in range(5)] + down_elapsed = time.monotonic() - started + + # the broker comes up: the kept client's background loop reconnects + broker.listen() + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + assert wait_until(block._connected.is_set) + up_result = block.run(message="after recovery", **run_kwargs) + + # then + assert first_result["error_status"] is True + assert client is not None + assert all(result["error_status"] for result in down_results) + assert down_elapsed < 0.5 + assert up_result["error_status"] is False + assert block.mqtt_client is client + assert wait_until( + lambda: any(b"after recovery" in message for message in broker.messages) + ) + finally: + block.close() + broker.finish() + if broker_thread is not None: + broker_thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_fire_and_forget_cold_start_delivers_qos1_once_broker_is_up(): + # given - the port refuses connections when the pipeline starts + broker = FakeMQTTBroker(listening=False) + broker.messages_count_to_wait_for = 1 + block = MQTTWriterSinkBlockV2() + broker_thread = None + + try: + # when - the first run cannot connect + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="queued during outage", + qos=1, + timeout=0.3, + fire_and_forget=True, + ) + broker.listen() + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + + # then - reported, yet paho delivers it once the background loop connects + assert result["error_status"] is True + assert "queued" in result["message"] + assert wait_until( + lambda: any( + b"queued during outage" in message for message in broker.messages + ) + ) + finally: + block.close() + broker.finish() + if broker_thread is not None: + broker_thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_fire_and_forget_does_not_wait_for_a_lost_broker(fake_mqtt_broker): + # given + block = MQTTWriterSinkBlockV2() + fake_mqtt_broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=fake_mqtt_broker.start) + broker_thread.start() + + try: + first_result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message="delivered", + timeout=2.0, + fire_and_forget=True, + ) + broker_thread.join(timeout=2) + fake_mqtt_broker.finish() + assert wait_until(lambda: not block._connected.is_set()) + + # when + started = time.monotonic() + second_result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message="lost", + timeout=2.0, + fire_and_forget=True, + ) + elapsed = time.monotonic() - started + + # then - reported at once instead of after the 2 s readiness wait + assert first_result["error_status"] is False + assert second_result["error_status"] is True + assert "dropped" in second_result["message"].lower() + assert elapsed < 1.0 + finally: + block.close() + + +@pytest.mark.timeout(10) +def test_second_broker_is_rejected_instead_of_publishing_to_first(fake_mqtt_broker): + # given + block = MQTTWriterSinkBlockV2() + other_broker = FakeMQTTBroker() + fake_mqtt_broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=fake_mqtt_broker.start) + broker_thread.start() + + try: + # when + first_result = block.run( + host=fake_mqtt_broker.host, + port=fake_mqtt_broker.port, + topic="RoboflowTopic", + message="for broker A", + timeout=2.0, + ) + broker_thread.join(timeout=2) + + second_result = block.run( + host=other_broker.host, + port=other_broker.port, + topic="RoboflowTopic", + message="for broker B", + timeout=2.0, + ) + + # then + assert first_result["error_status"] is False + assert second_result["error_status"] is True + assert "parameters" in second_result["message"].lower() + assert other_broker.messages == [] + assert all( + b"for broker B" not in message for message in fake_mqtt_broker.messages + ) + finally: + block.close() + other_broker.finish() + + +@pytest.mark.timeout(15) +def test_publish_over_tls_with_ca_certificate(mqtt_test_certificates): + # given + broker = FakeMQTTBroker(tls_context=mqtt_test_certificates.server_context) + broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + block = MQTTWriterSinkBlockV2(allow_access_to_file_system=True) + + try: + # when + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="encrypted payload", + timeout=5.0, + encryption="tls", + ca_certificate_path=mqtt_test_certificates.ca_path, + ) + broker_thread.join(timeout=2) + + # then + assert result["error_status"] is False, result["message"] + assert result["message"] == "Message published successfully" + assert b"encrypted payload" in broker.messages[-1] + assert broker.handshake_failures == 0 + finally: + block.close() + broker.finish() + + +MQTT_SINK_WORKFLOW = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowParameter", "name": "host"}, + {"type": "WorkflowParameter", "name": "port"}, + {"type": "WorkflowParameter", "name": "message"}, + ], + "steps": [ + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt_sink", + "host": "$inputs.host", + "port": "$inputs.port", + "topic": "RoboflowTopic", + "message": "$inputs.message", + "timeout": 0.2, + } + ], + "outputs": [ + { + "type": "JsonField", + "name": "status", + "selector": "$steps.mqtt_sink.error_status", + }, + ], +} + + +@pytest.fixture +def enterprise_blocks_enabled(monkeypatch): + from inference.core.env import ENTERPRISE_BLOCKS_PLUGIN + from inference.core.workflows.execution_engine.introspection import blocks_loader + from inference.core.workflows.execution_engine.v1.compiler.core import ( + COMPILATION_CACHE, + ) + + plugins = [p for p in os.getenv("WORKFLOWS_PLUGINS", "").split(",") if p] + if ENTERPRISE_BLOCKS_PLUGIN not in plugins: # env.py may have expanded it already; + plugins.append(ENTERPRISE_BLOCKS_PLUGIN) # a duplicate trips the clash check + monkeypatch.setenv("WORKFLOWS_PLUGINS", ",".join(plugins)) + blocks_loader.clear_caches() + yield + blocks_loader.clear_caches() + # drop graphs compiled with enterprise blocks from the process-wide cache + with COMPILATION_CACHE._cache_lock: + COMPILATION_CACHE._cache.clear() + COMPILATION_CACHE._keys_buffer.clear() + + +@pytest.mark.timeout(15) +def test_workflow_with_failing_mqtt_sink_logs_and_continues( + enterprise_blocks_enabled, +): + # given - nothing listens on the port, so every sink run fails + execution_engine = ExecutionEngine.init( + workflow_definition=MQTT_SINK_WORKFLOW, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + runtime_parameters = { + "host": "localhost", + "port": closed_port(), + "message": "frame payload", + } + + # when - two consecutive frames + with patch.object(v2, "logger") as mock_logger: + first_frame_result = execution_engine.run(runtime_parameters=runtime_parameters) + second_frame_result = execution_engine.run( + runtime_parameters=runtime_parameters + ) + + # then - failures are surfaced in logs and outputs, execution continues + assert any( + call.args[0].startswith("MQTT Writer failure") + for call in mock_logger.error.call_args_list + ) + assert first_frame_result[0]["status"] is True + assert second_frame_result[0]["status"] is True + + +@pytest.mark.timeout(15) +def test_workflow_with_unwired_failing_mqtt_sink_still_logs_and_continues( + enterprise_blocks_enabled, +): + # given - no MQTT outputs are wired, so logs are the only failure signal + workflow_definition = copy.deepcopy(MQTT_SINK_WORKFLOW) + workflow_definition["outputs"] = [ + {"type": "JsonField", "name": "echo", "selector": "$inputs.message"}, + ] + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + runtime_parameters = { + "host": "localhost", + "port": closed_port(), + "message": "frame payload", + } + + # when + with patch.object(v2, "logger") as mock_logger: + first_frame_result = execution_engine.run(runtime_parameters=runtime_parameters) + second_frame_result = execution_engine.run( + runtime_parameters=runtime_parameters + ) + + # then + assert any( + call.args[0].startswith("MQTT Writer failure") + for call in mock_logger.error.call_args_list + ) + assert first_frame_result[0]["echo"] == "frame payload" + assert second_frame_result[0]["echo"] == "frame payload" + + +@pytest.mark.timeout(15) +def test_workflow_with_failing_mqtt_sink_and_fail_fast_raises( + enterprise_blocks_enabled, +): + # given + workflow_definition = copy.deepcopy(MQTT_SINK_WORKFLOW) + workflow_definition["steps"][0]["fail_fast"] = True + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when / then + with pytest.raises(Exception, match="not connected"): + execution_engine.run( + runtime_parameters={ + "host": "localhost", + "port": closed_port(), + "message": "frame payload", + } + ) + + +@pytest.mark.timeout(15) +def test_workflows_core_disable_sinks_still_reaches_the_enterprise_sink( + enterprise_blocks_enabled, +): + # given - the server supplies `disable_sinks` namespaced under the core + # block source (http_api.py, inference_pipeline.py). An enterprise block + # re-tagged with its plugin module name would never see it and compilation + # would fail with BlockInitParameterNotProvidedError. + workflow_definition = copy.deepcopy(MQTT_SINK_WORKFLOW) + workflow_definition["outputs"].append( + { + "type": "JsonField", + "name": "message", + "selector": "$steps.mqtt_sink.message", + } + ) + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters={"workflows_core.disable_sinks": True}, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when - nothing listens on the port, so a sink that tried to connect + # would report a failure instead of the disabled response + with patch.object(v2.mqtt, "Client") as mqtt_client: + result = execution_engine.run( + runtime_parameters={ + "host": "localhost", + "port": closed_port(), + "message": "frame payload", + } + ) + + # then - no connection attempted, disabled response returned + mqtt_client.assert_not_called() + assert result[0]["status"] is False + assert result[0]["message"] == disabled_sink_message( + disabled_by_execution_policy=True + ) diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 2adace3b5a..f584041b17 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -16,6 +16,17 @@ for contributor and maintainer responsibilities. ## Unreleased +### Added + +- MQTT Writer v2 (`roboflow_enterprise/mqtt_writer_sink@v2`): v1 plus an optional `fire_and_forget` (default `False`) and outage hardening; v1 is unchanged. + - With `fire_and_forget` a run queues the message and returns without waiting for the send or the broker's acknowledgement, and after the first run it never waits for a reconnect, so a disconnected broker no longer stalls each frame by up to `timeout`. If the first connection attempt fails, the client is kept and retries in the background instead of reconnecting synchronously on every frame. Success means queued, not delivered. + - While the broker is disconnected a QoS 0 message is dropped and a QoS 1/2 message is queued and sent after reconnect, both reported with `error_status` set; these repeating failures are logged when they start or change and once on recovery, not on every run. + - The step declares the `fire_and_forget_hides_persistence_failures` and `connection_and_state_rebuilt_per_request` restrictions when the switch is on (over the HTTP API every request builds a fresh block, so each request is a first run and a message still queued when the request ends is dropped), and reports a selector-driven switch as unknown. + - In either mode the client keeps at most 1000 QoS 1/2 messages awaiting acknowledgement and 1000 packets waiting to be sent (paho buffers both without limit by default); a publish beyond that is dropped and reported as a queue-full error. The limit counts messages, not bytes. + - The background reconnect delay follows `timeout` as in v1 but never drops below 0.1 s for the first retry and 1 s between later retries, so a tiny `timeout` cannot retry a dead broker in a busy loop. + - Selector-supplied `retain`, `fail_fast` and `fire_and_forget` are coerced the way the manifest validates them (for example `"false"` is False; v1 treats any non-empty string as True), and so are `port` and `qos` (for example `"1883.0"`, which v1 rejects); booleans are still rejected as a port or QoS. + - While the broker stays unreachable or keeps answering "unavailable", the background reconnect loop logs the first failure as an error and later attempts at debug level until a connection succeeds, instead of an error about once per second. + ## `0.2.4` Bundled execution engine: `1.16.1`. diff --git a/workflows/roboflow_workflows/enterprise_blocks/loader.py b/workflows/roboflow_workflows/enterprise_blocks/loader.py index a9142f0a45..b9aeb6b32f 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/loader.py +++ b/workflows/roboflow_workflows/enterprise_blocks/loader.py @@ -16,6 +16,9 @@ from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( MQTTWriterSinkBlockV1, ) +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + MQTTWriterSinkBlockV2, +) from roboflow_workflows.enterprise_blocks.sinks.opc_writer.v1 import ( OPCWriterSinkBlockV1, ) @@ -44,6 +47,7 @@ def load_enterprise_blocks() -> List[Type[WorkflowBlock]]: return [ OPCWriterSinkBlockV1, MQTTWriterSinkBlockV1, + MQTTWriterSinkBlockV2, MQTTReaderBlockV1, PLCBlockV1, PLCReaderBlockV1, diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py new file mode 100644 index 0000000000..303f0546fa --- /dev/null +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py @@ -0,0 +1,1003 @@ +import logging +import math +import threading +from typing import Any, List, Literal, Optional, Tuple, Type, Union + +import paho.mqtt.client as mqtt +from pydantic import ConfigDict, Field, TypeAdapter, ValidationError +from typing_extensions import Annotated + +# Bind by literal name to the shared "inference" logger tree so server +# handlers / filters / propagation still apply when installed, without +# importing the server logger module. +logger = logging.getLogger("inference") +from roboflow_workflows.core_steps.common.workload_presets import ( + FIRE_AND_FORGET_RESTRICTION, +) +from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + PERMANENT_CONNACK_CODES, + TRANSIENT_CONNACK_CODE, + ConfigurationError, + configure_tls, + connection_refused_message, + resolve_broker_address, +) +from roboflow_workflows.execution_engine.entities.base import OutputDefinition +from roboflow_workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + Selector, +) +from roboflow_workflows.execution_engine.entities.workload import ( + Discovery, + RuntimeRestriction, + WorkOperation, + incomplete_discovery, + unresolved_selector_problem, +) +from roboflow_workflows.prototypes.block import ( + BlockResult, + DependentResource, + Runtime, + RuntimeInputMode, + Severity, + WorkflowBlock, + WorkflowBlockManifest, + actual_restrictions_of, + is_workflow_selector, +) + +LONG_DESCRIPTION = """ +MQTT Writer block for publishing messages to an MQTT broker. + +Compared with v1, this version adds `fire_and_forget`, bounds the client's +outgoing queue, and logs a broker outage once instead of on every retry. + +The first run connects synchronously: the TCP connect and the broker's +session acknowledgement are each bounded by `timeout` (DNS resolution is +bounded by the OS resolver instead), so a cold start may take up to twice +`timeout` - raise it for remote brokers. Afterwards a background network +loop maintains the connection and owns reconnects. + +Without `fire_and_forget`, a run whose TCP connect fails reports the failure +and leaves nothing behind; the next run on the same block instance (video +pipelines) retries from scratch. Once the TCP connection is open the client +is kept: a broker that does not answer within `timeout` or answers +"unavailable" is retried by the background loop, which also re-establishes an +established connection that later drops. Over the HTTP API every request builds a fresh block +instance, so each request pays the bounded connect and a failed request is +final for that request. One block instance publishes to a single broker +connection: changing host, port, credentials, timeout or the TLS settings +between runs is rejected as a configuration error. + +The server operator may restrict which brokers the MQTT blocks connect to +with `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` +entries) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (when False, the +workflow's host and port are ignored and the first allowlist entry is used); +a run the policy forbids is reported like any other failure. + +Set `encryption` to `tls` to encrypt the connection and verify the broker's +certificate against the system trust store; the port is not switched +automatically, so set it to the broker's TLS port (usually 8883). A broker +signed by a private CA needs `ca_certificate_path`, a PEM bundle on the +machine running inference, which requires local file system access for +Workflow blocks (`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). +Certificate verification cannot be disabled. A peer that accepts the TCP +connection but never completes the TLS handshake is given up on after the +MQTT keepalive interval (60 s), not after `timeout`; real brokers close such +a connection immediately. + +A broker that refuses the connection (bad user name or password, not +authorised, an unacceptable protocol version) is reported in the outputs with +the broker's reason, and the block stops reconnecting: retrying with the same +credentials cannot succeed and would only trip the broker's authentication +rate limiting. Fix the configuration and restart the pipeline. A broker +answering "unavailable" is a passing condition, so the client keeps retrying +in the background. + +Set `fire_and_forget` to True to stop runs from waiting on the broker. A run +hands the message to the client and returns without waiting for the send or +the broker's acknowledgement, so success means "queued for publishing", not +"delivered". The first run still connects synchronously; afterwards a run +never waits for a reconnect. If that first connection attempt fails, the +client is kept and retries in the background, so a broker that is down when +the pipeline starts does not block every frame either. While the broker is +disconnected a QoS 0 message is dropped (paho does not queue QoS 0), and a +QoS 1/2 message is queued by the client and sent once the background loop +reconnects; both are reported with `error_status` set on every run, but logged +only when the failure starts or changes, and the recovery is logged once. In +either mode the client holds at most 1000 QoS 1/2 messages awaiting +acknowledgement and 1000 packets waiting to be sent; a publish beyond that is +dropped and reported. The limit counts messages, not bytes: with large +payloads a long outage can still hold up to 1000 of them in memory, so prefer +QoS 0 for large messages. Fire and forget is meant for InferencePipelines: over +the HTTP API every request builds a fresh block instance, so each request is a +first run, and a message still queued when the request ends is dropped. + +Outputs: + - error_status (bool): Indicates if an error occurred during the MQTT publishing process. + True if there was an error, False if successful. + - message (str): Status message describing the result of the operation. + Contains error details if error_status is True, + or success confirmation if error_status is False. + A publish acknowledgement timeout on QoS 1/2 is reported + as delivery-unknown (the message may still be delivered); + an unconfirmed QoS 0 send is reported as lost. With + fire_and_forget a successful run only confirms the + message was queued. + +By default failures are returned in the outputs and logged, and the workflow +keeps running. Set fail_fast to True to raise the failure instead, stopping +the workflow run - intended for one-shot requests, not streaming pipelines. +""" + + +TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} + +# the fire-and-forget client and its queue live in the block instance, which +# the HTTP API rebuilds for every request +FIRE_AND_FORGET_PER_REQUEST_RESTRICTION = RuntimeRestriction( + code="connection_and_state_rebuilt_per_request", + severity=Severity.SOFT, + note=( + "With fire_and_forget the broker connection and the client's message queue " + "live in this block instance. Over HTTP every request builds a fresh " + "instance, so each request still connects before publishing, and a QoS 1/2 " + "message queued while the broker is down is dropped when the request ends. " + "Use an InferencePipeline for fire-and-forget publishing." + ), + applies_to_runtimes=[ + Runtime.HOSTED_SERVERLESS, + Runtime.SELF_HOSTED_CPU, + Runtime.SELF_HOSTED_GPU, + Runtime.DEDICATED_DEPLOYMENT, + ], + applies_to_input_modes=[RuntimeInputMode.IMAGE], +) + + +BOOLEAN_ADAPTER = TypeAdapter(bool) +INTEGER_ADAPTER = TypeAdapter(int) + + +def _coerce_integer(value: Any) -> int: + """Coerce a selector-resolved integer the way the manifest would. + + Accepts what the manifest's validation accepts (``1883``, ``"1883"``, + ``"1883.0"``, ``1883.0``) except booleans, which are never a port or QoS. + + Args: + value: The literal or selector-resolved value. + + Returns: + The integer value. + + Raises: + ValidationError: When the value is a boolean or not an integral number. + """ + if isinstance(value, bool): + return INTEGER_ADAPTER.validate_python(value, strict=True) + return INTEGER_ADAPTER.validate_python(value) + + +def _coerce_switch(value: Any) -> bool: + """Coerce a selector-resolved switch the way the manifest would. + + The engine validates selector values on a copy of the manifest and hands + ``run()`` the raw value, so a ``"false"`` string would otherwise be truthy. + + Args: + value: The literal or selector-resolved value; ``None`` (an unset + input) means the default, False. + + Returns: + The boolean the manifest's validation accepts the value as. + + Raises: + ValidationError: When the value is not a boolean or a boolean-like + string or number. + """ + if value is None: + return False + return BOOLEAN_ADAPTER.validate_python(value) + + +def _not_connected_message(error: BaseException, encryption: str) -> str: + if encryption == "tls": + return ( + f"MQTT broker not connected ({error}). TLS is enabled: check that port is " + "the broker's TLS port and that ca_certificate_path matches the broker's " + "certificate authority." + ) + return ( + f"MQTT broker not connected ({error}). Raise 'timeout' if the broker needs " + "longer to connect." + ) + + +WRITER_CONNECTION_INPUTS = "username and password" +NOT_CONNECTED_WITHIN_TIMEOUT = ( + "MQTT broker not connected (connection was not established within timeout); " + "the client keeps retrying in the background. Raise 'timeout' if the " + "broker needs longer to connect." +) +# paho keeps QoS 1/2 messages until the broker acknowledges them and QoS 0 +# packets until the socket takes them, both by default without limit; a +# fire-and-forget pipeline publishing through a long outage or to a slow broker +# would otherwise grow the buffer for as long as that lasts +MAX_QUEUED_MESSAGES = 1000 +# the reconnect schedule follows `timeout` (half of it for the first retry, up +# to twice it between retries), bounded below so a tiny timeout cannot retry a +# dead broker in a busy loop +MIN_RECONNECT_DELAY_SECONDS = 0.1 +MIN_MAX_RECONNECT_DELAY_SECONDS = 1.0 +QUEUE_FULL = ( + f"MQTT client queue is full ({MAX_QUEUED_MESSAGES} messages awaiting " + "delivery); the message was dropped." +) +FIRE_AND_FORGET_QUEUED = ( + "Message queued for publishing; delivery is not confirmed (fire_and_forget)." +) +FIRE_AND_FORGET_QUEUED_UNTIL_RECONNECT = ( + "MQTT broker not connected; the message is queued by the client and sent " + "once the background loop reconnects." +) +FIRE_AND_FORGET_QOS0_DROPPED = ( + "MQTT broker not connected; QoS 0 messages are not queued, so the message " + "was dropped. The client keeps reconnecting in the background." +) + + +class MQTTWriterState: + """Connection state shared between paho's network thread and the engine thread. + + Passed to paho as ``userdata`` so the callbacks never hold a reference to the + block itself. + + Attributes: + connected: Set once the broker accepted the session; cleared on disconnect. + connack: Set on every CONNACK, accepted or refused, so a run can stop + waiting as soon as the broker answered; cleared on a transport + disconnect, kept while a refusal is recorded. + refused_code: The CONNACK return code of the last refusal, or None after an + accepted CONNACK. A permanent code (see ``PERMANENT_CONNACK_CODES``) means + the network loop was stopped and every run reports the refusal. + connect_failing: Set by a failed connection attempt or a refused CONNACK + and cleared by an accepted CONNACK, so an outage is logged once + rather than on every background retry. + """ + + def __init__(self): + self.connected = threading.Event() + self.connack = threading.Event() + self.refused_code: Optional[int] = None + self.connect_failing = False + + def reset(self) -> None: + self.connected.clear() + self.connack.clear() + self.refused_code = None + self.connect_failing = False + + +def mqtt_on_connect( + client, state: MQTTWriterState, flags, reason_code, properties=None +): + # paho invokes on_connect for accepted and rejected CONNACK alike; + # only reason_code 0 means an established MQTT session. The outcome is + # recorded before `connack` wakes a waiting run, so the run never reads a + # half-updated state + if reason_code == 0: + state.refused_code = None + state.connect_failing = False + state.connected.set() + state.connack.set() + logger.info("MQTT client connected") + return + state.refused_code = reason_code + state.connected.clear() + state.connack.set() + # a transient refusal comes back on every background retry, so only the + # first one of an outage is an error; a permanent one ends the loop + repeated = state.connect_failing and reason_code not in PERMANENT_CONNACK_CODES + state.connect_failing = True + (logger.debug if repeated else logger.error)( + "MQTT connection refused: %s (code %s)", + mqtt.connack_string(reason_code), + reason_code, + ) + if reason_code in PERMANENT_CONNACK_CODES: + # paho would otherwise reconnect with the same credentials forever; + # disconnect() puts it in the disconnecting state, so its loop ends + # after this CONNACK instead of retrying + client.disconnect() + + +def mqtt_on_connect_fail(client, state: MQTTWriterState): + # paho 1.6.1 invokes this callback with exactly (client, userdata). The + # background loop retries about once per reconnect delay while the broker + # is down, so only the first failure of an outage is an error + if state.connect_failing: + logger.debug("MQTT client still failing to connect to broker") + else: + logger.error("MQTT client failed to establish connection with broker") + state.connect_failing = True + state.connected.clear() + + +def mqtt_on_disconnect(client, state: MQTTWriterState, reason_code, properties=None): + # paho also ends every refused background retry here; only the loss of an + # established session is worth an info line + (logger.info if state.connected.is_set() else logger.debug)( + "MQTT client disconnected, result code %s", reason_code + ) + state.connected.clear() + if state.refused_code is None: + # a transport drop: the next run waits for the reconnect's CONNACK; after + # a refusal the answer stays visible so run() reports it without waiting + state.connack.clear() + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "MQTT Writer", + "version": "v2", + "short_description": "Publishes messages to an MQTT broker.", + "long_description": LONG_DESCRIPTION, + "license": "Roboflow Enterprise License", + "block_type": "sink", + "ui_manifest": { + "section": "industrial", + "icon": "fal fa-network-wired", + "blockPriority": 10, + "enterprise_only": True, + "local_only": True, + }, + } + ) + type: Literal["roboflow_enterprise/mqtt_writer_sink@v2"] + host: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Host of the MQTT broker.", + examples=["localhost", "$inputs.mqtt_host"], + ) + port: Union[ + Annotated[int, Field(ge=1, le=65535)], + Selector(kind=[INTEGER_KIND]), + ] = Field( + description="Port of the MQTT broker (1-65535).", + examples=[1883, "$inputs.mqtt_port"], + ) + topic: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="MQTT topic to publish the message to.", + examples=["sensors/temperature", "$inputs.mqtt_topic"], + ) + message: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Message to be published.", + examples=["Hello, MQTT!", "$inputs.mqtt_message"], + ) + qos: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=0, + description="Quality of Service level for the message.", + examples=[0, 1, 2], + ) + retain: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Whether the message should be retained by the broker.", + examples=[True, False], + ) + timeout: Union[ + Annotated[float, Field(gt=0, allow_inf_nan=False)], + Selector(kind=[FLOAT_KIND]), + ] = Field( + default=0.5, + description="Timeout for connecting to the MQTT broker and for sending MQTT messages. " + "Must be a finite number greater than 0.", + examples=[0.5], + ) + username: Union[Selector(kind=[STRING_KIND]), str] = Field( + default=None, + description="Username for MQTT broker authentication.", + examples=["$inputs.mqtt_username"], + ) + password: Union[Selector(kind=[STRING_KIND]), str] = Field( + default=None, + description="Password for MQTT broker authentication.", + examples=["$inputs.mqtt_password"], + ) + fail_fast: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="If True, MQTT failures raise an error stopping the workflow run " + "instead of being returned in the block outputs. Intended for one-shot " + "requests; in streaming pipelines it stops processing.", + examples=[False], + ) + fire_and_forget: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="If True, a run queues the message and returns without waiting " + "for the broker to receive or acknowledge it, and does not wait for a " + "reconnect after the first run. Success then means queued, not delivered. " + "While the broker is disconnected QoS 0 messages are dropped and QoS 1/2 " + "messages are sent after reconnect.", + examples=[False, "$inputs.mqtt_fire_and_forget"], + ) + encryption: Literal["none", "tls"] = Field( + default="none", + description="Transport security for the broker connection. `tls` encrypts the " + "connection and verifies the broker's certificate; set port to the broker's " + "TLS port (usually 8883).", + examples=["none", "tls"], + json_schema_extra={ + "values_metadata": { + "none": {"name": "None", "description": "Plain TCP (default)."}, + "tls": { + "name": "TLS", + "description": "Encrypted, broker certificate verified.", + }, + }, + }, + ) + ca_certificate_path: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Path to a PEM CA bundle on the machine running inference, for " + "brokers whose certificate issuer is not in the system trust store. Requires " + "local file system access for Workflow blocks " + "(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`); otherwise the run " + "reports an error. Used only when encryption is `tls`.", + examples=["/etc/ssl/certs/factory-ca.pem"], + json_schema_extra={"relevant_for": TLS_RELEVANT}, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + def discover_work_operations(self) -> List[WorkOperation]: + return [WorkOperation.EXTERNAL_REQUEST] + + def get_actual_restrictions( + self, *, ignore_environment_restrictions: bool = False + ) -> Discovery[RuntimeRestriction]: + # the delivery caveat exists only when the run does not wait for the broker + declared: Union[List[RuntimeRestriction], Discovery[RuntimeRestriction]] + if is_workflow_selector(self.fire_and_forget): + # a runtime value decides whether delivery is awaited, so the caveat + # MAY apply: declare nothing complete instead of guessing either way + declared = incomplete_discovery( + items=[], + reasons=[ + unresolved_selector_problem( + node_id=f"$steps.{getattr(self, 'name', '')}", + declaration="restrictions", + field="fire_and_forget", + selector=self.fire_and_forget, + ) + ], + ) + elif self.fire_and_forget: + declared = [ + FIRE_AND_FORGET_RESTRICTION, + FIRE_AND_FORGET_PER_REQUEST_RESTRICTION, + ] + else: + declared = [] + return actual_restrictions_of( + declared=declared, + node_id=f"$steps.{getattr(self, 'name', '')}", + ignore_environment_restrictions=ignore_environment_restrictions, + ) + + def discover_dependent_resources(self) -> List[DependentResource]: + return [] + + +class MQTTWriterSinkBlockV2(WorkflowBlock): + def __init__( + self, disable_sinks: bool = False, allow_access_to_file_system: bool = False + ): + # gates `ca_certificate_path`, a server-side path chosen by the workflow; + # False by default so a hand-constructed block is safe + self._allow_access_to_file_system = allow_access_to_file_system + self.mqtt_client: Optional[mqtt.Client] = None + self._connection = MQTTWriterState() + # the same Event object as the state's, kept under its historical name + self._connected = self._connection.connected + self._connection_identity: Optional[Tuple] = None + self._lifecycle_lock = threading.Lock() + self._disable_sinks = disable_sinks + # the fire-and-forget failure last logged; repeats are not logged again + self._fire_and_forget_failure: Optional[str] = None + + def close(self) -> None: + with self._lifecycle_lock: + client = self.mqtt_client + if client is None: + return + self.mqtt_client = None + self._connection_identity = None + self._fire_and_forget_failure = None + try: + client.disconnect() + except Exception as e: + logger.error("Failed to disconnect MQTT client: %s", e) + finally: + try: + # loop_stop() joins the network thread without a timeout; + # accepted so the thread never outlives the block + client.loop_stop() + finally: + # only after the join can no callback re-set the events; a + # client rebuilt by a later run starts without a stale refusal + self._connection.reset() + + def __del__(self): + try: + self.close() + except Exception as e: + logger.error("Failed to close MQTT client: %s", e) + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["disable_sinks", "allow_access_to_file_system"] + + def run( + self, + host: str, + port: int, + topic: str, + message: str, + username: Optional[str] = None, + password: Optional[str] = None, + qos: int = 0, + retain: bool = False, + timeout: float = 0.5, + fail_fast: bool = False, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, + fire_and_forget: bool = False, + ) -> BlockResult: + if self._disable_sinks: + return disabled_sink_response() + try: + fail_fast = _coerce_switch(fail_fast) + except ValidationError: + return self._handle_failure( + f"Invalid fail_fast: {fail_fast!r}. Must be a boolean.", + fail_fast=False, + ) + # selector-resolved values bypass manifest constraints, so validate here + try: + timeout_seconds = float(timeout) + except (TypeError, ValueError, OverflowError): + timeout_seconds = math.nan + if ( + not math.isfinite(timeout_seconds) + or not 0 < timeout_seconds <= threading.TIMEOUT_MAX + ): + return self._handle_failure( + f"Invalid timeout: {timeout!r}. Timeout must be a positive finite " + "number of seconds within the platform limit.", + fail_fast=fail_fast, + ) + timeout = timeout_seconds + # selector-supplied ports arrive uncoerced (numeric strings, floats + # from JSON), so coerce like the manifest would before range-checking; + # paho validates only port <= 0 and above 65535 the socket call raises + # OverflowError inside the network loop, silently killing it + try: + port_number = _coerce_integer(port) + except ValidationError: + port_number = 0 + if not 1 <= port_number <= 65535: + return self._handle_failure( + f"Invalid port: {port!r}. Port must be an integer between 1 and 65535.", + fail_fast=fail_fast, + ) + port = port_number + try: + qos_number = _coerce_integer(qos) + except ValidationError: + qos_number = -1 + if qos_number not in (0, 1, 2): + return self._handle_failure( + f"Invalid qos: {qos!r}. QoS must be 0, 1 or 2.", + fail_fast=fail_fast, + ) + qos = qos_number + if password is not None and username is None: + return self._handle_failure( + "Password provided without username. Set username to enable MQTT authentication.", + fail_fast=fail_fast, + ) + if encryption not in ("none", "tls"): + return self._handle_failure( + f"Invalid encryption: {encryption!r}. Must be 'none' or 'tls'.", + fail_fast=fail_fast, + ) + # an empty editor field is "unset"; anything but a string is a wiring error + if ca_certificate_path is not None and not isinstance(ca_certificate_path, str): + return self._handle_failure( + f"Invalid ca_certificate_path: {ca_certificate_path!r}. Must be a path.", + fail_fast=fail_fast, + ) + ca_certificate_path = ca_certificate_path or None + try: + retain = _coerce_switch(retain) + except ValidationError: + return self._handle_failure( + f"Invalid retain: {retain!r}. Must be a boolean.", fail_fast=fail_fast + ) + try: + fire_and_forget = _coerce_switch(fire_and_forget) + except ValidationError: + return self._handle_failure( + f"Invalid fire_and_forget: {fire_and_forget!r}. Must be a boolean.", + fail_fast=fail_fast, + ) + try: + # the operator's broker policy; the RESOLVED address is what the + # block connects to and what the connection identity is built from + host, port = resolve_broker_address( + host, port, log_override=self.mqtt_client is None + ) + except ConfigurationError as e: + return self._handle_failure(str(e), fail_fast=fail_fast) + with self._lifecycle_lock: + return self._connect_and_publish( + host=host, + port=port, + topic=topic, + message=message, + username=username, + password=password, + qos=qos, + retain=retain, + timeout=timeout, + fail_fast=fail_fast, + encryption=encryption, + ca_certificate_path=ca_certificate_path, + fire_and_forget=fire_and_forget, + ) + + def _connect_and_publish( + self, + host: str, + port: int, + topic: str, + message: str, + username: Optional[str], + password: Optional[str], + qos: int, + retain: bool, + timeout: float, + fail_fast: bool, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, + fire_and_forget: bool = False, + ) -> BlockResult: + connection_identity = ( + host, + port, + username, + password, + timeout, + encryption, + ca_certificate_path, + ) + created_client = self.mqtt_client is None + if created_client: + client = None + try: + client = mqtt.Client(userdata=self._connection) + client.max_queued_messages_set(MAX_QUEUED_MESSAGES) + if username is not None: + client.username_pw_set(username, password) + # TLS must be configured before connect(); the CA path is gated + # by the engine's file-system permission + configure_tls( + client, + use_tls=encryption == "tls", + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=self._allow_access_to_file_system, + ) + client.on_connect = mqtt_on_connect + client.on_connect_fail = mqtt_on_connect_fail + client.on_disconnect = mqtt_on_disconnect + # min_delay stays below the readiness wait (= timeout) so the + # first reconnect attempt after a connection drop can finish + # within a single run's wait instead of starting as it expires; + # the floors keep a tiny timeout from turning the background + # retries against a dead broker into a busy loop + client.reconnect_delay_set( + min_delay=max(timeout / 2, MIN_RECONNECT_DELAY_SECONDS), + max_delay=max(2 * timeout, MIN_MAX_RECONNECT_DELAY_SECONDS), + ) + # paho 1.6.1 has no public setter for its synchronous connect + # timeout; without this the TCP phase runs under paho's 5s + # default instead of the block's timeout + client._connect_timeout = timeout + # the TCP connect happens here, bounded by _connect_timeout + # (DNS resolution is not - it runs under the OS resolver + # timeout); the CONNACK wait below covers the handshake rest + client.connect(host, port) + client.loop_start() + except ConfigurationError as e: + # TLS setup refused before any socket was opened: nothing kept + return self._handle_failure(str(e), fail_fast=fail_fast) + except OSError as e: + error = _not_connected_message(e, encryption=encryption) + if fire_and_forget: + return self._keep_retrying_after_failed_connect( + client, + connection_identity=connection_identity, + error=error, + topic=topic, + message=message, + qos=qos, + retain=retain, + fail_fast=fail_fast, + ) + # broker unreachable: nothing is kept and no loop was started, + # so a failed one-shot run leaves no background thread behind; + # the next run on this instance retries with a fresh client + return self._handle_failure(error, fail_fast=fail_fast) + except Exception as e: + if client is not None: + try: + client.loop_stop() + except Exception: + pass + return self._handle_failure( + f"Failed to initialize MQTT client: {e}", fail_fast=fail_fast + ) + self.mqtt_client = client + self._connection_identity = connection_identity + elif connection_identity != self._connection_identity: + return self._handle_failure( + "MQTT connection parameters (host, port, credentials, timeout, encryption " + "or ca_certificate_path) changed between runs; this block publishes only " + "to the connection configured on its first run.", + fail_fast=fail_fast, + ) + state = self._connection + # a permanent refusal stopped the network loop: report it on every run + # without waiting and without touching the broker again + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), + fail_fast=fail_fast, + ) + # fire and forget waits for the broker only on the run that built the + # client; afterwards a disconnected run returns at once and leaves the + # reconnect to the background loop + waits_for_broker = created_client or not fire_and_forget + if not waits_for_broker and qos == 0 and not state.connected.is_set(): + return self._fire_and_forget_outage( + FIRE_AND_FORGET_QOS0_DROPPED, fail_fast=fail_fast + ) + # the background loop owns (re)connecting; runs only wait for readiness. + # The wait ends on any CONNACK, so a refusal is reported as soon as the + # broker answers instead of after the timeout + if waits_for_broker and not state.connected.is_set(): + not_ready = self._wait_until_connected(timeout) + if not_ready is not None: + # a permanent refusal stopped the loop: nothing would be sent + if not fire_and_forget or state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure(not_ready, fail_fast=fail_fast) + return self._report_not_connected( + not_ready, + topic=topic, + message=message, + qos=qos, + retain=retain, + fail_fast=fail_fast, + ) + connected = state.connected.is_set() + # max_queued_messages_set() bounds the QoS 1/2 messages, not the encoded + # packets paho buffers while the socket cannot drain them (a broker + # slower than the stream); 1.6.1 exposes that buffer only privately. + # Checked for every QoS, so runs alternating QoS cannot exceed it + if len(self.mqtt_client._out_packet) >= MAX_QUEUED_MESSAGES: + if fire_and_forget: + return self._fire_and_forget_outage(QUEUE_FULL, fail_fast=fail_fast) + return self._handle_failure(QUEUE_FULL, fail_fast=fail_fast) + try: + res: mqtt.MQTTMessageInfo = self.mqtt_client.publish( + topic, message, qos=qos, retain=retain + ) + except Exception as e: + return self._handle_failure( + f"Failed to publish message: {e}", fail_fast=fail_fast + ) + if fire_and_forget: + return self._fire_and_forget_result( + res.rc, qos=qos, connected=connected, fail_fast=fail_fast + ) + if res.rc == mqtt.MQTT_ERR_QUEUE_SIZE: + return self._handle_failure(QUEUE_FULL, fail_fast=fail_fast) + if res.rc == mqtt.MQTT_ERR_NO_CONN and qos > 0: + # paho keeps QoS 1/2 messages queued for delivery after reconnect + return self._handle_failure( + "Connection lost before publish; the message is queued by the client, " + "delivery status unknown and the message may still be delivered.", + fail_fast=fail_fast, + ) + try: + res.wait_for_publish(timeout=timeout) + published = res.is_published() + except Exception as e: + return self._handle_failure( + f"Failed to publish message: {e}", fail_fast=fail_fast + ) + if published: + return { + "error_status": False, + "message": "Message published successfully", + } + if qos == 0: + # an unconfirmed QoS 0 send is a real loss: paho never queues + # QoS 0 messages and reconnect() clears the pending out packet + return self._handle_failure( + "Publish confirmation timed out; the connection dropped before " + "the message was fully sent and QoS 0 messages are not " + "retransmitted, so the message is lost.", + fail_fast=fail_fast, + ) + return self._handle_failure( + "Publish acknowledgement timed out; delivery status unknown " + "and the message may still be delivered.", + fail_fast=fail_fast, + ) + + def _keep_retrying_after_failed_connect( + self, + client: mqtt.Client, + connection_identity: Tuple, + error: str, + topic: str, + message: str, + qos: int, + retain: bool, + fail_fast: bool, + ) -> BlockResult: + # keep the client and let paho's loop retry the connection in the + # background; otherwise every later run would block on a fresh + # synchronous connect for as long as the broker is down. This run + # logs the failure, so the loop's retries do not log it again + self._connection.connect_failing = True + try: + client.loop_start() + except Exception as e: + self._connection.connect_failing = False + try: + client.loop_stop() + except Exception: + pass + return self._handle_failure( + f"Failed to initialize MQTT client: {e}", fail_fast=fail_fast + ) + self.mqtt_client = client + self._connection_identity = connection_identity + return self._report_not_connected( + error, + topic=topic, + message=message, + qos=qos, + retain=retain, + fail_fast=fail_fast, + ) + + def _wait_until_connected(self, timeout: float) -> Optional[str]: + """Wait for the background loop to establish the session. + + Args: + timeout: Bound for the CONNACK wait and, after a transient refusal, + for the reconnect. + + Returns: + None once connected, otherwise the reason the broker is not ready. + """ + state = self._connection + if not state.connack.wait(timeout=timeout): + return NOT_CONNECTED_WITHIN_TIMEOUT + if state.refused_code in PERMANENT_CONNACK_CODES: + return connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ) + if not state.connected.wait(timeout=timeout): + if state.refused_code == TRANSIENT_CONNACK_CODE: + return connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ) + return NOT_CONNECTED_WITHIN_TIMEOUT + return None + + def _report_not_connected( + self, + error: str, + topic: str, + message: str, + qos: int, + retain: bool, + fail_fast: bool, + ) -> BlockResult: + # fire and forget on a run that found no session: QoS 0 has nothing to + # queue, QoS 1/2 is handed to paho for the background connection + if qos == 0: + return self._fire_and_forget_outage( + f"{error} QoS 0 messages are not queued, so the message was dropped.", + fail_fast=fail_fast, + ) + try: + res = self.mqtt_client.publish(topic, message, qos=qos, retain=retain) + except Exception as e: + return self._handle_failure( + f"Failed to publish message: {e}", fail_fast=fail_fast + ) + if res.rc == mqtt.MQTT_ERR_QUEUE_SIZE: + return self._fire_and_forget_outage(QUEUE_FULL, fail_fast=fail_fast) + return self._fire_and_forget_outage( + f"{error} The message is queued by the client and sent once it connects; " + "over the HTTP API it is dropped when the request ends.", + fail_fast=fail_fast, + ) + + def _fire_and_forget_result( + self, rc: int, qos: int, connected: bool, fail_fast: bool + ) -> BlockResult: + if rc == mqtt.MQTT_ERR_SUCCESS and (connected or qos == 0): + if self._fire_and_forget_failure is not None: + logger.info( + "MQTT Writer publishing again after: %s", + self._fire_and_forget_failure, + ) + self._fire_and_forget_failure = None + return {"error_status": False, "message": FIRE_AND_FORGET_QUEUED} + if rc == mqtt.MQTT_ERR_QUEUE_SIZE: + return self._fire_and_forget_outage(QUEUE_FULL, fail_fast=fail_fast) + if qos > 0 and rc in (mqtt.MQTT_ERR_SUCCESS, mqtt.MQTT_ERR_NO_CONN): + # paho keeps QoS 1/2 messages and sends them after the reconnect + return self._fire_and_forget_outage( + FIRE_AND_FORGET_QUEUED_UNTIL_RECONNECT, fail_fast=fail_fast + ) + if rc == mqtt.MQTT_ERR_NO_CONN: + # the connection dropped between the readiness check and publish() + return self._fire_and_forget_outage( + FIRE_AND_FORGET_QOS0_DROPPED, fail_fast=fail_fast + ) + return self._handle_failure( + f"Failed to publish message: {mqtt.error_string(rc)}", + fail_fast=fail_fast, + ) + + def _fire_and_forget_outage(self, message: str, fail_fast: bool) -> BlockResult: + # without the readiness wait a pipeline reaches this on every frame of + # an outage; log when the failure starts or changes, not once per run + log = message != self._fire_and_forget_failure + self._fire_and_forget_failure = message + return self._handle_failure(message, fail_fast=fail_fast, log=log) + + def _handle_failure( + self, message: str, fail_fast: bool, log: bool = True + ) -> BlockResult: + if log: + logger.error("MQTT Writer failure: %s", message) + if fail_fast: + raise RuntimeError(message) + return {"error_status": True, "message": message} diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_fire_and_forget_sink_workload_declarations.py b/workflows/tests/unit_tests/core_steps/sinks/test_fire_and_forget_sink_workload_declarations.py index 3544879b49..c21be083f6 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_fire_and_forget_sink_workload_declarations.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_fire_and_forget_sink_workload_declarations.py @@ -1,7 +1,7 @@ -"""Fire-and-forget declarations of the Microsoft SQL Server, Event Writer and -OPC UA Writer sinks. +"""Fire-and-forget declarations of the Microsoft SQL Server, Event Writer, +OPC UA Writer and MQTT Writer v2 sinks. -All three dispatch their write in the background when ``fire_and_forget`` is +All four return before the write is confirmed when ``fire_and_forget`` is true, so a failed write is not returned. Their actual declarations used to omit that caveat. They now declare it for a literal ``True``, omit it for a literal ``False``, and report a selector as unknown rather than guessing either way. @@ -32,6 +32,12 @@ from roboflow_workflows.enterprise_blocks.sinks.microsoft_sql_server.v1 import ( MicrosoftSQLServerSinkBlockV1, ) +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + FIRE_AND_FORGET_PER_REQUEST_RESTRICTION, +) +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + BlockManifest as MQTTWriterManifest, +) from roboflow_workflows.enterprise_blocks.sinks.opc_writer.v1 import ( BlockManifest as OPCWriterManifest, ) @@ -57,6 +63,7 @@ SQL_SERVER_TYPE = "roboflow_core/microsoft_sql_server_sink@v1" EVENT_WRITER_TYPE = "roboflow_enterprise/event_writer_sink@v1" OPC_WRITER_TYPE = "roboflow_enterprise/opc_writer_sink@v1" +MQTT_WRITER_TYPE = "roboflow_enterprise/mqtt_writer_sink@v2" FIRE_AND_FORGET_CODE = "fire_and_forget_hides_persistence_failures" COOLDOWN_CODE = "cooldown_timer_resets_on_stateless_http" SELECTOR = "$inputs.fire_and_forget" @@ -101,6 +108,18 @@ def _opc_writer_step(fire_and_forget: Any) -> Dict[str, Any]: } +def _mqtt_writer_step(fire_and_forget: Any) -> Dict[str, Any]: + return { + "type": MQTT_WRITER_TYPE, + "name": STEP_NAME, + "host": "broker.example.invalid", + "port": 1883, + "topic": "line1/status", + "message": "ok", + "fire_and_forget": fire_and_forget, + } + + # (manifest class, step builder, restrictions known independently of the switch) BLOCKS = { SQL_SERVER_TYPE: (SQLServerManifest, _sql_server_step, []), @@ -110,8 +129,13 @@ def _opc_writer_step(fire_and_forget: Any) -> Dict[str, Any]: _opc_writer_step, [COOLDOWN_ACTUAL_RESTRICTION], ), + MQTT_WRITER_TYPE: (MQTTWriterManifest, _mqtt_writer_step, []), } +# restrictions a literal ``True`` declares besides the shared caveat: the MQTT +# Writer's client and queue live in the block instance HTTP rebuilds per request +SWITCHED_EXTRAS = {MQTT_WRITER_TYPE: [FIRE_AND_FORGET_PER_REQUEST_RESTRICTION]} + def _manifest(block_type: str, fire_and_forget: Any) -> Any: manifest_class, build_step, _ = BLOCKS[block_type] @@ -125,6 +149,9 @@ def _expected_items(block_type: str, fire_and_forget: Any) -> List[Any]: expected = list(unconditional) if fire_and_forget is True: expected.append(FIRE_AND_FORGET_RESTRICTION) + expected.extend(SWITCHED_EXTRAS.get(block_type, [])) + # discoveries order their items by code + expected.sort(key=lambda item: item.code) return expected @@ -223,7 +250,9 @@ def test_the_host_view_matches_the_portable_view( assert host_view == portable_view -@pytest.mark.parametrize("manifest_class", [SQLServerManifest, EventWriterManifest]) +@pytest.mark.parametrize( + "manifest_class", [SQLServerManifest, EventWriterManifest, MQTTWriterManifest] +) def test_legacy_editor_restrictions_stay_unchanged(manifest_class: type) -> None: # the OPC UA Writer is pinned in test_industrial_workload_restrictions.py assert manifest_class.get_restrictions() == [] diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py new file mode 100644 index 0000000000..34b9709dfe --- /dev/null +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py @@ -0,0 +1,1635 @@ +import threading +import time +from typing import get_args +from unittest.mock import MagicMock, patch + +import paho.mqtt.client as mqtt +import pytest +from pydantic import ValidationError +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v2 +from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + FIRE_AND_FORGET_QOS0_DROPPED, + FIRE_AND_FORGET_QUEUED, + FIRE_AND_FORGET_QUEUED_UNTIL_RECONNECT, + MAX_QUEUED_MESSAGES, + NOT_CONNECTED_WITHIN_TIMEOUT, + QUEUE_FULL, + BlockManifest, + MQTTWriterSinkBlockV2, + MQTTWriterState, + mqtt_on_connect, + mqtt_on_connect_fail, + mqtt_on_disconnect, +) + +CLIENT_CLASS_PATH = ( + "roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2.mqtt.Client" +) + + +def run_kwargs(**overrides) -> dict: + kwargs = { + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "timeout": 0.01, + } + kwargs.update(overrides) + return kwargs + + +@pytest.fixture +def mock_client_cls(): + with patch(CLIENT_CLASS_PATH) as client_cls: + publish_result = client_cls.return_value.publish.return_value + publish_result.rc = mqtt.MQTT_ERR_SUCCESS + publish_result.is_published.return_value = True + yield client_cls + + +@pytest.fixture +def block() -> MQTTWriterSinkBlockV2: + return MQTTWriterSinkBlockV2() + + +class TestManifest: + def test_primary_identifier_is_namespaced(self): + identifiers = get_args(BlockManifest.model_fields["type"].annotation) + + assert identifiers[0] == "roboflow_enterprise/mqtt_writer_sink@v2" + + @pytest.mark.parametrize( + "identifier", + ["roboflow_enterprise/mqtt_writer_sink@v1", "mqtt_writer_sink@v1"], + ) + def test_v1_identifiers_belong_to_v1(self, identifier): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": identifier, + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + + @pytest.mark.parametrize( + "timeout", [0, -1, float("nan"), float("inf"), float("-inf")] + ) + def test_literal_timeout_must_be_finite_and_positive(self, timeout): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "timeout": timeout, + } + ) + + @pytest.mark.parametrize("timeout", [0.5, "$inputs.timeout"]) + def test_valid_timeouts_accepted(self, timeout): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "timeout": timeout, + } + ) + + assert manifest.timeout == timeout + + @pytest.mark.parametrize("port", [0, -1, 65536]) + def test_literal_port_must_be_within_tcp_range(self, port): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": port, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + + @pytest.mark.parametrize("port", [1883, "$inputs.port"]) + def test_valid_ports_accepted(self, port): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": port, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + + assert manifest.port == port + + +class TestCallbacks: + def test_on_connect_sets_event_only_for_accepted_connack(self): + state = MQTTWriterState() + + mqtt_on_connect(MagicMock(), state, {}, 0) + + assert state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code is None + + @pytest.mark.parametrize("reason_code", [1, 2, 3, 4, 5]) + def test_on_connect_clears_event_for_rejected_connack(self, reason_code): + state = MQTTWriterState() + state.connected.set() + + mqtt_on_connect(MagicMock(), state, {}, reason_code) + + assert not state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code == reason_code + + @pytest.mark.parametrize("reason_code", [1, 2, 4, 5]) + def test_on_connect_stops_the_loop_for_a_permanent_refusal(self, reason_code): + client = MagicMock() + + mqtt_on_connect(client, MQTTWriterState(), {}, reason_code) + + client.disconnect.assert_called_once_with() + + def test_on_connect_keeps_the_loop_for_broker_unavailable(self): + client = MagicMock() + state = MQTTWriterState() + + mqtt_on_connect(client, state, {}, 3) + + client.disconnect.assert_not_called() + assert state.refused_code == 3 + + @pytest.mark.parametrize("reason_code", [0, 3, 5]) + def test_on_connect_completes_the_state_before_waking_a_waiter(self, reason_code): + # a run wakes on `connack`; what it then reads must already be final + state = MQTTWriterState() + seen = {} + + class RecordingEvent(threading.Event): + def set(self): + seen["refused_code"] = state.refused_code + seen["connected"] = state.connected.is_set() + super().set() + + state.connack = RecordingEvent() + + mqtt_on_connect(MagicMock(), state, {}, reason_code) + + assert seen == { + "refused_code": reason_code or None, + "connected": reason_code == 0, + } + + def test_on_connect_logs_the_refusal_reason(self, caplog): + with caplog.at_level("ERROR", logger="inference"): + mqtt_on_connect(MagicMock(), MQTTWriterState(), {}, 4) + + assert "bad user name or password" in caplog.text + assert "code 4" in caplog.text + + def test_accepted_connack_after_broker_unavailable_resets_refusal(self): + state = MQTTWriterState() + + mqtt_on_connect(MagicMock(), state, {}, 3) + mqtt_on_connect(MagicMock(), state, {}, 0) + + assert state.refused_code is None + assert state.connected.is_set() + + def test_on_connect_fail_matches_paho_two_argument_signature(self): + state = MQTTWriterState() + state.connected.set() + + # paho 1.6.1 invokes on_connect_fail with exactly (client, userdata) + mqtt_on_connect_fail(MagicMock(), state) + + assert not state.connected.is_set() + + def test_on_connect_fail_logs_an_outage_once(self): + # the background loop retries about once per reconnect delay + state = MQTTWriterState() + + with patch.object(v2, "logger") as mock_logger: + for _ in range(5): + mqtt_on_connect_fail(MagicMock(), state) + assert mock_logger.error.call_count == 1 + + mqtt_on_connect(MagicMock(), state, {}, 0) + mqtt_on_connect_fail(MagicMock(), state) + + assert mock_logger.error.call_count == 2 + + def test_repeated_broker_unavailable_logged_once_per_outage(self): + # code 3 keeps paho reconnecting, so the refusal repeats every retry + state = MQTTWriterState() + + with patch.object(v2, "logger") as mock_logger: + for _ in range(5): + mqtt_on_connect(MagicMock(), state, {}, 3) + assert mock_logger.error.call_count == 1 + + mqtt_on_connect(MagicMock(), state, {}, 0) + mqtt_on_connect(MagicMock(), state, {}, 3) + + assert mock_logger.error.call_count == 2 + + def test_refused_retries_do_not_log_a_disconnect_each_time(self): + # paho ends each refused CONNACK with on_disconnect before retrying + state = MQTTWriterState() + + with patch.object(v2, "logger") as mock_logger: + for _ in range(5): + mqtt_on_connect(MagicMock(), state, {}, 3) + mqtt_on_disconnect(MagicMock(), state, 5) + + assert mock_logger.error.call_count == 1 + mock_logger.info.assert_not_called() + + def test_losing_an_established_session_logs_info(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 0) + + with patch.object(v2, "logger") as mock_logger: + mqtt_on_disconnect(MagicMock(), state, 7) + + mock_logger.info.assert_called_once() + + def test_permanent_refusal_logged_after_transient_ones(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 3) + + with patch.object(v2, "logger") as mock_logger: + mqtt_on_connect(MagicMock(), state, {}, 5) + + mock_logger.error.assert_called_once() + + def test_on_disconnect_clears_event(self): + state = MQTTWriterState() + state.connected.set() + + mqtt_on_disconnect(MagicMock(), state, 1) + + assert not state.connected.is_set() + + def test_on_disconnect_after_transport_drop_clears_connack(self): + state = MQTTWriterState() + state.connack.set() + + mqtt_on_disconnect(MagicMock(), state, 1) + + assert not state.connack.is_set() + + def test_on_disconnect_after_refusal_keeps_connack(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) + + mqtt_on_disconnect(MagicMock(), state, 0) + + assert state.connack.is_set() + assert state.refused_code == 5 + + def test_on_disconnect_accepts_mqtt5_properties_argument(self): + state = MQTTWriterState() + state.connected.set() + + mqtt_on_disconnect(MagicMock(), state, 1, properties=None) + + assert not state.connected.is_set() + + def test_reset_clears_everything(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) + mqtt_on_connect_fail(MagicMock(), state) + + state.reset() + + assert not state.connected.is_set() + assert not state.connack.is_set() + assert state.refused_code is None + assert state.connect_failing is False + + +def _answer_connect_with(mock_client_cls, reason_code: int) -> None: + """Make the mocked client's loop_start deliver a CONNACK with `reason_code` + to the state the block handed paho as userdata.""" + mock_client = mock_client_cls.return_value + + def fire_connack(): + state = mock_client_cls.call_args.kwargs["userdata"] + mqtt_on_connect(mock_client, state, {}, reason_code) + + mock_client.loop_start.side_effect = fire_connack + + +class TestRefusedConnection: + def test_refusal_reported_with_reason_and_inputs(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 5) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "code 5" in result["message"] + assert "Check username and password" in result["message"] + assert "Raise 'timeout'" not in result["message"] + mock_client_cls.return_value.disconnect.assert_called_once_with() + mock_client_cls.return_value.publish.assert_not_called() + + @pytest.mark.parametrize( + "reason_code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + ], + ) + def test_other_permanent_refusals_name_their_reason( + self, mock_client_cls, block, reason_code, reason + ): + _answer_connect_with(mock_client_cls, reason_code) + + result = block.run(**run_kwargs()) + + assert reason in result["message"] + mock_client_cls.return_value.disconnect.assert_called_once_with() + + def test_refusal_is_reported_without_waiting_for_timeout( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert result["error_status"] is True + assert time.monotonic() - started < 0.5 + + def test_next_run_after_refusal_repeats_it_without_a_new_connection( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + first = block.run(**run_kwargs(timeout=2.0)) + + started = time.monotonic() + second = block.run(**run_kwargs(timeout=2.0)) + + assert second == first + assert time.monotonic() - started < 0.5 + assert mock_client_cls.call_count == 1 + mock_client_cls.return_value.connect.assert_called_once() + + def test_refusal_with_fail_fast_raises_with_the_reason( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + with pytest.raises(Exception, match="not authorised"): + block.run(**run_kwargs(fail_fast=True)) + + def test_broker_unavailable_keeps_retrying_and_names_the_reason( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 3) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "broker unavailable" in result["message"] + assert "retrying in the background" in result["message"] + mock_client_cls.return_value.disconnect.assert_not_called() + + def test_run_succeeds_once_an_unavailable_broker_accepts( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 3) + block.run(**run_kwargs()) + state = mock_client_cls.call_args.kwargs["userdata"] + mqtt_on_connect(mock_client_cls.return_value, state, {}, 0) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert state.refused_code is None + + def test_no_connack_at_all_reports_not_connected(self, mock_client_cls, block): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert result["message"] == NOT_CONNECTED_WITHIN_TIMEOUT + + def test_close_after_refusal_completes_and_resets_state( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + block.run(**run_kwargs()) + + block.close() + + assert block.mqtt_client is None + assert block._connection.refused_code is None + assert not block._connection.connack.is_set() + mock_client_cls.return_value.loop_stop.assert_called_once() + + +class TestRunValidation: + @pytest.mark.parametrize( + "timeout", + [0, -1, float("nan"), float("inf"), float("-inf"), 10**400, 1e308], + ) + def test_invalid_timeout_rejected_before_client_construction( + self, mock_client_cls, block, timeout + ): + result = block.run(**run_kwargs(timeout=timeout)) + + assert result["error_status"] is True + assert "timeout" in result["message"].lower() + mock_client_cls.assert_not_called() + + @pytest.mark.parametrize( + "port", [0, -1, 65536, "abc", "1883.5", 1883.7, True, None] + ) + def test_invalid_port_rejected_before_client_construction( + self, mock_client_cls, block, port + ): + result = block.run(**run_kwargs(port=port)) + + assert result["error_status"] is True + assert "port" in result["message"].lower() + mock_client_cls.assert_not_called() + + @pytest.mark.parametrize("port", ["1883", 1883.0, "1883.0", " 1883 "]) + def test_selector_supplied_port_coerced_like_manifest( + self, mock_client_cls, block, port + ): + # runtime validation coerces on a discarded manifest copy, so run() + # receives the raw selector value and must coerce it itself + block._connected.set() + + result = block.run(**run_kwargs(port=port)) + + assert result["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) + + @pytest.mark.parametrize("qos", [-1, 3, "abc", 1.5, "1.5", True, False, None]) + def test_invalid_qos_rejected_before_client_construction( + self, mock_client_cls, block, qos + ): + result = block.run(**run_kwargs(qos=qos)) + + assert result["error_status"] is True + assert "qos" in result["message"].lower() + mock_client_cls.assert_not_called() + + @pytest.mark.parametrize("qos, expected", [("1", 1), (2.0, 2), (0, 0), ("1.0", 1)]) + def test_selector_supplied_qos_coerced_like_other_parameters( + self, mock_client_cls, block, qos, expected + ): + block._connected.set() + + result = block.run(**run_kwargs(qos=qos)) + + assert result["error_status"] is False + mock_client_cls.return_value.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=expected, retain=False + ) + + def test_password_without_username_rejected(self, mock_client_cls, block): + result = block.run(**run_kwargs(password="secret")) + + assert result["error_status"] is True + assert "username" in result["message"].lower() + mock_client_cls.assert_not_called() + + def test_username_without_password_configures_authentication( + self, mock_client_cls, block + ): + block._connected.set() + + block.run(**run_kwargs(username="lenny")) + + mock_client_cls.return_value.username_pw_set.assert_called_once_with( + "lenny", None + ) + + def test_username_with_empty_password_configures_authentication( + self, mock_client_cls, block + ): + block._connected.set() + + block.run(**run_kwargs(username="lenny", password="")) + + mock_client_cls.return_value.username_pw_set.assert_called_once_with( + "lenny", "" + ) + + +class TestClientSetup: + def test_client_connects_synchronously_before_background_loop( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + mock_client_cls.assert_called_once_with(userdata=block._connection) + mock_client.connect.assert_called_once_with("localhost", 1883) + mock_client.connect_async.assert_not_called() + called_methods = [call[0] for call in mock_client.method_calls] + assert called_methods.index("connect") < called_methods.index("loop_start") + + def test_client_registers_static_callbacks(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + + block.run(**run_kwargs()) + + assert mock_client.on_connect is mqtt_on_connect + assert mock_client.on_connect_fail is mqtt_on_connect_fail + assert mock_client.on_disconnect is mqtt_on_disconnect + for callback in ( + mock_client.on_connect, + mock_client.on_connect_fail, + mock_client.on_disconnect, + ): + assert getattr(callback, "__self__", None) is not block + + def test_reconnect_delay_lets_first_retry_finish_within_one_run( + self, mock_client_cls, block + ): + # min_delay must be below the readiness wait (= timeout), otherwise + # the first reconnect attempt after a connection drop starts exactly + # as the waiting run's timeout expires and that run always fails + block._connected.set() + + block.run(**run_kwargs(timeout=0.5)) + + mock_client_cls.return_value.reconnect_delay_set.assert_called_once_with( + min_delay=0.25, max_delay=1.0 + ) + + def test_reconnect_delay_follows_a_long_timeout(self, mock_client_cls, block): + block._connected.set() + + block.run(**run_kwargs(timeout=4.0)) + + mock_client_cls.return_value.reconnect_delay_set.assert_called_once_with( + min_delay=2.0, max_delay=8.0 + ) + + def test_tiny_timeout_cannot_make_reconnects_a_busy_loop( + self, mock_client_cls, block + ): + # fire and forget keeps a client retrying after a failed first connect + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + + block.run(**run_kwargs(timeout=1e-6, fire_and_forget=True)) + + mock_client.reconnect_delay_set.assert_called_once_with( + min_delay=0.1, max_delay=1.0 + ) + + def test_paho_connect_timeout_bound_to_block_timeout(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + connect_timeout_at_connect_call = [] + mock_client.connect.side_effect = lambda *args, **kwargs: ( + connect_timeout_at_connect_call.append(mock_client._connect_timeout) + ) + + block.run(**run_kwargs(timeout=0.25)) + + assert connect_timeout_at_connect_call == [0.25] + + def test_unreachable_broker_on_first_run_leaves_no_background_thread( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + + first_result = block.run(**run_kwargs()) + + mock_client.connect.assert_called_once() + assert first_result["error_status"] is True + assert "not connected" in first_result["message"].lower() + # no client kept and no loop started: a failed one-shot run must not + # leave a reconnecting background thread behind + assert block.mqtt_client is None + mock_client.loop_start.assert_not_called() + mock_client.publish.assert_not_called() + + # broker becomes reachable, the next run retries with a fresh client + mock_client.connect.side_effect = None + block._connected.set() + + second_result = block.run(**run_kwargs()) + + assert second_result["error_status"] is False + assert mock_client_cls.call_count == 2 + mock_client.loop_start.assert_called_once() + + def test_unreachable_broker_with_fail_fast_raises(self, mock_client_cls, block): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = OSError("no route to host") + + with pytest.raises(RuntimeError, match="not connected"): + block.run(**run_kwargs(fail_fast=True)) + + mock_client.connect.assert_called_once() + mock_client.loop_start.assert_not_called() + assert block.mqtt_client is None + + def test_setup_failure_resets_client_so_next_run_can_retry( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + mock_client.loop_start.side_effect = Exception("boom") + + first_result = block.run(**run_kwargs()) + + assert first_result["error_status"] is True + assert block.mqtt_client is None + mock_client.loop_stop.assert_called_once() + + mock_client.loop_start.side_effect = None + block._connected.set() + + second_result = block.run(**run_kwargs()) + + assert second_result["error_status"] is False + assert mock_client_cls.call_count == 2 + + def test_first_run_connection_timeout_does_not_poison_client( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + + first_result = block.run(**run_kwargs()) + + assert first_result["error_status"] is True + assert "not connected" in first_result["message"].lower() + mock_client.publish.assert_not_called() + + # simulates the background loop establishing the connection later + block._connected.set() + + second_result = block.run(**run_kwargs()) + + assert second_result["error_status"] is False + mock_client_cls.assert_called_once() + mock_client.connect.assert_called_once() + mock_client.publish.assert_called_once() + + +class TestConnectionOwnership: + def test_disconnected_run_never_calls_manual_reconnect( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs()) + + block._connected.clear() + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not connected" in result["message"].lower() + mock_client.reconnect.assert_not_called() + mock_client.publish.assert_called_once() + + def test_changed_connection_parameters_rejected(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs(host="broker-a")) + + result = block.run(**run_kwargs(host="broker-b")) + + assert result["error_status"] is True + assert "parameters" in result["message"].lower() + mock_client.connect.assert_called_once_with("broker-a", 1883) + mock_client.publish.assert_called_once() + + def test_changed_credentials_rejected(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs(username="lenny", password="old")) + + result = block.run(**run_kwargs(username="lenny", password="new")) + + assert result["error_status"] is True + mock_client.publish.assert_called_once() + + def test_changed_timeout_rejected(self, mock_client_cls, block): + # the first timeout permanently configures the reconnect schedule + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs(timeout=0.5)) + + result = block.run(**run_kwargs(timeout=1.0)) + + assert result["error_status"] is True + assert "parameters" in result["message"].lower() + mock_client.publish.assert_called_once() + + +class TestPublishing: + def test_successful_publish(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run(**run_kwargs(qos=1, retain=True)) + + assert result["error_status"] is False + assert result["message"] == "Message published successfully" + mock_client.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=1, retain=True + ) + + def test_qos0_publish_confirmation_timeout_reported_as_lost( + self, mock_client_cls, block + ): + # paho never retransmits QoS 0: reconnect() clears _out_packet and + # QoS 0 messages are not queued, so an unconfirmed send is a real loss + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client.publish.return_value.is_published.return_value = False + + with patch.object(v2, "logger") as mock_logger: + result = block.run(**run_kwargs(qos=0)) + + assert result["error_status"] is True + assert "lost" in result["message"].lower() + assert "delivery status unknown" not in result["message"].lower() + assert mock_logger.error.called + + def test_qos1_publish_ack_timeout_reported_as_delivery_unknown( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client.publish.return_value.is_published.return_value = False + + with patch.object(v2, "logger") as mock_logger: + result = block.run(**run_kwargs(qos=1)) + + assert result["error_status"] is True + assert "delivery status unknown" in result["message"].lower() + assert mock_logger.error.called + + def test_qos1_publish_with_lost_connection_reported_as_delivery_unknown( + self, mock_client_cls, block + ): + # paho queues QoS 1/2 messages for redelivery when publish() returns + # MQTT_ERR_NO_CONN - not a final failure + import paho.mqtt.client as mqtt + + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client.publish.return_value.rc = mqtt.MQTT_ERR_NO_CONN + + result = block.run(**run_kwargs(qos=1)) + + assert result["error_status"] is True + assert "delivery status unknown" in result["message"].lower() + mock_client.publish.return_value.wait_for_publish.assert_not_called() + + def test_qos0_publish_with_lost_connection_reported_as_final_failure( + self, mock_client_cls, block + ): + # paho drops QoS 0 messages on MQTT_ERR_NO_CONN - final failure is honest + import paho.mqtt.client as mqtt + + block._connected.set() + mock_client = mock_client_cls.return_value + publish_result = mock_client.publish.return_value + publish_result.rc = mqtt.MQTT_ERR_NO_CONN + publish_result.wait_for_publish.side_effect = RuntimeError( + "Message publish failed: The client is not currently connected." + ) + + result = block.run(**run_kwargs(qos=0)) + + assert result["error_status"] is True + assert "failed to publish" in result["message"].lower() + + def test_publish_exception_returned_as_error(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client.publish.side_effect = ValueError("Invalid topic") + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "Invalid topic" in result["message"] + + def test_no_raw_prints_in_any_code_path(self, mock_client_cls, block): + with patch("builtins.print") as mock_print: + block._connected.set() + block.run(**run_kwargs()) + block._connected.clear() + block.run(**run_kwargs()) + + mock_print.assert_not_called() + + +class TestFailFast: + def test_fail_fast_raises_on_connection_failure(self, mock_client_cls, block): + with pytest.raises(Exception, match="not connected"): + block.run(**run_kwargs(fail_fast=True)) + + def test_fail_fast_raises_on_invalid_timeout(self, mock_client_cls, block): + with pytest.raises(Exception, match="[Tt]imeout"): + block.run(**run_kwargs(timeout=-1, fail_fast=True)) + + def test_fail_fast_does_not_affect_success(self, mock_client_cls, block): + block._connected.set() + + result = block.run(**run_kwargs(fail_fast=True)) + + assert result["error_status"] is False + + def test_fail_fast_ack_timeout_raises_once_with_clean_message( + self, mock_client_cls, block + ): + block._connected.set() + mock_client_cls.return_value.publish.return_value.is_published.return_value = ( + False + ) + + with patch.object(v2, "logger") as mock_logger: + with pytest.raises(RuntimeError, match="^Publish confirmation"): + block.run(**run_kwargs(fail_fast=True)) + + assert mock_logger.error.call_count == 1 + + +class TestQueueLimit: + def test_client_queue_is_bounded(self, mock_client_cls, block): + block._connected.set() + + block.run(**run_kwargs()) + + mock_client_cls.return_value.max_queued_messages_set.assert_called_once_with( + MAX_QUEUED_MESSAGES + ) + + @pytest.mark.parametrize("fire_and_forget", [False, True]) + def test_full_queue_reported_as_dropped( + self, mock_client_cls, block, fire_and_forget + ): + block._connected.set() + publish_result = mock_client_cls.return_value.publish.return_value + publish_result.rc = mqtt.MQTT_ERR_QUEUE_SIZE + + result = block.run(**run_kwargs(qos=1, fire_and_forget=fire_and_forget)) + + assert result == {"error_status": True, "message": QUEUE_FULL} + publish_result.wait_for_publish.assert_not_called() + + @pytest.mark.parametrize("qos", [0, 1, 2]) + @pytest.mark.parametrize("fire_and_forget", [False, True]) + def test_backlog_bounded_by_unsent_packets( + self, mock_client_cls, block, fire_and_forget, qos + ): + # max_queued_messages_set() bounds QoS 1/2 messages, not the packets + # paho buffers for a slow socket; a QoS 1/2 run after a QoS 0 backlog + # must not grow that buffer past the bound either + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs(qos=0, fire_and_forget=fire_and_forget)) + mock_client._out_packet = [object()] * MAX_QUEUED_MESSAGES + mock_client.publish.reset_mock() + + result = block.run(**run_kwargs(qos=qos, fire_and_forget=fire_and_forget)) + + assert result == {"error_status": True, "message": QUEUE_FULL} + mock_client.publish.assert_not_called() + + def test_qos0_publishes_below_the_packet_bound(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs(qos=0)) + mock_client._out_packet = [object()] * (MAX_QUEUED_MESSAGES - 1) + + result = block.run(**run_kwargs(qos=0, fire_and_forget=True)) + + assert result["error_status"] is False + assert mock_client.publish.call_count == 2 + + +def _connect_then_drop(mock_client_cls, block, **overrides) -> MagicMock: + """Run once against a connected broker, then simulate a dropped connection + the background loop has not re-established yet.""" + block._connected.set() + first = block.run(**run_kwargs(fire_and_forget=True, **overrides)) + assert first["error_status"] is False + block._connected.clear() + mock_client = mock_client_cls.return_value + mock_client.publish.reset_mock() + return mock_client + + +class TestFireAndForget: + def test_manifest_defaults_to_waiting(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "t", + "message": "m", + } + ) + + assert manifest.fire_and_forget is False + + def test_manifest_accepts_selector(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "t", + "message": "m", + "fire_and_forget": "$inputs.fire_and_forget", + } + ) + + assert manifest.fire_and_forget == "$inputs.fire_and_forget" + + def test_default_run_still_waits_for_confirmation(self, mock_client_cls, block): + block._connected.set() + + result = block.run(**run_kwargs(qos=1)) + + assert result["message"] == "Message published successfully" + mock_client_cls.return_value.publish.return_value.wait_for_publish.assert_called_once() + + @pytest.mark.parametrize("qos", [0, 1, 2]) + def test_publish_returns_without_waiting_for_confirmation( + self, mock_client_cls, block, qos + ): + block._connected.set() + mock_client = mock_client_cls.return_value + # an unconfirmed message must not matter: nothing waits for it + mock_client.publish.return_value.is_published.return_value = False + + result = block.run(**run_kwargs(qos=qos, retain=True, fire_and_forget=True)) + + assert result == {"error_status": False, "message": FIRE_AND_FORGET_QUEUED} + mock_client.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=qos, retain=True + ) + mock_client.publish.return_value.wait_for_publish.assert_not_called() + mock_client.publish.return_value.is_published.assert_not_called() + + def test_first_run_still_waits_for_the_broker(self, mock_client_cls, block): + started = time.monotonic() + result = block.run(**run_kwargs(timeout=0.2, fire_and_forget=True)) + + assert time.monotonic() - started >= 0.2 + assert result["error_status"] is True + assert result["message"].startswith(NOT_CONNECTED_WITHIN_TIMEOUT) + + def test_first_run_publishes_once_the_broker_accepts(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 0) + + result = block.run(**run_kwargs(fire_and_forget=True)) + + assert result == {"error_status": False, "message": FIRE_AND_FORGET_QUEUED} + + def test_unreachable_broker_on_first_run_keeps_client_retrying( + self, mock_client_cls, block + ): + # without the kept client every later frame would pay a fresh + # synchronous connect while the broker is down + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + + first = block.run(**run_kwargs(timeout=2.0, fire_and_forget=True)) + + assert first["error_status"] is True + assert "not connected" in first["message"].lower() + assert block.mqtt_client is mock_client + mock_client.loop_start.assert_called_once() + + started = time.monotonic() + second = block.run(**run_kwargs(timeout=2.0, fire_and_forget=True)) + + assert time.monotonic() - started < 0.5 + assert second == {"error_status": True, "message": FIRE_AND_FORGET_QOS0_DROPPED} + mock_client_cls.assert_called_once() + mock_client.connect.assert_called_once() + + @pytest.mark.parametrize("qos", [1, 2]) + def test_unreachable_broker_on_first_run_queues_qos_1_2( + self, mock_client_cls, block, qos + ): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + + result = block.run(**run_kwargs(qos=qos, retain=True, fire_and_forget=True)) + + assert result["error_status"] is True + assert "refused" in result["message"] + assert "queued by the client" in result["message"] + mock_client.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=qos, retain=True + ) + + def test_unreachable_broker_on_first_run_drops_qos0(self, mock_client_cls, block): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + + result = block.run(**run_kwargs(qos=0, fire_and_forget=True)) + + assert result["error_status"] is True + assert "refused" in result["message"] + assert "dropped" in result["message"] + mock_client.publish.assert_not_called() + + def test_first_run_loop_start_failure_reported_and_nothing_kept( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + mock_client.loop_start.side_effect = RuntimeError("no thread") + + result = block.run(**run_kwargs(fire_and_forget=True)) + + assert result["error_status"] is True + assert "Failed to initialize MQTT client: no thread" in result["message"] + assert block.mqtt_client is None + assert block._connection.connect_failing is False + mock_client.loop_stop.assert_called_once() + + @pytest.mark.parametrize("qos", [1, 2]) + def test_first_run_without_connack_queues_qos_1_2( + self, mock_client_cls, block, qos + ): + # TCP connected but the broker never answered within timeout + mock_client = mock_client_cls.return_value + + result = block.run(**run_kwargs(qos=qos, fire_and_forget=True)) + + assert result["error_status"] is True + assert result["message"].startswith(NOT_CONNECTED_WITHIN_TIMEOUT) + assert "queued by the client" in result["message"] + mock_client.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=qos, retain=False + ) + + def test_first_run_without_connack_drops_qos0(self, mock_client_cls, block): + result = block.run(**run_kwargs(qos=0, fire_and_forget=True)) + + assert result["error_status"] is True + assert result["message"].startswith(NOT_CONNECTED_WITHIN_TIMEOUT) + assert "dropped" in result["message"] + mock_client_cls.return_value.publish.assert_not_called() + + def test_first_run_broker_unavailable_queues_qos1(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 3) + + result = block.run(**run_kwargs(qos=1, fire_and_forget=True)) + + assert result["error_status"] is True + assert "broker unavailable" in result["message"] + assert "queued by the client" in result["message"] + mock_client_cls.return_value.publish.assert_called_once() + + def test_first_run_permanent_refusal_does_not_queue(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 5) + + result = block.run(**run_kwargs(qos=1, fire_and_forget=True)) + + assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "queued" not in result["message"] + mock_client_cls.return_value.publish.assert_not_called() + + def test_first_run_queue_full_while_not_connected(self, mock_client_cls, block): + mock_client = mock_client_cls.return_value + mock_client.publish.return_value.rc = mqtt.MQTT_ERR_QUEUE_SIZE + + result = block.run(**run_kwargs(qos=1, fire_and_forget=True)) + + assert result == {"error_status": True, "message": QUEUE_FULL} + + def test_failed_first_connect_is_not_logged_again_by_the_retry_loop( + self, mock_client_cls, block + ): + mock_client_cls.return_value.connect.side_effect = ConnectionRefusedError( + "refused" + ) + block.run(**run_kwargs(fire_and_forget=True)) + + with patch.object(v2, "logger") as mock_logger: + mqtt_on_connect_fail(mock_client_cls.return_value, block._connection) + + mock_logger.error.assert_not_called() + + def test_unreachable_broker_on_first_run_with_fail_fast_raises( + self, mock_client_cls, block + ): + mock_client_cls.return_value.connect.side_effect = ConnectionRefusedError( + "refused" + ) + + with pytest.raises(RuntimeError, match="not connected"): + block.run(**run_kwargs(fire_and_forget=True, fail_fast=True)) + + assert block.mqtt_client is mock_client_cls.return_value + + def test_permanent_refusal_still_reported(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 5) + block.run(**run_kwargs(fire_and_forget=True)) + + result = block.run(**run_kwargs(fire_and_forget=True)) + + assert result["error_status"] is True + assert "not authorised" in result["message"] + mock_client_cls.return_value.publish.assert_not_called() + + def test_disconnected_qos0_dropped_without_waiting(self, mock_client_cls, block): + mock_client = _connect_then_drop(mock_client_cls, block, timeout=2.0) + + started = time.monotonic() + result = block.run(**run_kwargs(qos=0, timeout=2.0, fire_and_forget=True)) + + assert time.monotonic() - started < 0.5 + assert result == {"error_status": True, "message": FIRE_AND_FORGET_QOS0_DROPPED} + mock_client.publish.assert_not_called() + mock_client.reconnect.assert_not_called() + + @pytest.mark.parametrize( + "rc", + [mqtt.MQTT_ERR_NO_CONN, mqtt.MQTT_ERR_SUCCESS], + ids=["no-socket", "socket"], + ) + def test_disconnected_qos1_queued_without_waiting(self, mock_client_cls, block, rc): + # paho returns NO_CONN without a socket and SUCCESS while the socket is + # still open; either way the message stays queued for the reconnect + mock_client = _connect_then_drop(mock_client_cls, block, timeout=2.0, qos=1) + mock_client.publish.return_value.rc = rc + + started = time.monotonic() + result = block.run(**run_kwargs(qos=1, timeout=2.0, fire_and_forget=True)) + + assert time.monotonic() - started < 0.5 + assert result == { + "error_status": True, + "message": FIRE_AND_FORGET_QUEUED_UNTIL_RECONNECT, + } + mock_client.publish.assert_called_once() + mock_client.publish.return_value.wait_for_publish.assert_not_called() + + def test_qos0_connection_lost_during_publish_reported_as_dropped( + self, mock_client_cls, block + ): + block._connected.set() + mock_client_cls.return_value.publish.return_value.rc = mqtt.MQTT_ERR_NO_CONN + + result = block.run(**run_kwargs(qos=0, fire_and_forget=True)) + + assert result == {"error_status": True, "message": FIRE_AND_FORGET_QOS0_DROPPED} + + def test_other_publish_error_reported(self, mock_client_cls, block): + block._connected.set() + mock_client_cls.return_value.publish.return_value.rc = mqtt.MQTT_ERR_PROTOCOL + + result = block.run(**run_kwargs(fire_and_forget=True)) + + assert result["error_status"] is True + assert result["message"].startswith("Failed to publish message") + + def test_outage_logged_when_it_starts_changes_and_ends( + self, mock_client_cls, block + ): + mock_client = _connect_then_drop(mock_client_cls, block, qos=1) + + with patch.object(v2, "logger") as mock_logger: + for _ in range(5): + result = block.run(**run_kwargs(qos=1, fire_and_forget=True)) + assert result["error_status"] is True + assert mock_logger.error.call_count == 1 + + mock_client.publish.return_value.rc = mqtt.MQTT_ERR_QUEUE_SIZE + for _ in range(3): + result = block.run(**run_kwargs(qos=1, fire_and_forget=True)) + assert result["message"] == QUEUE_FULL + assert mock_logger.error.call_count == 2 + + block._connected.set() + mock_client.publish.return_value.rc = mqtt.MQTT_ERR_SUCCESS + for _ in range(3): + result = block.run(**run_kwargs(qos=1, fire_and_forget=True)) + assert result["error_status"] is False + assert mock_logger.error.call_count == 2 + assert mock_logger.info.call_count == 1 + + def test_outage_with_fail_fast_raises_on_every_run(self, mock_client_cls, block): + _connect_then_drop(mock_client_cls, block) + + for _ in range(2): + with pytest.raises(RuntimeError, match="QoS 0 messages are not queued"): + block.run(**run_kwargs(fire_and_forget=True, fail_fast=True)) + + def test_switching_between_runs_keeps_the_connection(self, mock_client_cls, block): + block._connected.set() + + first = block.run(**run_kwargs(fire_and_forget=False)) + second = block.run(**run_kwargs(fire_and_forget=True)) + + assert first["message"] == "Message published successfully" + assert second["message"] == FIRE_AND_FORGET_QUEUED + mock_client_cls.assert_called_once() + + @pytest.mark.parametrize("value", ["maybe", 2, object()]) + def test_non_boolean_rejected_before_client_construction( + self, mock_client_cls, block, value + ): + result = block.run(**run_kwargs(fire_and_forget=value)) + + assert result["error_status"] is True + assert "Invalid fire_and_forget" in result["message"] + mock_client_cls.assert_not_called() + + @pytest.mark.parametrize("value", [None, "false", "0", 0, False]) + def test_selector_values_meaning_false_keep_waiting( + self, mock_client_cls, block, value + ): + block._connected.set() + + result = block.run(**run_kwargs(fire_and_forget=value)) + + assert result["message"] == "Message published successfully" + + @pytest.mark.parametrize("value", ["true", "1", 1, True]) + def test_selector_values_meaning_true_switch_it_on( + self, mock_client_cls, block, value + ): + block._connected.set() + + result = block.run(**run_kwargs(fire_and_forget=value)) + + assert result["message"] == FIRE_AND_FORGET_QUEUED + + +class TestSelectorSwitches: + """The engine hands run() the raw selector value, not the manifest's + coerced one, so the block coerces its switches the same way.""" + + @pytest.mark.parametrize("value, expected", [("false", False), ("true", True)]) + def test_retain_string_coerced_like_the_manifest( + self, mock_client_cls, block, value, expected + ): + block._connected.set() + + block.run(**run_kwargs(retain=value)) + + mock_client_cls.return_value.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=0, retain=expected + ) + + def test_invalid_retain_rejected_before_client_construction( + self, mock_client_cls, block + ): + result = block.run(**run_kwargs(retain="maybe")) + + assert result["error_status"] is True + assert "Invalid retain" in result["message"] + mock_client_cls.assert_not_called() + + def test_fail_fast_false_string_returns_the_failure(self, mock_client_cls, block): + result = block.run(**run_kwargs(fail_fast="false")) + + assert result["error_status"] is True + assert "not connected" in result["message"].lower() + + def test_fail_fast_true_string_raises(self, mock_client_cls, block): + with pytest.raises(RuntimeError, match="not connected"): + block.run(**run_kwargs(fail_fast="true")) + + def test_invalid_fail_fast_reported_without_raising(self, mock_client_cls, block): + result = block.run(**run_kwargs(fail_fast="maybe")) + + assert result["error_status"] is True + assert "Invalid fail_fast" in result["message"] + mock_client_cls.assert_not_called() + + def test_unset_switches_mean_their_defaults(self, mock_client_cls, block): + block._connected.set() + + result = block.run(**run_kwargs(retain=None, fail_fast=None)) + + assert result["message"] == "Message published successfully" + mock_client_cls.return_value.publish.assert_called_once_with( + "test/topic", "Hello, MQTT!", qos=0, retain=False + ) + + +class TestCleanup: + def test_close_disconnects_and_stops_loop_once(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + block.run(**run_kwargs()) + + block.close() + block.close() + block.__del__() + + mock_client.disconnect.assert_called_once() + mock_client.loop_stop.assert_called_once() + assert block.mqtt_client is None + assert not block._connected.is_set() + + def test_close_stops_loop_even_when_disconnect_raises(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client.disconnect.side_effect = Exception("socket already dead") + block.run(**run_kwargs()) + + block.close() + + mock_client.loop_stop.assert_called_once() + + def test_close_on_uninitialized_block_is_noop(self, block): + block.close() + + assert block.mqtt_client is None + + def test_close_clears_readiness_only_after_loop_thread_joined( + self, mock_client_cls, block + ): + # clearing before loop_stop() lets the dying loop's on_connect re-set + # the event; the event must stay authoritative until the join returns + block._connected.set() + mock_client = mock_client_cls.return_value + event_state_at_loop_stop = [] + mock_client.loop_stop.side_effect = lambda: event_state_at_loop_stop.append( + block._connected.is_set() + ) + block.run(**run_kwargs()) + + block.close() + + assert event_state_at_loop_stop == [True] + assert not block._connected.is_set() + + +class TestBrokerPolicy: + def test_allowlisted_broker_is_used(self, mock_client_cls, block, monkeypatch): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["localhost:1883"] + ) + block._connected.set() + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) + + def test_unlisted_broker_rejected_before_client_construction( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not permitted" in result["message"] + assert "broker.internal" not in result["message"] + mock_client_cls.assert_not_called() + + def test_unlisted_broker_with_fail_fast_raises( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + with pytest.raises(RuntimeError, match="not permitted"): + block.run(**run_kwargs(fail_fast=True)) + + mock_client_cls.assert_not_called() + + def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["operator:8883"] + ) + block._connected.set() + + first = block.run(**run_kwargs(host="workflow-a")) + second = block.run(**run_kwargs(host="workflow-b")) + + assert first["error_status"] is False + assert second["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with("operator", 8883) + + def test_user_host_not_allowed_without_operator_broker_disables_block( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "disabled" in result["message"] + mock_client_cls.assert_not_called() + + +def _call_names(mock_client) -> list: + return [call[0] for call in mock_client.method_calls] + + +@pytest.fixture +def ca_file(tmp_path): + # a real regular file: configure_tls() refuses anything else before paho + # (faked in these tests) would read it + path = tmp_path / "ca.pem" + path.write_text("content is irrelevant: paho is faked in these tests") + return path + + +class TestTLS: + def test_manifest_defaults_and_relevant_for(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + schema = BlockManifest.model_json_schema()["properties"] + + assert manifest.encryption == "none" + assert manifest.ca_certificate_path is None + assert schema["ca_certificate_path"]["relevant_for"] == { + "encryption": {"values": ["tls"], "required": True} + } + + @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None, ""]) + def test_manifest_rejects_other_encryption_values(self, encryption): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "encryption": encryption, + } + ) + + def test_init_parameters_include_file_system_access(self): + assert MQTTWriterSinkBlockV2.get_init_parameters() == [ + "disable_sinks", + "allow_access_to_file_system", + ] + + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + def test_default_encryption_never_configures_tls( + self, mock_client_cls, block, ca_certificate_path + ): + block._connected.set() + + result = block.run(**run_kwargs(ca_certificate_path=ca_certificate_path)) + + assert result["error_status"] is False + mock_client_cls.return_value.tls_set.assert_not_called() + mock_client_cls.return_value.tls_insecure_set.assert_not_called() + + def test_tls_without_ca_path_uses_system_store_before_connect( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + mock_client.tls_set.assert_called_once_with() + names = _call_names(mock_client) + assert names.index("tls_set") < names.index("connect") + mock_client.tls_insecure_set.assert_not_called() + + def test_tls_with_ca_path_and_file_system_access(self, mock_client_cls, ca_file): + block = MQTTWriterSinkBlockV2(allow_access_to_file_system=True) + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) + ) + + assert result["error_status"] is False + mock_client.tls_set.assert_called_once_with(ca_certs=str(ca_file)) + names = _call_names(mock_client) + assert names.index("tls_set") < names.index("connect") + + def test_tls_with_ca_path_refused_without_file_system_access( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd") + ) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["message"] + assert "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in result["message"] + assert block.mqtt_client is None + mock_client.tls_set.assert_not_called() + mock_client.connect.assert_not_called() + mock_client.loop_start.assert_not_called() + + def test_tls_refused_with_fail_fast_raises(self, mock_client_cls, block): + with pytest.raises(RuntimeError, match="ca_certificate_path"): + block.run( + **run_kwargs( + encryption="tls", ca_certificate_path="/etc/passwd", fail_fast=True + ) + ) + + assert block.mqtt_client is None + + def test_unloadable_ca_bundle_reported_and_nothing_kept( + self, mock_client_cls, ca_file + ): + block = MQTTWriterSinkBlockV2(allow_access_to_file_system=True) + mock_client = mock_client_cls.return_value + mock_client.tls_set.side_effect = FileNotFoundError("no such file") + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) + ) + + assert result["error_status"] is True + assert "could not load CA bundle" in result["message"] + assert str(ca_file) in result["message"] + assert block.mqtt_client is None + mock_client.connect.assert_not_called() + mock_client.loop_start.assert_not_called() + + @pytest.mark.parametrize( + "change", + [{"encryption": "tls"}, {"ca_certificate_path": "other.pem"}], + ) + def test_changed_tls_parameters_rejected(self, mock_client_cls, change): + block = MQTTWriterSinkBlockV2(allow_access_to_file_system=True) + block._connected.set() + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["message"] + mock_client_cls.assert_called_once() + + def test_invalid_encryption_at_run_time_rejected(self, mock_client_cls, block): + result = block.run(**run_kwargs(encryption="ssl")) + + assert result["error_status"] is True + assert "encryption" in result["message"] + mock_client_cls.assert_not_called() From 6ab92082c505adfa4b7ade6469b163c5cf753d1f Mon Sep 17 00:00:00 2001 From: Grzegorz Klimaszewski <166530809+grzegorz-roboflow@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:05:39 +0200 Subject: [PATCH 2/5] Fix MQTT Writer v2 refusal logging, keepalive and status messages --- workflows/CHANGELOG.md | 4 +- .../enterprise_blocks/sinks/mqtt_writer/v2.py | 92 +++++++---- .../core_steps/sinks/test_mqtt_writer_v2.py | 149 ++++++++++++++++-- 3 files changed, 203 insertions(+), 42 deletions(-) diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index f584041b17..b65496a508 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -22,8 +22,8 @@ for contributor and maintainer responsibilities. - With `fire_and_forget` a run queues the message and returns without waiting for the send or the broker's acknowledgement, and after the first run it never waits for a reconnect, so a disconnected broker no longer stalls each frame by up to `timeout`. If the first connection attempt fails, the client is kept and retries in the background instead of reconnecting synchronously on every frame. Success means queued, not delivered. - While the broker is disconnected a QoS 0 message is dropped and a QoS 1/2 message is queued and sent after reconnect, both reported with `error_status` set; these repeating failures are logged when they start or change and once on recovery, not on every run. - The step declares the `fire_and_forget_hides_persistence_failures` and `connection_and_state_rebuilt_per_request` restrictions when the switch is on (over the HTTP API every request builds a fresh block, so each request is a first run and a message still queued when the request ends is dropped), and reports a selector-driven switch as unknown. - - In either mode the client keeps at most 1000 QoS 1/2 messages awaiting acknowledgement and 1000 packets waiting to be sent (paho buffers both without limit by default); a publish beyond that is dropped and reported as a queue-full error. The limit counts messages, not bytes. - - The background reconnect delay follows `timeout` as in v1 but never drops below 0.1 s for the first retry and 1 s between later retries, so a tiny `timeout` cannot retry a dead broker in a busy loop. + - In either mode the block refuses a new publish once 1000 QoS 1/2 messages await acknowledgement or 1000 packets wait to be sent (paho buffers both without limit by default), and reports it as a queue-full error; after a reconnect up to about 2000 packets can be buffered briefly. The limits count messages, not bytes. + - The background reconnect backs off exponentially: the first retry waits half of `timeout` (at least 0.1 s), each further one doubles that, up to twice `timeout` (at least 1 s), plus the time each attempt takes, so a tiny `timeout` cannot retry a dead broker in a busy loop. - Selector-supplied `retain`, `fail_fast` and `fire_and_forget` are coerced the way the manifest validates them (for example `"false"` is False; v1 treats any non-empty string as True), and so are `port` and `qos` (for example `"1883.0"`, which v1 rejects); booleans are still rejected as a port or QoS. - While the broker stays unreachable or keeps answering "unavailable", the background reconnect loop logs the first failure as an error and later attempts at debug level until a connection succeeds, instead of an error about once per second. diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py index 303f0546fa..340cbcc2e6 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py @@ -1,6 +1,7 @@ import logging import math import threading +import time from typing import Any, List, Literal, Optional, Tuple, Type, Union import paho.mqtt.client as mqtt @@ -16,6 +17,7 @@ ) from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + MQTT_KEEPALIVE_SECONDS, PERMANENT_CONNACK_CODES, TRANSIENT_CONNACK_CODE, ConfigurationError, @@ -56,22 +58,26 @@ Compared with v1, this version adds `fire_and_forget`, bounds the client's outgoing queue, and logs a broker outage once instead of on every retry. -The first run connects synchronously: the TCP connect and the broker's -session acknowledgement are each bounded by `timeout` (DNS resolution is -bounded by the OS resolver instead), so a cold start may take up to twice -`timeout` - raise it for remote brokers. Afterwards a background network -loop maintains the connection and owns reconnects. +The first run connects synchronously: the TCP connect is bounded by `timeout` +per resolved address and the broker's session acknowledgement by `timeout`. +DNS resolution is bounded by the OS resolver and a TLS handshake by the +15 s MQTT keepalive, not by `timeout`. Raise `timeout` for remote brokers. +Afterwards a background network loop maintains the connection and owns +reconnects with an exponential backoff: the first retry waits half of +`timeout` (at least 0.1 s), each further one doubles that, up to twice +`timeout` (at least 1 s), plus the time each attempt takes. A large `timeout` +therefore also slows recovery after an outage. Without `fire_and_forget`, a run whose TCP connect fails reports the failure and leaves nothing behind; the next run on the same block instance (video pipelines) retries from scratch. Once the TCP connection is open the client is kept: a broker that does not answer within `timeout` or answers "unavailable" is retried by the background loop, which also re-establishes an -established connection that later drops. Over the HTTP API every request builds a fresh block -instance, so each request pays the bounded connect and a failed request is -final for that request. One block instance publishes to a single broker -connection: changing host, port, credentials, timeout or the TLS settings -between runs is rejected as a configuration error. +established connection that later drops. Over the HTTP API every request +builds a fresh block instance, so each request pays the bounded connect and a +failed request is final for that request. One block instance publishes to a +single broker connection: changing host, port, credentials, timeout or the TLS +settings between runs is rejected as a configuration error. The server operator may restrict which brokers the MQTT blocks connect to with `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` @@ -87,7 +93,7 @@ Workflow blocks (`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). Certificate verification cannot be disabled. A peer that accepts the TCP connection but never completes the TLS handshake is given up on after the -MQTT keepalive interval (60 s), not after `timeout`; real brokers close such +MQTT keepalive interval (15 s), not after `timeout`; real brokers close such a connection immediately. A broker that refuses the connection (bad user name or password, not @@ -109,9 +115,11 @@ QoS 1/2 message is queued by the client and sent once the background loop reconnects; both are reported with `error_status` set on every run, but logged only when the failure starts or changes, and the recovery is logged once. In -either mode the client holds at most 1000 QoS 1/2 messages awaiting -acknowledgement and 1000 packets waiting to be sent; a publish beyond that is -dropped and reported. The limit counts messages, not bytes: with large +either mode the block refuses a new publish, and reports it as dropped, once +1000 QoS 1/2 messages await acknowledgement or 1000 packets wait to be sent +(the client's own limit counts only QoS 1/2 messages). After a reconnect the +client re-sends its queued QoS 1/2 messages, so up to about 2000 packets can +be buffered briefly. The limits count messages, not bytes: with large payloads a long outage can still hold up to 1000 of them in memory, so prefer QoS 0 for large messages. Fire and forget is meant for InferencePipelines: over the HTTP API every request builds a fresh block instance, so each request is a @@ -125,7 +133,8 @@ or success confirmation if error_status is False. A publish acknowledgement timeout on QoS 1/2 is reported as delivery-unknown (the message may still be delivered); - an unconfirmed QoS 0 send is reported as lost. With + an unconfirmed QoS 0 send is reported as not fully sent + within timeout (it may still be sent or may be lost). With fire_and_forget a successful run only confirms the message was queued. @@ -212,10 +221,12 @@ def _not_connected_message(error: BaseException, encryption: str) -> str: "the broker's TLS port and that ca_certificate_path matches the broker's " "certificate authority." ) - return ( - f"MQTT broker not connected ({error}). Raise 'timeout' if the broker needs " - "longer to connect." - ) + if isinstance(error, TimeoutError): + return ( + f"MQTT broker not connected ({error}). Raise 'timeout' if the broker " + "needs longer to connect." + ) + return f"MQTT broker not connected ({error})." WRITER_CONNECTION_INPUTS = "username and password" @@ -326,6 +337,9 @@ def mqtt_on_connect_fail(client, state: MQTTWriterState): logger.error("MQTT client failed to establish connection with broker") state.connect_failing = True state.connected.clear() + if state.refused_code not in PERMANENT_CONNACK_CODES: + state.refused_code = None + state.connack.clear() def mqtt_on_disconnect(client, state: MQTTWriterState, reason_code, properties=None): @@ -581,6 +595,8 @@ def run( ) # selector-resolved values bypass manifest constraints, so validate here try: + if isinstance(timeout, bool): + raise TypeError("a boolean is not a timeout") timeout_seconds = float(timeout) except (TypeError, ValueError, OverflowError): timeout_seconds = math.nan @@ -733,7 +749,7 @@ def _connect_and_publish( # the TCP connect happens here, bounded by _connect_timeout # (DNS resolution is not - it runs under the OS resolver # timeout); the CONNACK wait below covers the handshake rest - client.connect(host, port) + client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) client.loop_start() except ConfigurationError as e: # TLS setup refused before any socket was opened: nothing kept @@ -777,11 +793,12 @@ def _connect_and_publish( # a permanent refusal stopped the network loop: report it on every run # without waiting and without touching the broker again if state.refused_code in PERMANENT_CONNACK_CODES: - return self._handle_failure( + return self._report_failure( connection_refused_message( state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS ), fail_fast=fail_fast, + fire_and_forget=fire_and_forget, ) # fire and forget waits for the broker only on the run that built the # client; afterwards a disconnected run returns at once and leaves the @@ -799,7 +816,11 @@ def _connect_and_publish( if not_ready is not None: # a permanent refusal stopped the loop: nothing would be sent if not fire_and_forget or state.refused_code in PERMANENT_CONNACK_CODES: - return self._handle_failure(not_ready, fail_fast=fail_fast) + return self._report_failure( + not_ready, + fail_fast=fail_fast, + fire_and_forget=fire_and_forget, + ) return self._report_not_connected( not_ready, topic=topic, @@ -838,6 +859,12 @@ def _connect_and_publish( "delivery status unknown and the message may still be delivered.", fail_fast=fail_fast, ) + if res.rc == mqtt.MQTT_ERR_NO_CONN and qos == 0: + return self._handle_failure( + "MQTT broker not connected; QoS 0 messages are not queued, so the " + "message was dropped.", + fail_fast=fail_fast, + ) try: res.wait_for_publish(timeout=timeout) published = res.is_published() @@ -851,12 +878,9 @@ def _connect_and_publish( "message": "Message published successfully", } if qos == 0: - # an unconfirmed QoS 0 send is a real loss: paho never queues - # QoS 0 messages and reconnect() clears the pending out packet return self._handle_failure( - "Publish confirmation timed out; the connection dropped before " - "the message was fully sent and QoS 0 messages are not " - "retransmitted, so the message is lost.", + "Publish confirmation timed out; the QoS 0 message was not fully " + "sent within timeout and may still be sent or may be lost.", fail_fast=fail_fast, ) return self._handle_failure( @@ -907,20 +931,21 @@ def _wait_until_connected(self, timeout: float) -> Optional[str]: """Wait for the background loop to establish the session. Args: - timeout: Bound for the CONNACK wait and, after a transient refusal, - for the reconnect. + timeout: Total bound for the CONNACK wait and, after a transient + refusal, the reconnect. Returns: None once connected, otherwise the reason the broker is not ready. """ state = self._connection + deadline = time.monotonic() + timeout if not state.connack.wait(timeout=timeout): return NOT_CONNECTED_WITHIN_TIMEOUT if state.refused_code in PERMANENT_CONNACK_CODES: return connection_refused_message( state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS ) - if not state.connected.wait(timeout=timeout): + if not state.connected.wait(timeout=max(0.0, deadline - time.monotonic())): if state.refused_code == TRANSIENT_CONNACK_CODE: return connection_refused_message( state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS @@ -986,6 +1011,13 @@ def _fire_and_forget_result( fail_fast=fail_fast, ) + def _report_failure( + self, message: str, fail_fast: bool, fire_and_forget: bool + ) -> BlockResult: + if fire_and_forget: + return self._fire_and_forget_outage(message, fail_fast=fail_fast) + return self._handle_failure(message, fail_fast=fail_fast) + def _fire_and_forget_outage(self, message: str, fail_fast: bool) -> BlockResult: # without the readiness wait a pipeline reaches this on every frame of # an outage; log when the failure starts or changes, not once per run diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py index 34b9709dfe..8d081bc58f 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py @@ -1,5 +1,6 @@ import threading import time +from types import SimpleNamespace from typing import get_args from unittest.mock import MagicMock, patch @@ -7,6 +8,9 @@ import pytest from pydantic import ValidationError from roboflow_workflows.enterprise_blocks.sinks import mqtt_common +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + MQTT_KEEPALIVE_SECONDS, +) from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v2 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( FIRE_AND_FORGET_QOS0_DROPPED, @@ -315,6 +319,24 @@ def test_on_disconnect_accepts_mqtt5_properties_argument(self): assert not state.connected.is_set() + def test_connect_fail_after_transient_refusal_forgets_the_refusal(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 3) + + mqtt_on_connect_fail(MagicMock(), state) + + assert state.refused_code is None + assert not state.connack.is_set() + + def test_connect_fail_after_permanent_refusal_keeps_the_refusal(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) + + mqtt_on_connect_fail(MagicMock(), state) + + assert state.refused_code == 5 + assert state.connack.is_set() + def test_reset_clears_everything(self): state = MQTTWriterState() mqtt_on_connect(MagicMock(), state, {}, 5) @@ -405,6 +427,32 @@ def test_refusal_with_fail_fast_raises_with_the_reason( with pytest.raises(Exception, match="not authorised"): block.run(**run_kwargs(fail_fast=True)) + def test_fire_and_forget_logs_a_permanent_refusal_once( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + with patch.object(v2, "logger") as mock_logger: + results = [block.run(**run_kwargs(fire_and_forget=True)) for _ in range(3)] + + assert all(result["error_status"] is True for result in results) + assert all("not authorised" in result["message"] for result in results) + failure_logs = [ + call + for call in mock_logger.error.call_args_list + if call.args[0] == "MQTT Writer failure: %s" + ] + assert len(failure_logs) == 1 + + def test_fire_and_forget_permanent_refusal_still_raises_with_fail_fast( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + block.run(**run_kwargs(fire_and_forget=True)) + + with pytest.raises(RuntimeError, match="not authorised"): + block.run(**run_kwargs(fire_and_forget=True, fail_fast=True)) + def test_broker_unavailable_keeps_retrying_and_names_the_reason( self, mock_client_cls, block ): @@ -464,6 +512,16 @@ def test_invalid_timeout_rejected_before_client_construction( assert "timeout" in result["message"].lower() mock_client_cls.assert_not_called() + @pytest.mark.parametrize("timeout", [True, False]) + def test_boolean_timeout_rejected_before_client_construction( + self, mock_client_cls, block, timeout + ): + result = block.run(**run_kwargs(timeout=timeout)) + + assert result["error_status"] is True + assert "Invalid timeout" in result["message"] + mock_client_cls.assert_not_called() + @pytest.mark.parametrize( "port", [0, -1, 65536, "abc", "1883.5", 1883.7, True, None] ) @@ -487,7 +545,9 @@ def test_selector_supplied_port_coerced_like_manifest( result = block.run(**run_kwargs(port=port)) assert result["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) + mock_client_cls.return_value.connect.assert_called_once_with( + "localhost", 1883, keepalive=MQTT_KEEPALIVE_SECONDS + ) @pytest.mark.parametrize("qos", [-1, 3, "abc", 1.5, "1.5", True, False, None]) def test_invalid_qos_rejected_before_client_construction( @@ -553,11 +613,69 @@ def test_client_connects_synchronously_before_background_loop( assert result["error_status"] is False mock_client_cls.assert_called_once_with(userdata=block._connection) - mock_client.connect.assert_called_once_with("localhost", 1883) + mock_client.connect.assert_called_once_with( + "localhost", 1883, keepalive=MQTT_KEEPALIVE_SECONDS + ) mock_client.connect_async.assert_not_called() called_methods = [call[0] for call in mock_client.method_calls] assert called_methods.index("connect") < called_methods.index("loop_start") + def test_client_connects_with_the_short_keepalive(self, mock_client_cls, block): + block._connected.set() + + block.run(**run_kwargs()) + + mock_client_cls.return_value.connect.assert_called_once_with( + "localhost", 1883, keepalive=MQTT_KEEPALIVE_SECONDS + ) + + def test_timeout_hint_only_for_a_connect_timeout(self, mock_client_cls, block): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = TimeoutError("timed out") + + result = block.run(**run_kwargs()) + + assert "Raise 'timeout'" in result["message"] + + def test_no_timeout_hint_when_the_broker_refuses_the_connection( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + mock_client.connect.side_effect = ConnectionRefusedError("refused") + + result = block.run(**run_kwargs()) + + assert "not connected" in result["message"] + assert "Raise 'timeout'" not in result["message"] + + def test_readiness_wait_shares_one_deadline(self, mock_client_cls, block): + clock = [100.0] + waits = [] + + class FakeEvent: + def __init__(self, spent, result): + self.spent = spent + self.result = result + + def is_set(self): + return False + + def wait(self, timeout=None): + waits.append(timeout) + clock[0] += self.spent + return self.result + + block._connection.connack = FakeEvent(spent=0.15, result=True) + block._connection.connected = FakeEvent(spent=0.05, result=False) + fake_time = SimpleNamespace(monotonic=lambda: clock[0]) + + with patch.object(v2, "time", fake_time): + result = block.run(**run_kwargs(timeout=0.2)) + + assert result["error_status"] is True + assert result["message"] == NOT_CONNECTED_WITHIN_TIMEOUT + assert waits == [pytest.approx(0.2), pytest.approx(0.05)] + def test_client_registers_static_callbacks(self, mock_client_cls, block): block._connected.set() mock_client = mock_client_cls.return_value @@ -727,7 +845,9 @@ def test_changed_connection_parameters_rejected(self, mock_client_cls, block): assert result["error_status"] is True assert "parameters" in result["message"].lower() - mock_client.connect.assert_called_once_with("broker-a", 1883) + mock_client.connect.assert_called_once_with( + "broker-a", 1883, keepalive=MQTT_KEEPALIVE_SECONDS + ) mock_client.publish.assert_called_once() def test_changed_credentials_rejected(self, mock_client_cls, block): @@ -766,11 +886,9 @@ def test_successful_publish(self, mock_client_cls, block): "test/topic", "Hello, MQTT!", qos=1, retain=True ) - def test_qos0_publish_confirmation_timeout_reported_as_lost( + def test_qos0_publish_confirmation_timeout_reported_as_unconfirmed( self, mock_client_cls, block ): - # paho never retransmits QoS 0: reconnect() clears _out_packet and - # QoS 0 messages are not queued, so an unconfirmed send is a real loss block._connected.set() mock_client = mock_client_cls.return_value mock_client.publish.return_value.is_published.return_value = False @@ -779,7 +897,10 @@ def test_qos0_publish_confirmation_timeout_reported_as_lost( result = block.run(**run_kwargs(qos=0)) assert result["error_status"] is True - assert "lost" in result["message"].lower() + assert result["message"] == ( + "Publish confirmation timed out; the QoS 0 message was not fully sent " + "within timeout and may still be sent or may be lost." + ) assert "delivery status unknown" not in result["message"].lower() assert mock_logger.error.called @@ -831,7 +952,11 @@ def test_qos0_publish_with_lost_connection_reported_as_final_failure( result = block.run(**run_kwargs(qos=0)) assert result["error_status"] is True - assert "failed to publish" in result["message"].lower() + assert result["message"] == ( + "MQTT broker not connected; QoS 0 messages are not queued, so the " + "message was dropped." + ) + publish_result.wait_for_publish.assert_not_called() def test_publish_exception_returned_as_error(self, mock_client_cls, block): block._connected.set() @@ -1410,7 +1535,9 @@ def test_allowlisted_broker_is_used(self, mock_client_cls, block, monkeypatch): result = block.run(**run_kwargs()) assert result["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) + mock_client_cls.return_value.connect.assert_called_once_with( + "localhost", 1883, keepalive=MQTT_KEEPALIVE_SECONDS + ) def test_unlisted_broker_rejected_before_client_construction( self, mock_client_cls, block, monkeypatch @@ -1454,7 +1581,9 @@ def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( assert first["error_status"] is False assert second["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with("operator", 8883) + mock_client_cls.return_value.connect.assert_called_once_with( + "operator", 8883, keepalive=MQTT_KEEPALIVE_SECONDS + ) def test_user_host_not_allowed_without_operator_broker_disables_block( self, mock_client_cls, block, monkeypatch From 2d79411b3670f4873c956f8b5be406174e0de288 Mon Sep 17 00:00:00 2001 From: Grzegorz Klimaszewski <166530809+grzegorz-roboflow@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:37:28 +0200 Subject: [PATCH 3/5] Bound MQTT Writer v2 close() when it races a background reconnect --- .../integration_tests/execution/conftest.py | 7 + .../test_workflow_with_mqtt_writer_v2.py | 204 ++++++++++++++++++ workflows/CHANGELOG.md | 1 + .../enterprise_blocks/sinks/mqtt_writer/v2.py | 65 +++++- .../core_steps/sinks/test_mqtt_writer_v2.py | 149 +++++++++++++ 5 files changed, 421 insertions(+), 5 deletions(-) diff --git a/tests/workflows/integration_tests/execution/conftest.py b/tests/workflows/integration_tests/execution/conftest.py index 9f57e0021e..206121f661 100644 --- a/tests/workflows/integration_tests/execution/conftest.py +++ b/tests/workflows/integration_tests/execution/conftest.py @@ -288,6 +288,7 @@ def __init__( keep_serving: bool = False, tls_context: Optional[ssl.SSLContext] = None, session_aware: bool = False, + connack_gate: Optional[threading.Event] = None, ): # Bind to "localhost" for maximum performance, as described in: # http://docs.python.org/howto/sockets.html#ipc @@ -308,6 +309,9 @@ def __init__( # wrapped with it before any MQTT packet is read self.tls_context = tls_context self.session_aware = session_aware + # connack_gate: every CONNACK is withheld until the event is set + self.connack_gate = connack_gate + self.disconnects = 0 self.handshake_failures = 0 self.subscriptions = [] self.retained = {} @@ -462,6 +466,7 @@ def _handle_packet(self, connection: _Connection, packet: bytes) -> bool: self._send(connection, bytes([MQTT_PUBCOMP, 2]) + packet[2:4]) return False if packet_type == MQTT_DISCONNECT: + self.disconnects += 1 return False self.messages.append(packet) if packet_type == MQTT_PUBLISH and self.session_aware: @@ -483,6 +488,8 @@ def _read_string(body: bytes, index: int) -> Tuple[str, int]: return body[index : index + length].decode("utf-8"), index + length def _handle_connect(self, connection: _Connection, packet: bytes) -> None: + if self.connack_gate is not None: + self.connack_gate.wait() print("sending CONNACK") if not self.session_aware: self._send(connection, b"\x20\x02\x00" + bytes([self.connack_reason_code])) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py index a9f79abc0b..85529e7a15 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py @@ -6,6 +6,9 @@ from unittest.mock import patch import pytest +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + MQTT_KEEPALIVE_SECONDS, +) from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v2 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( MQTTWriterSinkBlockV2, @@ -501,6 +504,207 @@ def test_publish_over_tls_with_ca_certificate(mqtt_test_certificates): broker.finish() +class ReconnectGate: + def __init__(self, client): + self.entered = threading.Event() + self.released = threading.Event() + create_socket_connection = client._create_socket_connection + disconnect = client.disconnect + + def gated_create_socket_connection(): + self.entered.set() + self.released.wait(timeout=10) + return create_socket_connection() + + def releasing_disconnect(*args, **kwargs): + rc = disconnect(*args, **kwargs) + self.released.set() + return rc + + client._create_socket_connection = gated_create_socket_connection + client.disconnect = releasing_disconnect + + +def hold_reconnect_until_close_disconnects(block, broker) -> ReconnectGate: + gate = ReconnectGate(block.mqtt_client) + broker.drop_connection() + assert gate.entered.wait(timeout=5) + return gate + + +@pytest.mark.timeout(15) +def test_close_leaves_no_network_thread_or_socket(): + broker = FakeMQTTBroker(keep_serving=True) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + + try: + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="delivered", + timeout=2.0, + ) + client = block.mqtt_client + thread = client._thread + + block.close() + + assert result["error_status"] is False + assert not thread.is_alive() + assert client._sock is None + assert wait_until(lambda: broker.disconnects == 1) + finally: + block.close() + broker.finish() + + +@pytest.mark.timeout(15) +def test_close_drops_queued_messages_without_waiting_for_acknowledgement(): + broker = FakeMQTTBroker(keep_serving=True) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + kwargs = dict( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + qos=1, + timeout=1.0, + fire_and_forget=True, + ) + + try: + results = [block.run(message=f"unacknowledged {i}", **kwargs) for i in range(3)] + client = block.mqtt_client + assert wait_until(lambda: len(broker.messages) == 3) + assert len(client._out_messages) == 3 + + block.close() + + assert all(result["error_status"] is False for result in results) + assert len(client._out_messages) == 3 + assert wait_until(lambda: broker.disconnects == 1) + assert len(broker.messages) == 3 + finally: + block.close() + broker.finish() + + +@pytest.mark.timeout(30) +def test_close_during_reconnect_with_unacknowledged_messages_returns(): + broker = FakeMQTTBroker(keep_serving=True) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + gate = None + + try: + first = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="in flight", + qos=1, + timeout=0.5, + fire_and_forget=True, + ) + assert first["error_status"] is False + client = block.mqtt_client + thread = client._thread + assert wait_until(lambda: len(broker.messages) == 1) + gate = hold_reconnect_until_close_disconnects(block, broker) + assert client._out_messages + + block.close() + + assert not thread.is_alive() + assert client._sock is None + assert client._out_messages + assert not block._connected.is_set() + assert wait_until(lambda: broker.disconnects == 1) + assert broker.connections_accepted == 2 + assert len(broker.messages) == 1 + finally: + if gate is not None: + gate.released.set() + block.close() + broker.finish() + + +@pytest.mark.timeout(30) +def test_close_during_reconnect_ends_a_loop_waiting_for_a_connack_that_never_comes(): + connack_gate = threading.Event() + broker = FakeMQTTBroker(keep_serving=True, connack_gate=connack_gate) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + gate = None + + try: + first = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="queued", + qos=1, + timeout=0.5, + fire_and_forget=True, + ) + assert first["error_status"] is True + client = block.mqtt_client + thread = client._thread + assert wait_until(lambda: broker.connections_accepted == 1) + gate = hold_reconnect_until_close_disconnects(block, broker) + assert client._out_messages + + started = time.monotonic() + block.close() + elapsed = time.monotonic() - started + + assert elapsed < MQTT_KEEPALIVE_SECONDS + assert not thread.is_alive() + assert client._sock is None + assert wait_until(lambda: broker.connections_accepted == 2) + finally: + if gate is not None: + gate.released.set() + connack_gate.set() + block.close() + broker.finish() + + +@pytest.mark.timeout(15) +def test_close_from_a_callback_on_the_network_thread_lets_the_loop_end_itself(): + broker = FakeMQTTBroker(keep_serving=True) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + + try: + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="delivered", + timeout=2.0, + ) + client = block.mqtt_client + thread = client._thread + client.on_disconnect = lambda _client, _state, _rc: block.close() + + with patch.object(v2, "logger") as mock_logger: + broker.drop_connection() + assert wait_until(lambda: not thread.is_alive()) + + assert result["error_status"] is False + assert block.mqtt_client is None + assert block._connection.closing.is_set() + assert client._sock is None + assert broker.connections_accepted == 1 + mock_logger.error.assert_not_called() + finally: + block.close() + broker.finish() + + MQTT_SINK_WORKFLOW = { "version": "1.0", "inputs": [ diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index b65496a508..8abe8a42f6 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -26,6 +26,7 @@ for contributor and maintainer responsibilities. - The background reconnect backs off exponentially: the first retry waits half of `timeout` (at least 0.1 s), each further one doubles that, up to twice `timeout` (at least 1 s), plus the time each attempt takes, so a tiny `timeout` cannot retry a dead broker in a busy loop. - Selector-supplied `retain`, `fail_fast` and `fire_and_forget` are coerced the way the manifest validates them (for example `"false"` is False; v1 treats any non-empty string as True), and so are `port` and `qos` (for example `"1883.0"`, which v1 rejects); booleans are still rejected as a port or QoS. - While the broker stays unreachable or keeps answering "unavailable", the background reconnect loop logs the first failure as an error and later attempts at debug level until a connection succeeds, instead of an error about once per second. + - Closing the block (when a pipeline stops) is bounded even when it races a background reconnect: a connection the broker accepts after the close began is disconnected at once, and the close never waits for acknowledgements or for the broker; messages not yet acknowledged are dropped, though packets already queued may still be sent before the disconnect. A close during an in-flight reconnect first waits for that connection attempt (TCP bounded by `timeout`, a TLS handshake by the 15 s keepalive, DNS by the OS resolver). ## `0.2.4` diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py index 340cbcc2e6..8d588b6435 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py @@ -1,5 +1,6 @@ import logging import math +import socket import threading import time from typing import Any, List, Literal, Optional, Tuple, Type, Union @@ -123,7 +124,13 @@ payloads a long outage can still hold up to 1000 of them in memory, so prefer QoS 0 for large messages. Fire and forget is meant for InferencePipelines: over the HTTP API every request builds a fresh block instance, so each request is a -first run, and a message still queued when the request ends is dropped. +first run, and a message still queued when the request ends is dropped. In +either mode, closing the block (when the pipeline stops) never waits for +acknowledgements or for the broker: messages not yet acknowledged are dropped, +though packets already queued for sending may still be written before the +disconnect. A close that lands during a background reconnect first waits for +that connection attempt (TCP bounded by `timeout`, a TLS handshake by the 15 s +keepalive, DNS by the OS resolver). Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. @@ -245,6 +252,9 @@ def _not_connected_message(error: BaseException, encryption: str) -> str: # dead broker in a busy loop MIN_RECONNECT_DELAY_SECONDS = 0.1 MIN_MAX_RECONNECT_DELAY_SECONDS = 1.0 +# paho's network loop checks for a stop request once per select timeout (1 s) +CLOSE_JOIN_SECONDS = 1.0 +CLOSE_JOIN_ATTEMPTS = 3 QUEUE_FULL = ( f"MQTT client queue is full ({MAX_QUEUED_MESSAGES} messages awaiting " "delivery); the message was dropped." @@ -279,6 +289,10 @@ class MQTTWriterState: connect_failing: Set by a failed connection attempt or a refused CONNACK and cleared by an accepted CONNACK, so an outage is logged once rather than on every background retry. + closing: Set by ``close()`` before it asks the client to disconnect, so a + CONNACK landing after that point disconnects the client instead of + reviving the session; cleared by ``reset()`` once the network thread + is gone. """ def __init__(self): @@ -286,17 +300,23 @@ def __init__(self): self.connack = threading.Event() self.refused_code: Optional[int] = None self.connect_failing = False + self.closing = threading.Event() def reset(self) -> None: self.connected.clear() self.connack.clear() self.refused_code = None self.connect_failing = False + self.closing.clear() def mqtt_on_connect( client, state: MQTTWriterState, flags, reason_code, properties=None ): + if state.closing.is_set(): + # a reconnect that raced close(): this CONNACK must end the loop + client.disconnect() + return # paho invokes on_connect for accepted and rejected CONNACK alike; # only reason_code 0 means an established MQTT session. The outcome is # recorded before `connack` wakes a waiting run, so the run never reads a @@ -355,6 +375,38 @@ def mqtt_on_disconnect(client, state: MQTTWriterState, reason_code, properties=N state.connack.clear() +def _stop_network_loop(client: mqtt.Client) -> bool: + """Stop paho's network thread within a bounded time, then close its socket. + + Returns False when called on that thread itself: termination is requested + and the thread ends on its own after the caller returns. + """ + thread = client._thread + if thread is threading.current_thread(): + client._thread_terminate = True + return False + if thread is not None: + client._thread_terminate = True + for _ in range(CLOSE_JOIN_ATTEMPTS): + thread.join(timeout=CLOSE_JOIN_SECONDS) + if not thread.is_alive(): + break + _shut_down_socket(client) + client.loop_stop() + client._sock_close() + return True + + +def _shut_down_socket(client: mqtt.Client) -> None: + sock = client._sock + if sock is None: + return + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + + class BlockManifest(WorkflowBlockManifest): model_config = ConfigDict( json_schema_extra={ @@ -540,19 +592,22 @@ def close(self) -> None: self.mqtt_client = None self._connection_identity = None self._fire_and_forget_failure = None + self._connection.closing.set() + # _reconnect_on_failure also gates paho's inline reconnect on CONNACK 1/2 + client._reconnect_on_failure = False try: client.disconnect() except Exception as e: logger.error("Failed to disconnect MQTT client: %s", e) finally: + loop_stopped = True try: - # loop_stop() joins the network thread without a timeout; - # accepted so the thread never outlives the block - client.loop_stop() + loop_stopped = _stop_network_loop(client) finally: # only after the join can no callback re-set the events; a # client rebuilt by a later run starts without a stale refusal - self._connection.reset() + if loop_stopped: + self._connection.reset() def __del__(self): try: diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py index 8d081bc58f..e29fcc1e2c 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py @@ -1,3 +1,4 @@ +import socket import threading import time from types import SimpleNamespace @@ -13,6 +14,7 @@ ) from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v2 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + CLOSE_JOIN_ATTEMPTS, FIRE_AND_FORGET_QOS0_DROPPED, FIRE_AND_FORGET_QUEUED, FIRE_AND_FORGET_QUEUED_UNTIL_RECONNECT, @@ -341,6 +343,7 @@ def test_reset_clears_everything(self): state = MQTTWriterState() mqtt_on_connect(MagicMock(), state, {}, 5) mqtt_on_connect_fail(MagicMock(), state) + state.closing.set() state.reset() @@ -348,6 +351,35 @@ def test_reset_clears_everything(self): assert not state.connack.is_set() assert state.refused_code is None assert state.connect_failing is False + assert not state.closing.is_set() + + @pytest.mark.parametrize("reason_code", [0, 3, 5]) + def test_on_connect_while_closing_disconnects_without_marking_connected( + self, reason_code + ): + client = MagicMock() + state = MQTTWriterState() + state.closing.set() + + mqtt_on_connect(client, state, {}, reason_code) + + assert not state.connected.is_set() + assert not state.connack.is_set() + assert state.refused_code is None + client.disconnect.assert_called_once_with() + + @pytest.mark.timeout(10) + def test_callbacks_complete_while_the_lifecycle_lock_is_held(self, block): + client = MagicMock() + state = block._connection + state.closing.set() + + with block._lifecycle_lock: + mqtt_on_connect(client, state, {}, 0) + mqtt_on_connect_fail(client, state) + mqtt_on_disconnect(client, state, 0) + + client.disconnect.assert_called_once_with() def _answer_connect_with(mock_client_cls, reason_code: int) -> None: @@ -665,16 +697,22 @@ def wait(self, timeout=None): clock[0] += self.spent return self.result + def clear(self): + pass + block._connection.connack = FakeEvent(spent=0.15, result=True) block._connection.connected = FakeEvent(spent=0.05, result=False) fake_time = SimpleNamespace(monotonic=lambda: clock[0]) with patch.object(v2, "time", fake_time): result = block.run(**run_kwargs(timeout=0.2)) + with patch.object(v2, "logger") as mock_logger: + block.close() assert result["error_status"] is True assert result["message"] == NOT_CONNECTED_WITHIN_TIMEOUT assert waits == [pytest.approx(0.2), pytest.approx(0.05)] + mock_logger.error.assert_not_called() def test_client_registers_static_callbacks(self, mock_client_cls, block): block._connected.set() @@ -1524,6 +1562,117 @@ def test_close_clears_readiness_only_after_loop_thread_joined( assert event_state_at_loop_stop == [True] assert not block._connected.is_set() + def test_close_marks_the_state_closing_before_disconnecting( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + closing_at_disconnect = [] + mock_client.disconnect.side_effect = lambda: closing_at_disconnect.append( + block._connection.closing.is_set() + ) + block.run(**run_kwargs()) + + block.close() + + assert closing_at_disconnect == [True] + assert not block._connection.closing.is_set() + + def test_run_after_close_builds_a_fresh_client_that_is_not_closing( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 0) + block.run(**run_kwargs()) + block.close() + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert mock_client_cls.call_count == 2 + assert block._connection.connected.is_set() + assert not block._connection.closing.is_set() + + def test_close_shuts_the_socket_down_while_the_loop_thread_outlives_the_join( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client._thread.is_alive.return_value = True + block.run(**run_kwargs()) + + block.close() + + assert mock_client._thread.join.call_count == CLOSE_JOIN_ATTEMPTS + assert mock_client._sock.shutdown.call_count == CLOSE_JOIN_ATTEMPTS + mock_client._sock.shutdown.assert_called_with(socket.SHUT_RDWR) + mock_client.loop_stop.assert_called_once() + mock_client._sock_close.assert_called_once() + + def test_close_does_not_touch_the_socket_once_the_loop_thread_ended( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client._thread.is_alive.return_value = False + block.run(**run_kwargs()) + + block.close() + + assert mock_client._thread.join.call_count == 1 + mock_client._sock.shutdown.assert_not_called() + mock_client.loop_stop.assert_called_once() + mock_client._sock_close.assert_called_once() + + def test_close_from_the_network_thread_only_requests_termination( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + mock_client._thread = threading.current_thread() + block.run(**run_kwargs()) + + with patch.object(v2, "logger") as mock_logger: + block.close() + + assert mock_client._thread_terminate is True + assert block.mqtt_client is None + assert block._connection.closing.is_set() + mock_client._sock_close.assert_not_called() + mock_client._sock.shutdown.assert_not_called() + mock_logger.error.assert_not_called() + + def test_close_disables_paho_reconnects_before_disconnecting( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + reconnect_flag_at_disconnect = [] + mock_client.disconnect.side_effect = ( + lambda: reconnect_flag_at_disconnect.append( + mock_client._reconnect_on_failure + ) + ) + block.run(**run_kwargs()) + + block.close() + + assert reconnect_flag_at_disconnect == [False] + + @pytest.mark.parametrize("reason_code", [1, 2]) + def test_connack_refusal_after_close_does_not_reconnect(self, block, reason_code): + client = mqtt.Client(userdata=block._connection) + client.reconnect = MagicMock() + block.mqtt_client = client + block.close() + client._in_packet["remaining_length"] = 2 + client._in_packet["packet"] = bytes([0, reason_code]) + + rc = client._handle_connack() + + client.reconnect.assert_not_called() + assert rc == mqtt.MQTT_ERR_PROTOCOL + assert not block._connection.connack.is_set() + class TestBrokerPolicy: def test_allowlisted_broker_is_used(self, mock_client_cls, block, monkeypatch): From dccc7266bb01361f029487bea52fa2c2a9a0b83d Mon Sep 17 00:00:00 2001 From: Grzegorz Klimaszewski <166530809+grzegorz-roboflow@users.noreply.github.com> Date: Fri, 9 Oct 2026 11:48:06 +0200 Subject: [PATCH 4/5] Tighten MQTT Writer v2 manifest types and tidy comments --- .../core_steps/common/workload_presets.py | 4 +- .../enterprise_blocks/sinks/mqtt_writer/v2.py | 130 ++++++------------ .../core_steps/sinks/test_mqtt_writer_v2.py | 66 +++++++++ 3 files changed, 111 insertions(+), 89 deletions(-) diff --git a/workflows/roboflow_workflows/core_steps/common/workload_presets.py b/workflows/roboflow_workflows/core_steps/common/workload_presets.py index 5b79c969e6..8ca75c7024 100644 --- a/workflows/roboflow_workflows/core_steps/common/workload_presets.py +++ b/workflows/roboflow_workflows/core_steps/common/workload_presets.py @@ -210,8 +210,8 @@ # Fire-and-forget writes hide delivery / persistence failures and pile up when # the destination is slower than the video stream. Shared by the sinks that # offer the switch (PostgreSQL, Kafka producer, Microsoft SQL Server, Event -# Writer, OPC UA Writer), so the note stays neutral about what the destination -# is. +# Writer, OPC UA Writer, MQTT Writer), so the note stays neutral about what the +# destination is. FIRE_AND_FORGET_RESTRICTION = RuntimeRestriction( code="fire_and_forget_hides_persistence_failures", severity=Severity.SOFT, diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py index 8d588b6435..37a4193438 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v2.py @@ -7,12 +7,6 @@ import paho.mqtt.client as mqtt from pydantic import ConfigDict, Field, TypeAdapter, ValidationError -from typing_extensions import Annotated - -# Bind by literal name to the shared "inference" logger tree so server -# handlers / filters / propagation still apply when installed, without -# importing the server logger module. -logger = logging.getLogger("inference") from roboflow_workflows.core_steps.common.workload_presets import ( FIRE_AND_FORGET_RESTRICTION, ) @@ -52,6 +46,10 @@ actual_restrictions_of, is_workflow_selector, ) +from typing_extensions import Annotated + +# reuse server logging configuration without importing the server package +logger = logging.getLogger("inference") LONG_DESCRIPTION = """ MQTT Writer block for publishing messages to an MQTT broker. @@ -153,8 +151,7 @@ TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} -# the fire-and-forget client and its queue live in the block instance, which -# the HTTP API rebuilds for every request +# the client and its queue live in the block instance, rebuilt per HTTP request FIRE_AND_FORGET_PER_REQUEST_RESTRICTION = RuntimeRestriction( code="connection_and_state_rebuilt_per_request", severity=Severity.SOFT, @@ -221,7 +218,7 @@ def _coerce_switch(value: Any) -> bool: return BOOLEAN_ADAPTER.validate_python(value) -def _not_connected_message(error: BaseException, encryption: str) -> str: +def _not_connected_message(error: BaseException, *, encryption: str) -> str: if encryption == "tls": return ( f"MQTT broker not connected ({error}). TLS is enabled: check that port is " @@ -242,14 +239,9 @@ def _not_connected_message(error: BaseException, encryption: str) -> str: "the client keeps retrying in the background. Raise 'timeout' if the " "broker needs longer to connect." ) -# paho keeps QoS 1/2 messages until the broker acknowledges them and QoS 0 -# packets until the socket takes them, both by default without limit; a -# fire-and-forget pipeline publishing through a long outage or to a slow broker -# would otherwise grow the buffer for as long as that lasts +# paho's buffers are unbounded by default; cap them so a long outage cannot grow memory MAX_QUEUED_MESSAGES = 1000 -# the reconnect schedule follows `timeout` (half of it for the first retry, up -# to twice it between retries), bounded below so a tiny timeout cannot retry a -# dead broker in a busy loop +# floors keep a tiny timeout from retrying a dead broker in a busy loop MIN_RECONNECT_DELAY_SECONDS = 0.1 MIN_MAX_RECONNECT_DELAY_SECONDS = 1.0 # paho's network loop checks for a stop request once per select timeout (1 s) @@ -317,10 +309,7 @@ def mqtt_on_connect( # a reconnect that raced close(): this CONNACK must end the loop client.disconnect() return - # paho invokes on_connect for accepted and rejected CONNACK alike; - # only reason_code 0 means an established MQTT session. The outcome is - # recorded before `connack` wakes a waiting run, so the run never reads a - # half-updated state + # only reason_code 0 is a session; record state before connack wakes a waiting run if reason_code == 0: state.refused_code = None state.connect_failing = False @@ -331,8 +320,7 @@ def mqtt_on_connect( state.refused_code = reason_code state.connected.clear() state.connack.set() - # a transient refusal comes back on every background retry, so only the - # first one of an outage is an error; a permanent one ends the loop + # transient refusals repeat on every retry; only the first of an outage is an error repeated = state.connect_failing and reason_code not in PERMANENT_CONNACK_CODES state.connect_failing = True (logger.debug if repeated else logger.error)( @@ -341,16 +329,12 @@ def mqtt_on_connect( reason_code, ) if reason_code in PERMANENT_CONNACK_CODES: - # paho would otherwise reconnect with the same credentials forever; - # disconnect() puts it in the disconnecting state, so its loop ends - # after this CONNACK instead of retrying + # paho would retry forever with the same credentials; disconnect() ends its loop client.disconnect() def mqtt_on_connect_fail(client, state: MQTTWriterState): - # paho 1.6.1 invokes this callback with exactly (client, userdata). The - # background loop retries about once per reconnect delay while the broker - # is down, so only the first failure of an outage is an error + # paho 1.6.1 passes (client, userdata); only an outage's first failure is an error if state.connect_failing: logger.debug("MQTT client still failing to connect to broker") else: @@ -363,15 +347,13 @@ def mqtt_on_connect_fail(client, state: MQTTWriterState): def mqtt_on_disconnect(client, state: MQTTWriterState, reason_code, properties=None): - # paho also ends every refused background retry here; only the loss of an - # established session is worth an info line + # refused retries end here too; only losing an established session merits info (logger.info if state.connected.is_set() else logger.debug)( "MQTT client disconnected, result code %s", reason_code ) state.connected.clear() if state.refused_code is None: - # a transport drop: the next run waits for the reconnect's CONNACK; after - # a refusal the answer stays visible so run() reports it without waiting + # after a transport drop the next run waits for CONNACK; a refusal stays visible state.connack.clear() @@ -445,7 +427,10 @@ class BlockManifest(WorkflowBlockManifest): description="Message to be published.", examples=["Hello, MQTT!", "$inputs.mqtt_message"], ) - qos: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + qos: Union[ + Annotated[int, Field(ge=0, le=2)], + Selector(kind=[INTEGER_KIND]), + ] = Field( default=0, description="Quality of Service level for the message.", examples=[0, 1, 2], @@ -464,12 +449,12 @@ class BlockManifest(WorkflowBlockManifest): "Must be a finite number greater than 0.", examples=[0.5], ) - username: Union[Selector(kind=[STRING_KIND]), str] = Field( + username: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( default=None, description="Username for MQTT broker authentication.", examples=["$inputs.mqtt_username"], ) - password: Union[Selector(kind=[STRING_KIND]), str] = Field( + password: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( default=None, description="Password for MQTT broker authentication.", examples=["$inputs.mqtt_password"], @@ -537,8 +522,7 @@ def get_actual_restrictions( # the delivery caveat exists only when the run does not wait for the broker declared: Union[List[RuntimeRestriction], Discovery[RuntimeRestriction]] if is_workflow_selector(self.fire_and_forget): - # a runtime value decides whether delivery is awaited, so the caveat - # MAY apply: declare nothing complete instead of guessing either way + # a runtime value decides, so the caveat MAY apply: declare nothing complete declared = incomplete_discovery( items=[], reasons=[ @@ -571,8 +555,7 @@ class MQTTWriterSinkBlockV2(WorkflowBlock): def __init__( self, disable_sinks: bool = False, allow_access_to_file_system: bool = False ): - # gates `ca_certificate_path`, a server-side path chosen by the workflow; - # False by default so a hand-constructed block is safe + # gates `ca_certificate_path`, a server-side path chosen by the workflow self._allow_access_to_file_system = allow_access_to_file_system self.mqtt_client: Optional[mqtt.Client] = None self._connection = MQTTWriterState() @@ -604,8 +587,7 @@ def close(self) -> None: try: loop_stopped = _stop_network_loop(client) finally: - # only after the join can no callback re-set the events; a - # client rebuilt by a later run starts without a stale refusal + # reset only after the join, or a callback could re-set stale events if loop_stopped: self._connection.reset() @@ -665,10 +647,7 @@ def run( fail_fast=fail_fast, ) timeout = timeout_seconds - # selector-supplied ports arrive uncoerced (numeric strings, floats - # from JSON), so coerce like the manifest would before range-checking; - # paho validates only port <= 0 and above 65535 the socket call raises - # OverflowError inside the network loop, silently killing it + # selector ports arrive uncoerced; >65535 would kill paho's network loop try: port_number = _coerce_integer(port) except ValidationError: @@ -720,8 +699,7 @@ def run( fail_fast=fail_fast, ) try: - # the operator's broker policy; the RESOLVED address is what the - # block connects to and what the connection identity is built from + # the resolved address is what the connection identity is built from host, port = resolve_broker_address( host, port, log_override=self.mqtt_client is None ) @@ -777,8 +755,7 @@ def _connect_and_publish( client.max_queued_messages_set(MAX_QUEUED_MESSAGES) if username is not None: client.username_pw_set(username, password) - # TLS must be configured before connect(); the CA path is gated - # by the engine's file-system permission + # TLS must be configured before connect() configure_tls( client, use_tls=encryption == "tls", @@ -788,22 +765,14 @@ def _connect_and_publish( client.on_connect = mqtt_on_connect client.on_connect_fail = mqtt_on_connect_fail client.on_disconnect = mqtt_on_disconnect - # min_delay stays below the readiness wait (= timeout) so the - # first reconnect attempt after a connection drop can finish - # within a single run's wait instead of starting as it expires; - # the floors keep a tiny timeout from turning the background - # retries against a dead broker into a busy loop + # min_delay below timeout lets a reconnect finish within one run's wait client.reconnect_delay_set( min_delay=max(timeout / 2, MIN_RECONNECT_DELAY_SECONDS), max_delay=max(2 * timeout, MIN_MAX_RECONNECT_DELAY_SECONDS), ) - # paho 1.6.1 has no public setter for its synchronous connect - # timeout; without this the TCP phase runs under paho's 5s - # default instead of the block's timeout + # no public setter in paho 1.6.1; else TCP connect uses its 5 s default client._connect_timeout = timeout - # the TCP connect happens here, bounded by _connect_timeout - # (DNS resolution is not - it runs under the OS resolver - # timeout); the CONNACK wait below covers the handshake rest + # DNS is not bounded by _connect_timeout; the OS resolver governs it client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) client.loop_start() except ConfigurationError as e: @@ -822,9 +791,7 @@ def _connect_and_publish( retain=retain, fail_fast=fail_fast, ) - # broker unreachable: nothing is kept and no loop was started, - # so a failed one-shot run leaves no background thread behind; - # the next run on this instance retries with a fresh client + # nothing kept, no loop started; next run retries with a fresh client return self._handle_failure(error, fail_fast=fail_fast) except Exception as e: if client is not None: @@ -845,8 +812,7 @@ def _connect_and_publish( fail_fast=fail_fast, ) state = self._connection - # a permanent refusal stopped the network loop: report it on every run - # without waiting and without touching the broker again + # permanent refusal stopped the loop: report each run, skip the broker if state.refused_code in PERMANENT_CONNACK_CODES: return self._report_failure( connection_refused_message( @@ -855,17 +821,13 @@ def _connect_and_publish( fail_fast=fail_fast, fire_and_forget=fire_and_forget, ) - # fire and forget waits for the broker only on the run that built the - # client; afterwards a disconnected run returns at once and leaves the - # reconnect to the background loop + # fire-and-forget waits only on the first run; later ones leave it to paho waits_for_broker = created_client or not fire_and_forget if not waits_for_broker and qos == 0 and not state.connected.is_set(): return self._fire_and_forget_outage( FIRE_AND_FORGET_QOS0_DROPPED, fail_fast=fail_fast ) - # the background loop owns (re)connecting; runs only wait for readiness. - # The wait ends on any CONNACK, so a refusal is reported as soon as the - # broker answers instead of after the timeout + # the wait ends on any CONNACK, so a refusal is reported once the broker answers if waits_for_broker and not state.connected.is_set(): not_ready = self._wait_until_connected(timeout) if not_ready is not None: @@ -885,10 +847,7 @@ def _connect_and_publish( fail_fast=fail_fast, ) connected = state.connected.is_set() - # max_queued_messages_set() bounds the QoS 1/2 messages, not the encoded - # packets paho buffers while the socket cannot drain them (a broker - # slower than the stream); 1.6.1 exposes that buffer only privately. - # Checked for every QoS, so runs alternating QoS cannot exceed it + # paho 1.6.1's private packet backlog needs its own cap, at every QoS if len(self.mqtt_client._out_packet) >= MAX_QUEUED_MESSAGES: if fire_and_forget: return self._fire_and_forget_outage(QUEUE_FULL, fail_fast=fail_fast) @@ -947,6 +906,7 @@ def _connect_and_publish( def _keep_retrying_after_failed_connect( self, client: mqtt.Client, + *, connection_identity: Tuple, error: str, topic: str, @@ -955,10 +915,7 @@ def _keep_retrying_after_failed_connect( retain: bool, fail_fast: bool, ) -> BlockResult: - # keep the client and let paho's loop retry the connection in the - # background; otherwise every later run would block on a fresh - # synchronous connect for as long as the broker is down. This run - # logs the failure, so the loop's retries do not log it again + # keep the client so paho retries in the background; runs must not block self._connection.connect_failing = True try: client.loop_start() @@ -1011,14 +968,14 @@ def _wait_until_connected(self, timeout: float) -> Optional[str]: def _report_not_connected( self, error: str, + *, topic: str, message: str, qos: int, retain: bool, fail_fast: bool, ) -> BlockResult: - # fire and forget on a run that found no session: QoS 0 has nothing to - # queue, QoS 1/2 is handed to paho for the background connection + # QoS 0 has nothing to queue; QoS 1/2 goes to paho for the background connection if qos == 0: return self._fire_and_forget_outage( f"{error} QoS 0 messages are not queued, so the message was dropped.", @@ -1039,7 +996,7 @@ def _report_not_connected( ) def _fire_and_forget_result( - self, rc: int, qos: int, connected: bool, fail_fast: bool + self, rc: int, *, qos: int, connected: bool, fail_fast: bool ) -> BlockResult: if rc == mqtt.MQTT_ERR_SUCCESS and (connected or qos == 0): if self._fire_and_forget_failure is not None: @@ -1067,21 +1024,20 @@ def _fire_and_forget_result( ) def _report_failure( - self, message: str, fail_fast: bool, fire_and_forget: bool + self, message: str, *, fail_fast: bool, fire_and_forget: bool ) -> BlockResult: if fire_and_forget: return self._fire_and_forget_outage(message, fail_fast=fail_fast) return self._handle_failure(message, fail_fast=fail_fast) - def _fire_and_forget_outage(self, message: str, fail_fast: bool) -> BlockResult: - # without the readiness wait a pipeline reaches this on every frame of - # an outage; log when the failure starts or changes, not once per run + def _fire_and_forget_outage(self, message: str, *, fail_fast: bool) -> BlockResult: + # log when the failure starts or changes, not once per frame of an outage log = message != self._fire_and_forget_failure self._fire_and_forget_failure = message return self._handle_failure(message, fail_fast=fail_fast, log=log) def _handle_failure( - self, message: str, fail_fast: bool, log: bool = True + self, message: str, *, fail_fast: bool, log: bool = True ) -> BlockResult: if log: logger.error("MQTT Writer failure: %s", message) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py index e29fcc1e2c..670737e357 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py @@ -145,6 +145,72 @@ def test_valid_ports_accepted(self, port): assert manifest.port == port + @pytest.mark.parametrize("qos", [-1, 3]) + def test_literal_qos_must_be_zero_to_two(self, qos): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "qos": qos, + } + ) + + @pytest.mark.parametrize("qos", [0, 1, 2, "$inputs.qos"]) + def test_valid_qos_accepted(self, qos): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "qos": qos, + } + ) + + assert manifest.qos == qos + + def test_credentials_may_be_omitted(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + schema = BlockManifest.model_json_schema()["properties"] + + assert manifest.username is None + assert manifest.password is None + for field in ("username", "password"): + assert {"type": "null"} in schema[field]["anyOf"] + + def test_credentials_accept_explicit_null(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v2", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "username": None, + "password": None, + } + ) + + assert manifest.username is None + assert manifest.password is None + class TestCallbacks: def test_on_connect_sets_event_only_for_accepted_connack(self): From 31e300a6f5391d06a5195b7a05e6c98e991f849f Mon Sep 17 00:00:00 2001 From: Grzegorz Klimaszewski <166530809+grzegorz-roboflow@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:03:12 +0200 Subject: [PATCH 5/5] Cover MQTT Writer v2 queue-full, invalid input and close paths with real paho tests --- .../test_workflow_with_mqtt_writer_v2.py | 143 +++++++++++++++++- .../core_steps/sinks/test_mqtt_writer_v2.py | 95 ++++++++++++ 2 files changed, 236 insertions(+), 2 deletions(-) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py index 85529e7a15..05847e7288 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer_v2.py @@ -11,6 +11,8 @@ ) from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v2 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v2 import ( + MAX_QUEUED_MESSAGES, + QUEUE_FULL, MQTTWriterSinkBlockV2, ) @@ -131,9 +133,9 @@ def test_refused_connection_is_reported_once_and_not_retried(): assert "not authorised" in first["message"] assert "Check username and password" in first["message"] assert "Raise 'timeout'" not in first["message"] - assert first_elapsed < 0.4 + assert first_elapsed < kwargs["timeout"] assert second == first - assert second_elapsed < 0.1 + assert second_elapsed < kwargs["timeout"] assert broker.connections_accepted == 1 assert broker.messages == [] finally: @@ -733,6 +735,143 @@ def test_close_from_a_callback_on_the_network_thread_lets_the_loop_end_itself(): } +@pytest.mark.timeout(60) +def test_fire_and_forget_reports_a_full_client_queue_once(): + broker = FakeMQTTBroker(keep_serving=True) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + kwargs = dict( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + qos=1, + timeout=2.0, + fire_and_forget=True, + ) + + try: + with patch.object(v2, "logger") as mock_logger: + queued = [ + block.run(message=f"queued {i}", **kwargs) + for i in range(MAX_QUEUED_MESSAGES) + ] + client = block.mqtt_client + dropped = [block.run(message=f"dropped {i}", **kwargs) for i in range(4)] + + assert all(result["error_status"] is False for result in queued) + assert all( + result == {"error_status": True, "message": QUEUE_FULL} + for result in dropped + ) + assert len(client._out_messages) == MAX_QUEUED_MESSAGES + mock_logger.error.assert_called_once_with("MQTT Writer failure: %s", QUEUE_FULL) + finally: + block.close() + broker.finish() + + +@pytest.mark.timeout(15) +@pytest.mark.parametrize("fire_and_forget", [False, True]) +@pytest.mark.parametrize("topic", ["a/#", "a/+/b", ""]) +def test_invalid_topic_is_reported_by_the_real_client(topic, fire_and_forget): + broker = FakeMQTTBroker(keep_serving=True) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + + try: + result = block.run( + host=broker.host, + port=broker.port, + topic=topic, + message="never sent", + timeout=2.0, + fire_and_forget=fire_and_forget, + ) + + assert result["error_status"] is True + assert "Failed to publish message" in result["message"] + assert block._connected.is_set() + assert broker.messages == [] + finally: + block.close() + broker.finish() + + +@pytest.mark.timeout(15) +def test_close_waits_for_a_run_waiting_for_the_broker_and_stops_the_client(): + connack_gate = threading.Event() + broker = FakeMQTTBroker(keep_serving=True, connack_gate=connack_gate) + threading.Thread(target=broker.serve, daemon=True).start() + block = MQTTWriterSinkBlockV2() + outcome = {} + entered = threading.Event() + close_started = threading.Event() + close_body_started = threading.Event() + run_finished = threading.Event() + run_finished_when_close_body_started = [] + wait_until_connected = block._wait_until_connected + connect_and_publish = block._connect_and_publish + + def gated_wait_until_connected(timeout): + entered.set() + close_started.wait(timeout=5) + close_body_started.wait(timeout=1) + return wait_until_connected(timeout) + + def marking_connect_and_publish(*args, **kwargs): + try: + return connect_and_publish(*args, **kwargs) + finally: + run_finished.set() + + block._wait_until_connected = gated_wait_until_connected + block._connect_and_publish = marking_connect_and_publish + + def run_block(): + try: + outcome["result"] = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="never sent", + timeout=2.0, + ) + except Exception as e: + outcome["exception"] = e + + worker = threading.Thread(target=run_block) + + try: + worker.start() + assert entered.wait(timeout=5) + client = block.mqtt_client + thread = client._thread + disconnect = client.disconnect + + def snapshotting_disconnect(*args, **kwargs): + run_finished_when_close_body_started.append(run_finished.is_set()) + close_body_started.set() + return disconnect(*args, **kwargs) + + client.disconnect = snapshotting_disconnect + + close_started.set() + block.close() + worker.join(timeout=5) + + assert "exception" not in outcome + assert outcome["result"]["error_status"] is True + assert run_finished_when_close_body_started == [True] + assert not thread.is_alive() + assert block.mqtt_client is None + finally: + close_started.set() + connack_gate.set() + block.close() + broker.finish() + worker.join(timeout=5) + + @pytest.fixture def enterprise_blocks_enabled(monkeypatch): from inference.core.env import ENTERPRISE_BLOCKS_PLUGIN diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py index 670737e357..f10dbc9319 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer_v2.py @@ -34,6 +34,9 @@ ) +PAHO_CLIENT_CLASS = mqtt.Client + + def run_kwargs(**overrides) -> dict: kwargs = { "host": "localhost", @@ -977,6 +980,19 @@ def test_changed_timeout_rejected(self, mock_client_cls, block): mock_client.publish.assert_called_once() +def use_paho_publish(mock_client_cls) -> None: + paho_client = PAHO_CLIENT_CLASS() + paho_client._sock = MagicMock() + paho_client._thread = MagicMock() + + def publish(*args, **kwargs): + info = paho_client.publish(*args, **kwargs) + info._set_as_published() + return info + + mock_client_cls.return_value.publish.side_effect = publish + + class TestPublishing: def test_successful_publish(self, mock_client_cls, block): block._connected.set() @@ -1072,6 +1088,58 @@ def test_publish_exception_returned_as_error(self, mock_client_cls, block): assert result["error_status"] is True assert "Invalid topic" in result["message"] + @pytest.mark.parametrize("fire_and_forget", [False, True]) + @pytest.mark.parametrize("topic", ["a/#", "a/+/b", ""]) + def test_invalid_topic_is_reported_as_failed_publish( + self, mock_client_cls, block, topic, fire_and_forget + ): + use_paho_publish(mock_client_cls) + block._connected.set() + block.run(**run_kwargs(fire_and_forget=fire_and_forget)) + + with patch.object(v2, "logger") as mock_logger: + result = block.run( + **run_kwargs(topic=topic, fire_and_forget=fire_and_forget) + ) + + assert result["error_status"] is True + assert "Failed to publish message" in result["message"] + assert mock_logger.error.call_count == 1 + + @pytest.mark.parametrize("fire_and_forget", [False, True]) + @pytest.mark.parametrize("message", [5, None]) + def test_non_string_message_paho_converts_is_published( + self, mock_client_cls, block, message, fire_and_forget + ): + use_paho_publish(mock_client_cls) + block._connected.set() + + result = block.run( + **run_kwargs(message=message, fire_and_forget=fire_and_forget) + ) + + assert result["error_status"] is False + mock_client_cls.return_value.publish.assert_called_once_with( + "test/topic", message, qos=0, retain=False + ) + + @pytest.mark.parametrize("fire_and_forget", [False, True]) + def test_dict_message_is_reported_as_failed_publish( + self, mock_client_cls, block, fire_and_forget + ): + use_paho_publish(mock_client_cls) + block._connected.set() + + result = block.run( + **run_kwargs(message={"a": 1}, fire_and_forget=fire_and_forget) + ) + + assert result["error_status"] is True + assert result["message"] == ( + "Failed to publish message: payload must be a string, bytearray, " + "int, float or None." + ) + def test_no_raw_prints_in_any_code_path(self, mock_client_cls, block): with patch("builtins.print") as mock_print: block._connected.set() @@ -1526,6 +1594,22 @@ def test_selector_values_meaning_true_switch_it_on( assert result["message"] == FIRE_AND_FORGET_QUEUED + def test_qos_change_between_runs_keeps_the_client(self, mock_client_cls, block): + block._connected.set() + mock_client = mock_client_cls.return_value + + results = [ + block.run(**run_kwargs(qos=qos, fire_and_forget=True)) for qos in (0, 1, 2) + ] + + assert all(result["error_status"] is False for result in results) + assert mock_client_cls.call_count == 1 + assert [call.kwargs["qos"] for call in mock_client.publish.call_args_list] == [ + 0, + 1, + 2, + ] + class TestSelectorSwitches: """The engine hands run() the raw selector value, not the manifest's @@ -1605,6 +1689,17 @@ def test_close_stops_loop_even_when_disconnect_raises(self, mock_client_cls, blo mock_client.loop_stop.assert_called_once() + def test_del_swallows_a_failing_close_and_logs_it(self, block): + with patch.object( + MQTTWriterSinkBlockV2, "close", side_effect=RuntimeError("boom") + ): + with patch.object(v2, "logger") as mock_logger: + block.__del__() + + mock_logger.error.assert_called_once() + assert mock_logger.error.call_args.args[0] == "Failed to close MQTT client: %s" + assert str(mock_logger.error.call_args.args[1]) == "boom" + def test_close_on_uninitialized_block_is_noop(self, block): block.close()