From 4b89972f552c7a1ad7c4eaaae47dfe669020de94 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Fri, 2 Oct 2026 18:50:56 +0000 Subject: [PATCH] fix(redis): reject subscription close signals when receiving (#42799) --- api/libs/broadcast_channel/redis/_subscription.py | 2 ++ api/services/workflow_event_snapshot_service.py | 3 +++ .../broadcast_channel/redis/test_sharded_channel.py | 6 +++++- .../redis/test_channel_unit_tests.py | 6 ++++++ .../workflow/test_workflow_event_snapshot_service.py | 12 ++++++++++++ 5 files changed, 28 insertions(+), 1 deletion(-) diff --git a/api/libs/broadcast_channel/redis/_subscription.py b/api/libs/broadcast_channel/redis/_subscription.py index c155b9e1973..94af7a784c3 100644 --- a/api/libs/broadcast_channel/redis/_subscription.py +++ b/api/libs/broadcast_channel/redis/_subscription.py @@ -196,6 +196,8 @@ class RedisSubscriptionBase(Subscription): except queue.Empty: return None + if item == SIG_CLOSE: + raise SubscriptionClosedError(f"The Redis {self._get_subscription_type()} subscription is closed") return item @override diff --git a/api/services/workflow_event_snapshot_service.py b/api/services/workflow_event_snapshot_service.py index 5fb86206ec9..b25079b8762 100644 --- a/api/services/workflow_event_snapshot_service.py +++ b/api/services/workflow_event_snapshot_service.py @@ -43,6 +43,7 @@ from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.runtime import GraphRuntimeState from graphon.runtime.graph_runtime_state_protocol import ReadOnlyVariablePool from graphon.workflow_type_encoder import WorkflowRuntimeTypeConverter +from libs.broadcast_channel.exc import SubscriptionClosedError from libs.datetime_utils import to_utc_timestamp from models.human_input import HumanInputForm from models.model import AppMode, Message @@ -669,6 +670,8 @@ def _start_buffering(subscription) -> BufferState: except queue.Full: continue logger.warning("Dropped buffered workflow event, total_dropped=%s", dropped_count) + except SubscriptionClosedError: + pass except Exception: logger.exception("Failed while buffering workflow events") finally: diff --git a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py b/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py index 84c1d0ca41b..7e0ed7cd366 100644 --- a/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py +++ b/api/tests/test_containers_integration_tests/libs/broadcast_channel/redis/test_sharded_channel.py @@ -237,6 +237,7 @@ class TestShardedRedisBroadcastChannelIntegration: def consumer_thread() -> set[bytes]: received_msgs: set[bytes] = set() with subscription: + subscription.receive(timeout=0.1) consumer_ready.set() while True: try: @@ -260,7 +261,10 @@ class TestShardedRedisBroadcastChannelIntegration: for future in as_completed(producer_futures, timeout=30.0): sent_msgs.update(future.result()) - consumer_received_msgs = consumer_future.result(timeout=60.0) + try: + consumer_received_msgs = consumer_future.result(timeout=60.0) + finally: + subscription.close() assert sent_msgs == consumer_received_msgs diff --git a/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py b/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py index 2a9bc002607..a7b32d986c4 100644 --- a/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py +++ b/api/tests/unit_tests/libs/broadcast_channel/redis/test_channel_unit_tests.py @@ -1441,6 +1441,12 @@ class TestRedisSubscriptionCommon: with pytest.raises(SubscriptionClosedError): subscription.receive() + def test_receive_does_not_leak_close_signal(self, subscription, subscription_params): + subscription._queue.put_nowait(SIG_CLOSE) + + with pytest.raises(SubscriptionClosedError): + subscription.receive() + # ==================== Table-driven Tests ==================== @pytest.mark.parametrize( diff --git a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py index 9f914556800..e52ce4ca02d 100644 --- a/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py +++ b/api/tests/unit_tests/services/workflow/test_workflow_event_snapshot_service.py @@ -27,6 +27,7 @@ from core.workflow.nodes.human_input.enums import ValueSourceType from core.workflow.nodes.human_input.pause_reason import HumanInputRequired from graphon.enums import WorkflowExecutionStatus, WorkflowNodeExecutionStatus from graphon.runtime import GraphRuntimeState, VariablePool +from libs.broadcast_channel.exc import SubscriptionClosedError from libs.datetime_utils import to_utc_timestamp from models.enums import ConversationFromSource, CreatorUserRole from models.human_input import HumanInputForm, HumanInputFormRecipient, RecipientType @@ -701,6 +702,17 @@ def test_start_buffering_should_set_done_event_when_subscription_raises() -> Non assert finished is True +def test_start_buffering_should_treat_closed_subscription_as_done(caplog: pytest.LogCaptureFixture) -> None: + class Subscription: + def receive(self, timeout: int = 1) -> bytes | None: + raise SubscriptionClosedError("closed") + + buffer_state = service_module._start_buffering(Subscription()) + + assert buffer_state.done_event.wait(timeout=1) is True + assert "Failed while buffering workflow events" not in caplog.text + + def test_build_workflow_event_stream_should_emit_ping_and_terminal_snapshot_event( monkeypatch: pytest.MonkeyPatch, unbound_session_factory: sessionmaker[Session],