diff --git a/python/pyspark/sql/connect/streaming/query.py b/python/pyspark/sql/connect/streaming/query.py index ea1a7b07bfd6..0f3700f125d9 100644 --- a/python/pyspark/sql/connect/streaming/query.py +++ b/python/pyspark/sql/connect/streaming/query.py @@ -278,7 +278,11 @@ 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 + # _listeners_state_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 +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._lock: + with self._lifecycle_lock, self._listeners_state_lock: self._listener_bus.append(listener) if len(self._listener_bus) == 1: @@ -321,11 +325,19 @@ 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 + with self._lifecycle_lock: + with self._listeners_state_lock: + 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) + + if is_last_listener: cmd = pb2.StreamingQueryListenerBusCommand() cmd.remove_listener_bus_listener = True exec_cmd = pb2.Command() @@ -338,11 +350,15 @@ def remove(self, listener: StreamingQueryListener) -> None: f"The listener is not removed, please remove it again." ) return - if self._execution_thread is not None: - self._execution_thread.join() - self._execution_thread = None - 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( @@ -400,7 +416,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 @@ -431,7 +447,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): 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 81a116eff228..f8b1290ce771 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,126 @@ 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) + + # 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_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_dispatched.wait(5): + raise TimeoutError("pending event was not dispatched") + + 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") + 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()