From c7e3c9cc4175331b40936b15371cdb7947f28db8 Mon Sep 17 00:00:00 2001 From: Saksham Kapoor Date: Sun, 18 Oct 2026 02:15:25 -0400 Subject: [PATCH] [v3-4-test] Drain result_queue before dispatch in LocalExecutor to prevent runtime dispatch deadlock (#70545) Fixes the runtime/dispatch-path deadlock reported in apache/airflow#70526: during _process_workloads the scheduler could block in activity_queue.put() while workers were blocked writing results into a full result_queue pipe, a circular wait. Draining the result queue before each put breaks the cycle. The end()/join path is already handled by #67881; this covers normal dispatch. Adds a call-order unit test and a bounded real-queue regression test. (cherry picked from commit b00cfeb11f0f7f45f03317a57cb4364f33c09f49) Co-authored-by: Saksham Kapoor --- .../src/airflow/executors/local_executor.py | 1 + .../unit/executors/test_local_executor.py | 80 +++++++++++++++++++ 2 files changed, 81 insertions(+) diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 9e1f1f8bbf8dc..aebadbf66ad72 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -337,6 +337,7 @@ def _terminate_worker_process(self, proc: multiprocessing.Process) -> None: def _process_workloads(self, workload_list): for workload in workload_list: key = self.get_workload_key(workload) + self._read_results() self.activity_queue.put(workload) removed = self.executor_queues[workload.type].pop(key, None) if not removed: diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 9f5d196d6cd94..0b4916c139076 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -21,6 +21,7 @@ import multiprocessing import os import signal +import threading import time from pathlib import Path from unittest import mock @@ -106,6 +107,17 @@ def _write_large_results_to_queue(result_queue, activity_queue, unread_messages, result_queue.put((os.getpid(), key, State.SUCCESS, payload)) +def _refill_worker(activity_queue, result_queue, barrier, count, payload_size): + """Stand-in worker: consume activity items and emit large results over real pipes.""" + payload = b"x" * payload_size + barrier.wait() + for _ in range(count): + item = activity_queue.get() + if item is None: + return + result_queue.put((os.getpid(), None, None, payload)) + + def _make_workload(kind): if kind == "task": return _make_task_workload() @@ -443,6 +455,74 @@ def test_end_drains_result_queue_to_avoid_join_deadlock(self, mocker): assert not executor._worker_tasks assert executor._unread_messages.value == 0 + @pytest.mark.parametrize("num_workloads", [1, 3]) + def test_process_workloads_drains_results_before_put(self, num_workloads): + """#70545: _process_workloads must drain the result queue before each activity_queue.put.""" + executor = LocalExecutor(parallelism=1) + executor.activity_queue = mock.MagicMock() + executor._read_results = mock.MagicMock() + executor._unread_messages = multiprocessing.Value("I", 0) + executor._check_workers = mock.MagicMock() + + call_order = [] + executor._read_results.side_effect = lambda: call_order.append("_read_results") + executor.activity_queue.put.side_effect = lambda *args, **kwargs: call_order.append("put") + + workloads_list = [] + for _ in range(num_workloads): + workload = _make_task_workload() + key = executor.get_workload_key(workload) + executor.executor_queues.setdefault(workload.type, {})[key] = workload + workloads_list.append(workload) + + executor._process_workloads(workloads_list) + + assert call_order == ["_read_results", "put"] * num_workloads + + def test_process_workloads_progresses_under_concurrent_result_refill(self, mocker): + """#70545: with real queues, _process_workloads must keep draining so a concurrently + refilling worker cannot deadlock the dispatch loop via a full result_queue pipe.""" + ctx = multiprocessing.get_context("fork") + executor = LocalExecutor(parallelism=1) + executor.activity_queue = ctx.SimpleQueue() + executor.result_queue = ctx.SimpleQueue() + executor._unread_messages = multiprocessing.Value("I", 0) + executor._worker_tasks = {} + executor._dispatch_counts = {} + mocker.patch.object(executor, "_check_workers", autospec=True) + + count = 200 + submitted = [_make_task_workload() for _ in range(count)] + for workload in submitted: + executor.executor_queues.setdefault(workload.type, {})[executor.get_workload_key(workload)] = ( + workload + ) + + barrier = ctx.Barrier(2) + proc = ctx.Process( + target=_refill_worker, + args=(executor.activity_queue, executor.result_queue, barrier, count, 16 * 1024), + ) + proc.start() + executor.workers = {proc.pid: proc} + barrier.wait(timeout=30) + + completed = threading.Event() + + def _run(): + try: + executor._process_workloads(submitted) + finally: + completed.set() + + runner = threading.Thread(target=_run, daemon=True) + runner.start() + finished = completed.wait(timeout=30) + proc.terminate() + proc.join(timeout=5) + + assert finished, "dispatch did not complete within 30s (result-queue refill deadlock)" + @pytest.mark.parametrize( ("conf_values", "expected_server"), [