[SPARK-59623][SS][PYTHON] Fix streaming query listener removal deadlock

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,28 +325,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._listeners_state_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._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 +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
@@ -431,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):

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,128 @@ 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 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()
+        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()