fix(redis): reject subscription close signals when receiving (#42799)

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

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:

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
 

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(

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],