Drain result_queue in LocalExecutor to prevent process join deadlock

airflow-core/src/airflow/executors/local_executor.py

@@ -25,6 +25,7 @@
 
 from __future__ import annotations
 
+import contextlib
 import ctypes
 import multiprocessing
 import multiprocessing.sharedctypes
@@ -243,10 +244,12 @@ def sync(self) -> None:
         self._check_workers()
 
     def _read_results(self):
-        while not self.result_queue.empty():
-            key, state, exc = self.result_queue.get()
-
-            self.change_state(key, state)
+        try:
+            while not self.result_queue.empty():
+                key, state, exc = self.result_queue.get()
+                self.change_state(key, state)
+        except (OSError, EOFError):
+            self.log.exception("Error reading from result queue")
 
     def end(self) -> None:
         """End the executor."""
@@ -263,19 +266,47 @@ def end(self) -> None:
             if proc.is_alive():
                 self.activity_queue.put(None)
 
-        for proc in self.workers.values():
-            if proc.is_alive():
-                proc.join()
-            proc.close()
+        # To prevent deadlock, we should consume results from result_queue while waiting for processes to join.
+        # Otherwise, a worker blocked on putting results into a full result_queue pipe will never exit,
+        # and an unbounded proc.join() will hang the scheduler indefinitely.
+        try:
+            for proc in self.workers.values():
+                while proc.is_alive():
+                    self._read_results()
+                    proc.join(timeout=0.05)
+        except (KeyboardInterrupt, SystemExit):
+            self.log.error("KeyboardInterrupt received during shutdown. Force terminating workers.")
+            for proc in self.workers.values():
+                self._terminate_worker_process(proc)
+            raise
+        finally:
+            # Process any extra results before closing
+            self._read_results()
 
-        # Process any extra results before closing
-        self._read_results()
+            for proc in self.workers.values():
+                with contextlib.suppress(ValueError):
+                    proc.close()
 
-        self.activity_queue.close()
-        self.result_queue.close()
+            self.activity_queue.close()
+            self.result_queue.close()
 
     def terminate(self):
-        """Terminate the executor is not doing anything."""
+        """Terminate all worker processes under control of the executor forcefully."""
+        self.log.info("Terminating all LocalExecutor worker processes.")
+        for proc in self.workers.values():
+            self._terminate_worker_process(proc)
+
+    def _terminate_worker_process(self, proc: multiprocessing.Process) -> None:
+        """Terminate a worker process, escalating to kill if it stays alive."""
+        if not proc.is_alive():
+            return
+
+        proc.terminate()
+        proc.join(timeout=0.2)
+        if proc.is_alive():
+            self.log.warning("Worker process %s did not stop after SIGTERM. Sending SIGKILL.", proc.pid)
+            proc.kill()
+            proc.join(timeout=0.2)
 
     def _process_workloads(self, workload_list):
         for workload in workload_list:

airflow-core/tests/unit/executors/test_local_executor.py

@@ -35,6 +35,7 @@
 from airflow.executors.workloads.callback import CallbackDTO
 from airflow.executors.workloads.task import TaskInstanceDTO
 from airflow.models.callback import CallbackFetchMethod
+from airflow.models.taskinstancekey import TaskInstanceKey
 from airflow.settings import Session
 from airflow.utils.state import State
 
@@ -90,6 +91,13 @@ def _make_task_workload():
     )
 
 
+def _write_large_results_to_queue(result_queue, result_count, payload_size):
+    payload = RuntimeError("x" * payload_size)
+    for index in range(result_count):
+        key = TaskInstanceKey("test_dag", f"test_task_{index}", "test_run")
+        result_queue.put((key, State.SUCCESS, payload))
+
+
 class TestLocalExecutor:
     """
     When the executor is started, end() must be called before the test finishes.
@@ -274,6 +282,87 @@ def test_clean_stop_on_signal(self):
         finally:
             executor.end()
 
+    def test_end_drains_results_while_joining_workers(self):
+        executor = LocalExecutor(parallelism=1)
+        executor.activity_queue = mock.MagicMock()
+        executor.result_queue = mock.MagicMock()
+        proc = mock.MagicMock(spec=multiprocessing.Process)
+        proc.is_alive.side_effect = [True, True, True, False]
+        executor.workers = {1: proc}
+
+        with mock.patch.object(executor, "_read_results") as mock_read_results:
+            executor.end()
+
+        executor.activity_queue.put.assert_called_once_with(None)
+        assert proc.join.call_args_list == [mock.call(timeout=0.05), mock.call(timeout=0.05)]
+        assert mock_read_results.call_count == 3
+        proc.close.assert_called_once()
+        executor.activity_queue.close.assert_called_once()
+        executor.result_queue.close.assert_called_once()
+
+    def test_end_terminates_workers_and_closes_resources_on_interrupt(self):
+        executor = LocalExecutor(parallelism=1)
+        executor.activity_queue = mock.MagicMock()
+        executor.result_queue = mock.MagicMock()
+        proc = mock.MagicMock(spec=multiprocessing.Process)
+        proc.is_alive.side_effect = [True, True]
+        executor.workers = {1: proc}
+
+        with (
+            mock.patch.object(executor, "_read_results", side_effect=[KeyboardInterrupt, None]),
+            mock.patch.object(executor, "_terminate_worker_process") as mock_terminate_worker_process,
+            pytest.raises(KeyboardInterrupt),
+        ):
+            executor.end()
+
+        mock_terminate_worker_process.assert_called_once_with(proc)
+        proc.close.assert_called_once()
+        executor.activity_queue.close.assert_called_once()
+        executor.result_queue.close.assert_called_once()
+
+    def test_terminate_joins_worker_after_sigterm(self):
+        executor = LocalExecutor(parallelism=1)
+        proc = mock.MagicMock(spec=multiprocessing.Process)
+        proc.is_alive.side_effect = [True, False]
+        executor.workers = {1: proc}
+
+        executor.terminate()
+
+        proc.terminate.assert_called_once_with()
+        proc.join.assert_called_once_with(timeout=0.2)
+        proc.kill.assert_not_called()
+
+    def test_terminate_kills_worker_that_ignores_sigterm(self):
+        executor = LocalExecutor(parallelism=1)
+        proc = mock.MagicMock(spec=multiprocessing.Process)
+        proc.pid = 123
+        proc.is_alive.side_effect = [True, True]
+        executor.workers = {1: proc}
+
+        executor.terminate()
+
+        proc.terminate.assert_called_once_with()
+        proc.kill.assert_called_once_with()
+        assert proc.join.call_args_list == [mock.call(timeout=0.2), mock.call(timeout=0.2)]
+
+    @pytest.mark.execution_timeout(10)
+    def test_end_drains_result_queue_to_avoid_join_deadlock(self):
+        executor = LocalExecutor(parallelism=1)
+        executor.activity_queue = multiprocessing.SimpleQueue()
+        executor.result_queue = multiprocessing.SimpleQueue()
+        result_count = 8
+        payload_size = 128 * 1024
+        proc = multiprocessing.Process(
+            target=_write_large_results_to_queue,
+            args=(executor.result_queue, result_count, payload_size),
+        )
+        proc.start()
+        executor.workers = {proc.pid: proc}
+
+        executor.end()
+
+        assert len(executor.event_buffer) == result_count
+
     @pytest.mark.parametrize(
         ("conf_values", "expected_server"),
         [