dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4252087367
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -241,53 +237,35 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
_thread_lock: threading.Lock = attrs.field(factory=threading.Lock,
repr=False)
_async_lock: asyncio.Lock = attrs.field(factory=asyncio.Lock, repr=False)
- _loop_thread_id: int | None = attrs.field(default=None, repr=False,
init=False)
+ # The thread exchanging a request while it holds _thread_lock, so send()
can tell a call that
+ # would wait for its own thread from one that waits for another.
+ _channel_owner: 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
-
- @property
- def _asend_in_flight_on_this_thread(self) -> bool:
+ def send(self, msg: SendMsgType) -> ReceiveMsgType | None:
"""
- Whether an ``asend()`` coroutine of a loop that runs on the current
thread is in flight.
-
- ``asend()`` holds ``_async_lock`` for its whole duration, from before
it takes
- ``_thread_lock`` (in an executor thread) until after it releases it,
so the async lock
- being held is the reliable signal that the thread lock is, or is about
to be, held by a
- coroutine. Such a coroutine only makes progress when its loop runs,
and that loop runs on
- this thread. This is thread-scoped on purpose: a loop running on
*another* thread keeps
- going while this one blocks, so ``send()`` from there can safely wait
for the lock.
+ Send a request to the parent and block until the response is received.
+
+ Called on an event loop's thread while an ``asend()`` is in flight,
this waits for that
+ request's round trip: ``asend()`` runs on a worker thread that needs
no loop to finish, so
+ the wait ends, but it stalls every coroutine on the loop meanwhile.
Async code should
+ prefer ``asend()`` and the async SDK calls built on it. Only a call on
a thread that is
+ itself in the middle of a request raises
:class:`DeadlockImminentError`, as waiting there
+ would never end.
"""
- return threading.get_ident() == self._loop_thread_id and
self._async_lock.locked()
-
- 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()
-
- # An asend() in flight on this thread's loop can only release
_thread_lock once that loop
- # runs again, and the loop cannot run while this thread blocks in
send(). Raise instead of
- # blocking. This covers the loop being paused between two
run_until_complete() calls, where
- # the running-loop check below does not apply:
asyncio.get_running_loop() raises, so send()
- # would otherwise fall through to a blocking acquire and freeze the
process with every
- # thread idle (seen with a sync XCom read pulling the next iterated
sub-task's input while a
- # sub-task's asend() was parked mid-I/O).
- if self._asend_in_flight_on_this_thread:
+ if self._channel_owner == threading.get_ident():
raise DeadlockImminentError(msg)
+ frame_bytes = self._make_frame(msg).as_bytes()
+ with self._thread_lock:
+ return self._exchange(msg, frame_bytes)
- # When called from the event loop thread while the loop is running,
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)
+ def _exchange(self, msg: SendMsgType, frame_bytes: bytes | bytearray) ->
ReceiveMsgType | None:
+ """Write one request and read its response; the caller holds
``_thread_lock``."""
+ self._channel_owner = threading.get_ident()
Review Comment:
Right, closed in 345e3cb74c as you suggested: `_thread_lock` is now an
`RLock`, which records its owner as it is taken, and `send()` reads an
`_exchanging` flag once it holds the lock instead of the owner set inside
`_exchange`. Another thread still waits for the lock. The owner thread is let
in again: in the middle of a request it raises `DeadlockImminentError` as
before, and before the request is written or after its response is read,
nothing of it is on the wire, so the handler's call is exchanged whole in
between and both callers get their own response. Two new tests land a nested
`send()` in each of those windows; both hang on the previous code, so they
carry `execution_timeout(10)` to fail instead. CodSpeed simulation on the comms
benchmarks of #74562, base vs this change:
| Benchmark | Instructions before | After | Change |
|---|---|---|---|
| `send` round trip | 115,881 | 114,059 | -1.6% |
| `asend` round trip | 445,395 | 443,415 | -0.4% |
| 50 concurrent `asend` | 18,732,959 | 18,769,872 | +0.2% |
Drafted-by: Claude Opus 5.5; reviewed by @dabla before posting
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
