dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4252453717
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -306,42 +284,45 @@ def send(self, msg: SendMsgType) -> ReceiveMsgType | None:
return self._get_response()
finally:
- self._thread_lock.release()
+ self._channel_owner = None
+
+ def _locked_exchange(
+ self, msg: SendMsgType, frame_bytes: bytes | bytearray, abandoned:
threading.Event
+ ) -> ReceiveMsgType | None:
+ """
+ Take ``_thread_lock``, exchange one request and its response, and
release it.
+
+ A request whose caller gave up while this thread waited for the lock
is not sent: nothing
+ reached the parent yet, and its side effects (an XCom push, a
checkpoint) belong to a
+ coroutine that was cancelled. Once sent, its response is always read,
so the next request
+ never reads it as its own.
+ """
+ with self._thread_lock:
+ if abandoned.is_set():
+ return None
+ return self._exchange(msg, frame_bytes)
async def asend(self, msg: SendMsgType) -> ReceiveMsgType | None:
"""
Send a request to the parent without blocking.
- Uses async lock for coroutine safety and thread lock for socket safety.
+ The async lock orders the coroutines; the request itself runs in a
worker thread, which
+ takes ``_thread_lock``, writes the request, reads its response and
releases the lock.
+ Cancelling the await cancels the future, not that thread: a request
not sent yet is
+ dropped, one already sent has its response drained before the thread
lets go of the
+ channel, so the next request never reads a response that was not its
own, nor finds the
+ lock taken with nobody left to release it.
"""
- self._loop_thread_id = threading.get_ident()
-
frame_bytes = self._make_frame(msg).as_bytes()
async with self._async_lock:
Review Comment:
Thanks, reproduced, and it is wider than two loops at the same time: the
lock binds to the first loop that has to wait on it, so any later loop in the
process (another `asyncio.run()`) fails the same way once two of its calls
overlap. Both cases are tests now.
I went with one `asyncio.Lock` per event loop instead of dropping it
(364829bfbf), and measured both before choosing. Without the lock every
concurrent `asend()` takes a thread of the loop's default executor and blocks
it on `_thread_lock` until the channel is free. That executor is the one the
iteration's other `asyncio.to_thread` work uses (rendering each item's
templates, reading lazy sources), so that work queues behind threads that only
wait for the channel. With a supervisor answering in 2 ms and 50 calls in
flight, a `to_thread` job waited (16 CPUs here, so 20 default-executor threads;
with fewer CPUs there are fewer threads and the wait grows):
| | median wait |
|---|---|
| lock dropped | 76.7 ms |
| one lock per loop | 0.2 ms |
Dropping it is cheaper on the comms path itself (CodSpeed simulation,
instructions: 50 concurrent `asend` -8.4%, one `asend` -2.9%), the per-loop
lock keeps today's numbers (-0.0% and +0.8%), but the threads it saves on the
loop are then spent waiting for the channel, where the iteration's other work
needs them. Each loop now queues its own calls on its own lock in a
`WeakKeyDictionary`, which drops the lock with its loop, and `_thread_lock`
orders the loops, as you point out.
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]
