From 39df1adcac119f9abd3004c6972f3c19d512ea3c Mon Sep 17 00:00:00 2001 From: zifeif2 Date: Thu, 17 Sep 2026 17:22:08 +0000 Subject: [PATCH 1/5] [CONNECT][PYTHON] Fix streaming query listener removal deadlock --- python/pyspark/sql/connect/streaming/query.py | 56 +++++--- .../connect/streaming/test_parity_listener.py | 124 +++++++++++++++++- 2 files changed, 158 insertions(+), 22 deletions(-) diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py index ea1a7b07bfd6a..285d312115621 100644 --- a/python/pyspark/sql/connect/streaming/query.py +++ b/python/pyspark/sql/connect/streaming/query.py @@ -279,6 +279,9 @@ def __init__(self, sqm: "StreamingQueryManager") -> None: self._listener_bus: List[StreamingQueryListener] = [] self._execution_thread: Optional[Thread] = None self._lock = Lock() + # Serialize listener lifecycle changes while allowing the event thread to acquire + # _lock and drain pending events when the last listener is removed. + self._lifecycle_lock = Lock() def close(self) -> None: for listener in self._listener_bus: @@ -290,7 +293,7 @@ def append(self, listener: StreamingQueryListener) -> None: the first listener, request the server to create the server side listener and start a thread to handle query events. """ - with self._lock: + with self._lifecycle_lock, self._lock: self._listener_bus.append(listener) if len(self._listener_bus) == 1: @@ -321,28 +324,39 @@ def remove(self, listener: StreamingQueryListener) -> None: will return after processing remaining listener events. This function blocks until all events are processed. """ - with self._lock: - if listener not in self._listener_bus: - return - - if len(self._listener_bus) == 1: - cmd = pb2.StreamingQueryListenerBusCommand() - cmd.remove_listener_bus_listener = True - exec_cmd = pb2.Command() - exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd) - try: - self._sqm._session.client.execute_command(exec_cmd) - except Exception as e: - warnings.warn( - f"Failed to remove the listener because of exception: {e}\n" - f"The listener is not removed, please remove it again." - ) + with self._lifecycle_lock: + with self._lock: + if listener not in self._listener_bus: return - if self._execution_thread is not None: - self._execution_thread.join() - self._execution_thread = None - self._listener_bus.remove(listener) + is_last_listener = len(self._listener_bus) == 1 + execution_thread = None + if is_last_listener: + cmd = pb2.StreamingQueryListenerBusCommand() + cmd.remove_listener_bus_listener = True + exec_cmd = pb2.Command() + exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd) + try: + self._sqm._session.client.execute_command(exec_cmd) + except Exception as e: + warnings.warn( + f"Failed to remove the listener because of exception: {e}\n" + f"The listener is not removed, please remove it again." + ) + return + execution_thread = self._execution_thread + else: + self._listener_bus.remove(listener) + + if is_last_listener: + if execution_thread is not None: + execution_thread.join() + with self._lock: + if self._execution_thread is execution_thread: + self._execution_thread = None + # The event thread may have cleared the listener bus after an exception. + if listener in self._listener_bus: + self._listener_bus.remove(listener) @staticmethod def _iter_listener_events( diff --git a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py index 81a116eff2285..2bd9978e1c20e 100644 --- a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py +++ b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py @@ -15,12 +15,17 @@ # limitations under the License. # +import threading import time +import unittest +import uuid +from unittest.mock import MagicMock import pyspark.cloudpickle from pyspark.errors import AnalysisException +from pyspark.sql.connect.streaming.query import StreamingQueryListenerBus from pyspark.sql.functions import count, lit -from pyspark.sql.streaming.listener import StreamingQueryListener +from pyspark.sql.streaming.listener import QueryStartedEvent, StreamingQueryListener from pyspark.sql.tests.streaming.test_streaming_listener import StreamingListenerTestsMixin from pyspark.testing.connectutils import ReusedConnectTestCase from pyspark.testing.utils import eventually @@ -84,6 +89,123 @@ def onQueryTerminated(self, event): self.terminated.append(event) +class StreamingQueryListenerBusTests(unittest.TestCase): + def test_remove_last_listener_with_pending_event(self): + listener = TestListenerLocalV2() + next_listener = TestListenerLocalV2() + + client = MagicMock() + sqm = MagicMock() + sqm._session.client = client + listener_bus = StreamingQueryListenerBus(sqm) + listener_bus._listener_bus.append(listener) + + remove_command_started = threading.Event() + event_dispatch_started = threading.Event() + event_dispatched = threading.Event() + event_thread_can_exit = threading.Event() + removal_finished = threading.Event() + append_finished = threading.Event() + lifecycle_wait_started = threading.Event() + thread_errors = [] + threads = [] + + class TrackingLifecycleLock: + def __init__(self): + self._lock = threading.Lock() + + def __enter__(self): + if self._lock.locked(): + lifecycle_wait_started.set() + self._lock.acquire() + + def __exit__(self, exc_type, exc_value, traceback): + self._lock.release() + + listener_bus._lifecycle_lock = TrackingLifecycleLock() + + def execute_command(_): + remove_command_started.set() + if not event_dispatch_started.wait(5): + raise TimeoutError("event dispatch did not start") + + client.execute_command.side_effect = execute_command + + event = QueryStartedEvent( + id=uuid.uuid4(), + runId=uuid.uuid4(), + name="pending-event", + timestamp="2026-09-17T00:00:00.000Z", + jobTags=set(), + ) + + def dispatch_pending_event(): + try: + if not remove_command_started.wait(5): + raise TimeoutError("listener removal did not start") + event_dispatch_started.set() + listener_bus.post_to_all(event) + event_dispatched.set() + event_thread_can_exit.wait() + except BaseException as error: + thread_errors.append(error) + + def remove_listener(): + try: + listener_bus.remove(listener) + except BaseException as error: + thread_errors.append(error) + finally: + removal_finished.set() + + register_server_side_listener = MagicMock(return_value=iter(())) + listener_bus._register_server_side_listener = register_server_side_listener + + def append_listener(): + try: + listener_bus.append(next_listener) + except BaseException as error: + thread_errors.append(error) + finally: + append_finished.set() + + event_thread = threading.Thread(target=dispatch_pending_event, daemon=True) + listener_bus._execution_thread = event_thread + removal_thread = threading.Thread(target=remove_listener, daemon=True) + threads.extend([event_thread, removal_thread]) + + try: + event_thread.start() + removal_thread.start() + + self.assertTrue( + event_dispatched.wait(5), + "removeListener blocked the event thread while waiting for it to exit", + ) + + append_thread = threading.Thread(target=append_listener, daemon=True) + threads.append(append_thread) + append_thread.start() + self.assertTrue( + lifecycle_wait_started.wait(5), + "addListener did not wait for the listener shutdown", + ) + self.assertFalse(append_finished.is_set()) + + event_thread_can_exit.set() + self.assertTrue(removal_finished.wait(5)) + self.assertTrue(append_finished.wait(5)) + + self.assertEqual(thread_errors, []) + self.assertEqual(listener.start, [event]) + self.assertEqual(listener_bus._listener_bus, [next_listener]) + register_server_side_listener.assert_called_once_with() + finally: + event_thread_can_exit.set() + for thread in threads: + thread.join(timeout=1) + + class StreamingListenerParityTests(StreamingListenerTestsMixin, ReusedConnectTestCase): def test_listener_management(self): listener1 = TestListenerLocalV1() From 43787b1e7c5a626dcc5d7241e83a9399c5eb404c Mon Sep 17 00:00:00 2001 From: zifeif2 Date: Thu, 17 Sep 2026 22:25:56 +0000 Subject: [PATCH 2/5] [CONNECT][PYTHON] Document listener deadlock test ordering --- .../sql/tests/connect/streaming/test_parity_listener.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py index 2bd9978e1c20e..e0f88c6a329aa 100644 --- a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py +++ b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py @@ -100,6 +100,11 @@ def test_remove_last_listener_with_pending_event(self): listener_bus = StreamingQueryListenerBus(sqm) listener_bus._listener_bus.append(listener) + # Reproduce the ordering that used to deadlock: removal holds _lock while + # requesting server-side shutdown, and the event thread starts dispatching a + # pending event that also needs _lock. If removal keeps _lock while joining the + # event thread, each thread waits for the other. After dispatch, keep the event + # thread alive long enough to verify that append waits for shutdown to finish. remove_command_started = threading.Event() event_dispatch_started = threading.Event() event_dispatched = threading.Event() From df461b70fa0f58091c09a425ecf535c3dd6bd430 Mon Sep 17 00:00:00 2001 From: zifeif2 Date: Fri, 18 Sep 2026 21:10:01 +0000 Subject: [PATCH 3/5] [CONNECT][PYTHON] Clarify listener state lock name --- python/pyspark/sql/connect/streaming/query.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py index 285d312115621..7d0b54014e69f 100644 --- a/python/pyspark/sql/connect/streaming/query.py +++ b/python/pyspark/sql/connect/streaming/query.py @@ -278,9 +278,10 @@ def __init__(self, sqm: "StreamingQueryManager") -> None: self._sqm = sqm self._listener_bus: List[StreamingQueryListener] = [] self._execution_thread: Optional[Thread] = None - self._lock = Lock() + # Protects _listener_bus and _execution_thread shared by API callers and the event thread. + self._listeners_state_lock = Lock() # Serialize listener lifecycle changes while allowing the event thread to acquire - # _lock and drain pending events when the last listener is removed. + # _listeners_state_lock and drain pending events when the last listener is removed. self._lifecycle_lock = Lock() def close(self) -> None: @@ -293,7 +294,7 @@ def append(self, listener: StreamingQueryListener) -> None: the first listener, request the server to create the server side listener and start a thread to handle query events. """ - with self._lifecycle_lock, self._lock: + with self._lifecycle_lock, self._listeners_state_lock: self._listener_bus.append(listener) if len(self._listener_bus) == 1: @@ -325,7 +326,7 @@ def remove(self, listener: StreamingQueryListener) -> None: all events are processed. """ with self._lifecycle_lock: - with self._lock: + with self._listeners_state_lock: if listener not in self._listener_bus: return @@ -351,7 +352,7 @@ def remove(self, listener: StreamingQueryListener) -> None: if is_last_listener: if execution_thread is not None: execution_thread.join() - with self._lock: + with self._listeners_state_lock: if self._execution_thread is execution_thread: self._execution_thread = None # The event thread may have cleared the listener bus after an exception. @@ -414,7 +415,7 @@ def _query_event_handler(self, iter: Iterator["_ExecutePlanResponseItem"]) -> No "StreamingQueryListenerBus Handler thread received exception, all client side " f"listeners are removed and handler thread is terminated. The error is: {e}" ) - with self._lock: + with self._listeners_state_lock: self._execution_thread = None self._listener_bus.clear() return @@ -445,7 +446,7 @@ def post_to_all( Post listener events to all active listeners, note that if one listener throws, it should not affect other listeners. """ - with self._lock: + with self._listeners_state_lock: for listener in self._listener_bus: try: if isinstance(event, QueryStartedEvent): From a490cb824f2504713e10c20847e14a8400e253dc Mon Sep 17 00:00:00 2001 From: zifeif2 Date: Mon, 21 Sep 2026 17:08:10 +0000 Subject: [PATCH 4/5] [CONNECT][PYTHON] Avoid holding listener state lock during RPC --- python/pyspark/sql/connect/streaming/query.py | 51 +++++++++---------- .../connect/streaming/test_parity_listener.py | 16 +++--- 2 files changed, 32 insertions(+), 35 deletions(-) diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py index 7d0b54014e69f..bcbd08a9d7a36 100644 --- a/python/pyspark/sql/connect/streaming/query.py +++ b/python/pyspark/sql/connect/streaming/query.py @@ -330,34 +330,33 @@ def remove(self, listener: StreamingQueryListener) -> None: if listener not in self._listener_bus: return - is_last_listener = len(self._listener_bus) == 1 - execution_thread = None - if is_last_listener: - cmd = pb2.StreamingQueryListenerBusCommand() - cmd.remove_listener_bus_listener = True - exec_cmd = pb2.Command() - exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd) - try: - self._sqm._session.client.execute_command(exec_cmd) - except Exception as e: - warnings.warn( - f"Failed to remove the listener because of exception: {e}\n" - f"The listener is not removed, please remove it again." - ) - return - execution_thread = self._execution_thread - else: + if len(self._listener_bus) != 1: self._listener_bus.remove(listener) + return + + execution_thread = self._execution_thread + + cmd = pb2.StreamingQueryListenerBusCommand() + cmd.remove_listener_bus_listener = True + exec_cmd = pb2.Command() + exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd) + try: + self._sqm._session.client.execute_command(exec_cmd) + except Exception as e: + warnings.warn( + f"Failed to remove the listener because of exception: {e}\n" + f"The listener is not removed, please remove it again." + ) + return - if is_last_listener: - if execution_thread is not None: - execution_thread.join() - with self._listeners_state_lock: - if self._execution_thread is execution_thread: - self._execution_thread = None - # The event thread may have cleared the listener bus after an exception. - if listener in self._listener_bus: - self._listener_bus.remove(listener) + if execution_thread is not None: + execution_thread.join() + with self._listeners_state_lock: + if self._execution_thread is execution_thread: + self._execution_thread = None + # The event thread may have cleared the listener bus after an exception. + if listener in self._listener_bus: + self._listener_bus.remove(listener) @staticmethod def _iter_listener_events( diff --git a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py index e0f88c6a329aa..f8b1290ce771c 100644 --- a/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py +++ b/python/pyspark/sql/tests/connect/streaming/test_parity_listener.py @@ -100,13 +100,12 @@ def test_remove_last_listener_with_pending_event(self): listener_bus = StreamingQueryListenerBus(sqm) listener_bus._listener_bus.append(listener) - # Reproduce the ordering that used to deadlock: removal holds _lock while - # requesting server-side shutdown, and the event thread starts dispatching a - # pending event that also needs _lock. If removal keeps _lock while joining the - # event thread, each thread waits for the other. After dispatch, keep the event - # thread alive long enough to verify that append waits for shutdown to finish. + # Reproduce the ordering that used to deadlock: removal requests server-side + # shutdown while the event thread dispatches a pending event. The request must + # not hold the listener state lock, which the event thread also needs. After + # dispatch, keep the event thread alive long enough to verify that append waits + # for shutdown to finish. remove_command_started = threading.Event() - event_dispatch_started = threading.Event() event_dispatched = threading.Event() event_thread_can_exit = threading.Event() removal_finished = threading.Event() @@ -131,8 +130,8 @@ def __exit__(self, exc_type, exc_value, traceback): def execute_command(_): remove_command_started.set() - if not event_dispatch_started.wait(5): - raise TimeoutError("event dispatch did not start") + if not event_dispatched.wait(5): + raise TimeoutError("pending event was not dispatched") client.execute_command.side_effect = execute_command @@ -148,7 +147,6 @@ def dispatch_pending_event(): try: if not remove_command_started.wait(5): raise TimeoutError("listener removal did not start") - event_dispatch_started.set() listener_bus.post_to_all(event) event_dispatched.set() event_thread_can_exit.wait() From bde1847b7d2fcc169ce3cc854c02c8cc2cf9a2d2 Mon Sep 17 00:00:00 2001 From: zifeif2 Date: Mon, 21 Sep 2026 17:23:14 +0000 Subject: [PATCH 5/5] [CONNECT][PYTHON] Preserve listener removal branch structure --- python/pyspark/sql/connect/streaming/query.py | 50 ++++++++++--------- 1 file changed, 26 insertions(+), 24 deletions(-) diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py index bcbd08a9d7a36..0f3700f125d95 100644 --- a/python/pyspark/sql/connect/streaming/query.py +++ b/python/pyspark/sql/connect/streaming/query.py @@ -330,33 +330,35 @@ def remove(self, listener: StreamingQueryListener) -> None: if listener not in self._listener_bus: return - if len(self._listener_bus) != 1: + is_last_listener = len(self._listener_bus) == 1 + execution_thread = None + if is_last_listener: + execution_thread = self._execution_thread + else: self._listener_bus.remove(listener) - return - execution_thread = self._execution_thread - - cmd = pb2.StreamingQueryListenerBusCommand() - cmd.remove_listener_bus_listener = True - exec_cmd = pb2.Command() - exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd) - try: - self._sqm._session.client.execute_command(exec_cmd) - except Exception as e: - warnings.warn( - f"Failed to remove the listener because of exception: {e}\n" - f"The listener is not removed, please remove it again." - ) - return + if is_last_listener: + cmd = pb2.StreamingQueryListenerBusCommand() + cmd.remove_listener_bus_listener = True + exec_cmd = pb2.Command() + exec_cmd.streaming_query_listener_bus_command.CopyFrom(cmd) + try: + self._sqm._session.client.execute_command(exec_cmd) + except Exception as e: + warnings.warn( + f"Failed to remove the listener because of exception: {e}\n" + f"The listener is not removed, please remove it again." + ) + return - if execution_thread is not None: - execution_thread.join() - with self._listeners_state_lock: - if self._execution_thread is execution_thread: - self._execution_thread = None - # The event thread may have cleared the listener bus after an exception. - if listener in self._listener_bus: - self._listener_bus.remove(listener) + if execution_thread is not None: + execution_thread.join() + with self._listeners_state_lock: + if self._execution_thread is execution_thread: + self._execution_thread = None + # The event thread may have cleared the listener bus after an exception. + if listener in self._listener_bus: + self._listener_bus.remove(listener) @staticmethod def _iter_listener_events(