This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6a73afb1c15 Added deadlock detection in CommsDecoder (#68377)
6a73afb1c15 is described below
commit 6a73afb1c159b4ada036a4df455355d42d72d1b2
Author: David Blain <[email protected]>
AuthorDate: Fri Jul 3 06:50:41 2026 +0200
Added deadlock detection in CommsDecoder (#68377)
* refactor: Fix deadlock in CommsDecoder.asend() when sync send() runs
concurrently on event loop
---
task-sdk/src/airflow/sdk/execution_time/comms.py | 52 ++++++-
task-sdk/tests/task_sdk/execution_time/conftest.py | 12 ++
.../tests/task_sdk/execution_time/test_comms.py | 155 ++++++++++++++++++++-
3 files changed, 206 insertions(+), 13 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py
b/task-sdk/src/airflow/sdk/execution_time/comms.py
index c70b3e9b518..a97df1df8b0 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -51,7 +51,9 @@ from __future__ import annotations
import asyncio
import itertools
import threading
+import traceback
from collections.abc import Iterator
+from contextlib import suppress
from datetime import datetime
from functools import cached_property
from pathlib import Path
@@ -118,6 +120,31 @@ SendMsgType = TypeVar("SendMsgType", bound=BaseModel)
ReceiveMsgType = TypeVar("ReceiveMsgType", bound=BaseModel)
+class DeadlockImminentError(BaseException):
+ """
+ Raised when ``send()`` is called from the event loop thread while
``asend()`` holds the lock.
+
+ Inherits from :class:`BaseException` so it escapes
``contextlib.suppress(Exception)``
+ and always surfaces to the caller.
+ """
+
+ def __init__(self, msg: object) -> None:
+ self.msg_type = type(msg).__name__
+ self.stack = "".join(traceback.format_stack())
+ super().__init__()
+
+ def __str__(self) -> str:
+ return (
+ f"comms.send() called from the event loop thread for message
'{self.msg_type}' "
+ "— deadlock is imminent (asend() is concurrently in-flight). "
+ "Likely cause: BaseHook.get_hook() or BaseHook.get_connection()
was called "
+ "from inside an async task. "
+ "Use the async equivalents instead: "
+ "await BaseHook.aget_hook() or await BaseHook.aget_connection()."
+ f"\nOffending call stack:\n{self.stack}"
+ )
+
+
def _msgpack_enc_hook(obj: Any) -> Any:
import pendulum
@@ -204,23 +231,32 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
err_decoder: TypeAdapter[ErrorResponse] = attrs.field(factory=lambda:
TypeAdapter(ToTask), repr=False)
- # Threading lock for sync operations
_thread_lock: threading.Lock = attrs.field(factory=threading.Lock,
repr=False)
- # Async lock for async operations
_async_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
+ _loop_thread_id: int | None = attrs.field(default=None, repr=False,
init=False)
def _make_frame(self, msg: SendMsgType) -> _RequestFrame:
carrier: dict[str, str] = {}
_trace_propagator.inject(carrier)
return _RequestFrame(id=next(self.id_counter), body=msg.model_dump(),
context_carrier=carrier or None)
+ @property
+ def _is_on_loop_thread(self) -> bool:
+ if threading.get_ident() == self._loop_thread_id:
+ with suppress(RuntimeError):
+ return bool(asyncio.get_running_loop())
+ return False
+
def send(self, msg: SendMsgType) -> ReceiveMsgType | None:
"""Send a request to the parent and block until the response is
received."""
frame_bytes = self._make_frame(msg).as_bytes()
- # We must make sure sockets aren't intermixed between sync and async
calls,
- # thus we need a dual locking mechanism to ensure that.
- with self._thread_lock:
+ # When called from the event loop thread, use non-blocking acquire to
detect
+ # an imminent deadlock: an asend() coroutine currently holds
_thread_lock and
+ # is waiting for the event loop to complete its I/O.
+ if not self._thread_lock.acquire(blocking=not self._is_on_loop_thread):
+ raise DeadlockImminentError(msg)
+ try:
self.socket.sendall(frame_bytes)
if isinstance(msg, ResendLoggingFD):
if recv_fds is None:
@@ -232,11 +268,13 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
if TYPE_CHECKING:
assert isinstance(resp, SentFDs)
resp.fds = fds
- # Since we know this is an expliclt SendFDs, and since this
class is generic SendFDs might not
+ # Since we know this is an explicit ResendLoggingFD, and since
this class is generic SentFDs might not
# always be in the return type union
return resp # type: ignore[return-value]
return self._get_response()
+ finally:
+ self._thread_lock.release()
async def asend(self, msg: SendMsgType) -> ReceiveMsgType | None:
"""
@@ -244,6 +282,8 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
Uses async lock for coroutine safety and thread lock for socket safety.
"""
+ self._loop_thread_id = threading.get_ident()
+
frame_bytes = self._make_frame(msg).as_bytes()
async with self._async_lock:
diff --git a/task-sdk/tests/task_sdk/execution_time/conftest.py
b/task-sdk/tests/task_sdk/execution_time/conftest.py
index 4a537373363..2d60b635f63 100644
--- a/task-sdk/tests/task_sdk/execution_time/conftest.py
+++ b/task-sdk/tests/task_sdk/execution_time/conftest.py
@@ -18,10 +18,22 @@
from __future__ import annotations
import sys
+from socket import socketpair
import pytest
[email protected]
+def socket_pair():
+ """Yield a connected ``(read_sock, write_sock)`` pair and guarantee both
are closed on teardown."""
+ r, w = socketpair()
+ try:
+ yield r, w
+ finally:
+ r.close()
+ w.close()
+
+
@pytest.fixture
def disable_capturing():
old_in, old_out, old_err = sys.stdin, sys.stdout, sys.stderr
diff --git a/task-sdk/tests/task_sdk/execution_time/test_comms.py
b/task-sdk/tests/task_sdk/execution_time/test_comms.py
index 51782796b82..97b9a79931f 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_comms.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_comms.py
@@ -19,7 +19,6 @@ from __future__ import annotations
import threading
import uuid
-from socket import socketpair
import msgspec
import pytest
@@ -29,6 +28,8 @@ from airflow.sdk import timezone
from airflow.sdk.execution_time.comms import (
BundleInfo,
CommsDecoder,
+ DeadlockImminentError,
+ GetVariable,
MaskSecret,
StartupDetails,
VariableResult,
@@ -75,8 +76,8 @@ class TestCommsDecoder:
"""Test the communication between the subprocess and the "supervisor"."""
@pytest.mark.usefixtures("disable_capturing")
- def test_recv_StartupDetails(self):
- r, w = socketpair()
+ def test_recv_StartupDetails(self, socket_pair):
+ r, w = socket_pair
msg = {
"type": "StartupDetails",
@@ -131,8 +132,8 @@ class TestCommsDecoder:
assert msg.bundle_info == BundleInfo(name="any-name",
version="any-version")
assert msg.start_date == timezone.datetime(2024, 12, 1, 1)
- def test_huge_payload(self):
- r, w = socketpair()
+ def test_huge_payload(self, socket_pair):
+ r, w = socket_pair
msg = {
"type": "XComResult",
@@ -160,8 +161,8 @@ class TestCommsDecoder:
assert len(msg.value) == 10 * 1024 * 1024 + 1
assert msg.value[-1] == "b"
- def test_send_thread_safety(self):
- r, w = socketpair()
+ def test_send_thread_safety(self, socket_pair):
+ r, w = socket_pair
decoder = CommsDecoder(socket=r, log=structlog.get_logger())
num_threads = 5
results = [None] * num_threads
@@ -208,3 +209,143 @@ class TestCommsDecoder:
assert errors[idx] is None, f"Thread {idx} error: {errors[idx]}"
assert results[idx].key == f"key{idx}", f"Out-of-order or missing
response for thread {idx}"
assert results[idx].value == f"value{idx}", f"Incorrect value for
thread {idx}"
+
+ @pytest.mark.asyncio
+ async def
test_send_from_event_loop_raises_deadlock_imminent_error_when_asend_in_flight(
+ self, socket_pair
+ ):
+ """
+ Regression: send() called from the event loop thread while asend()
holds
+ _thread_lock in a thread-pool worker must raise DeadlockImminentError
+ (deadlock_imminent=True) before attempting to acquire the lock.
+
+ DeadlockImminentError inherits from BaseException so it escapes
+ contextlib.suppress(Exception) in mask_secret() and similar helpers.
+ """
+ r, _ = socket_pair
+ decoder = CommsDecoder(socket=r, log=structlog.get_logger())
+
+ # Simulate asend() having recorded this thread as the event-loop
thread.
+ decoder._loop_thread_id = threading.get_ident()
+
+ # Hold _thread_lock from a background thread to simulate an in-flight
asend().
+ lock_held = threading.Event()
+ lock_release = threading.Event()
+
+ def _hold_lock():
+ decoder._thread_lock.acquire()
+ lock_held.set()
+ lock_release.wait()
+ decoder._thread_lock.release()
+
+ holder = threading.Thread(target=_hold_lock, daemon=True)
+ holder.start()
+ assert lock_held.wait(timeout=2), "Background thread never acquired
_thread_lock"
+
+ try:
+ # send() from the event loop thread while lock is held must raise
immediately.
+ with pytest.raises(DeadlockImminentError) as exc_info:
+ decoder.send(GetVariable(key="should_fail"))
+
+ msg = str(exc_info.value)
+ assert msg.startswith("comms.send() called from the event loop
thread")
+ assert "deadlock is imminent" in msg
+ finally:
+ lock_release.set()
+ holder.join(timeout=2)
+
+ @pytest.mark.asyncio
+ async def test_send_from_event_loop_succeeds_when_lock_free(self,
socket_pair):
+ """
+ send() called from the event loop thread when no asend() is currently
+ in-flight (lock not held) must proceed safely and return a result —
+ even after _loop_thread_id has been recorded by a prior asend() call.
+
+ The non-blocking acquire succeeds (lock is free), so the sync socket
+ I/O completes without needing the event loop, avoiding a deadlock.
+ DeadlockImminentError must NOT be raised in this scenario.
+ """
+ import asyncio
+
+ r, w = socket_pair
+ decoder = CommsDecoder(socket=r, log=structlog.get_logger())
+
+ def _serve_one():
+ length = int.from_bytes(w.recv(4), "big")
+ body = b""
+ while len(body) < length:
+ body += w.recv(length - len(body))
+ req = msgspec.msgpack.decode(body, type=_RequestFrame)
+ resp = {"type": "VariableResult", "key": req.body["key"], "value":
"v"}
+ encoded = msgspec.msgpack.encode(_ResponseFrame(req.id, resp,
None))
+ w.sendall(len(encoded).to_bytes(4, "big") + encoded)
+
+ # Phase 1: complete one asend() to record _loop_thread_id on the
decoder
+ server = threading.Thread(target=_serve_one, daemon=True)
+ server.start()
+ await asyncio.wait_for(decoder.asend(GetVariable(key="prime")),
timeout=5)
+ server.join(timeout=2)
+
+ # Phase 2: lock is free, _loop_thread_id is set.
+ # send() from the event loop thread must succeed (non-blocking acquire
+ # gets the free lock and the sync I/O completes without the event
loop).
+ server2 = threading.Thread(target=_serve_one, daemon=True)
+ server2.start()
+ result = decoder.send(GetVariable(key="should_succeed"))
+ server2.join(timeout=2)
+
+ assert result is not None
+ assert result.key == "should_succeed"
+
+ def test_send_after_event_loop_closes_does_not_raise(self, socket_pair):
+ """
+ Regression test: send() called from the same thread that previously ran
+ the event loop — but after the loop has finished — must NOT raise
+ DeadlockImminentError.
+
+ Scenario: AsyncOperator.execute() runs aexecute() via
+ loop.run_until_complete(), which calls await
SUPERVISOR_COMMS.asend(...)
+ (recording _loop_thread_id). After execute() returns the loop is
closed.
+ The task_runner then calls SUPERVISOR_COMMS.send() to push XCom
results.
+ Without this fix that call raised DeadlockImminentError spuriously.
+ """
+ import asyncio
+
+ r, w = socket_pair
+ decoder = CommsDecoder(socket=r, log=structlog.get_logger())
+
+ # Helper: serve exactly one request from the given socket pair.
+ def _serve_one(sock):
+ length = int.from_bytes(sock.recv(4), "big")
+ body = b""
+ while len(body) < length:
+ body += sock.recv(length - len(body))
+ req = msgspec.msgpack.decode(body, type=_RequestFrame)
+ resp = {"type": "VariableResult", "key": req.body["key"], "value":
"v"}
+ encoded = msgspec.msgpack.encode(_ResponseFrame(req.id, resp,
None))
+ sock.sendall(len(encoded).to_bytes(4, "big") + encoded)
+
+ # --- Phase 1: run the async part (simulates aexecute()) ---
+ async def _async_phase():
+ server = threading.Thread(target=_serve_one, args=(w,),
daemon=True)
+ server.start()
+ await
asyncio.wait_for(decoder.asend(GetVariable(key="async_key")), timeout=5)
+ server.join(timeout=2)
+
+ loop = asyncio.new_event_loop()
+ try:
+ loop.run_until_complete(_async_phase())
+ finally:
+ loop.close()
+ asyncio.set_event_loop(None)
+
+ # --- Phase 2: sync send() after the loop is closed (simulates
_push_xcom_if_needed) ---
+ # _loop_thread_id is still set from phase 1, but the loop is no longer
running.
+ # This must NOT raise DeadlockImminentError.
+ server2 = threading.Thread(target=_serve_one, args=(w,), daemon=True)
+ server2.start()
+
+ result = decoder.send(GetVariable(key="sync_key"))
+ server2.join(timeout=2)
+
+ assert result is not None
