dabla commented on code in PR #62922:
URL: https://github.com/apache/airflow/pull/62922#discussion_r4268274101
##########
task-sdk/src/airflow/sdk/execution_time/comms.py:
##########
@@ -239,55 +236,63 @@ class CommsDecoder(Generic[ReceiveMsgType, SendMsgType]):
err_decoder: TypeAdapter[ErrorResponse] = attrs.field(factory=lambda:
TypeAdapter(ToTask), repr=False)
- _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)
+ # Reentrant, so that a send() which interrupts its own thread's (a signal
handler on the main
+ # thread) is let in instead of waiting for itself, and can tell from
_exchanging whether that
+ # thread was in the middle of a request. The RLock records its owner when
it is taken, so no
+ # moment is left where the thread holds the lock but is not yet known as
its holder.
+ _thread_lock: threading.RLock = attrs.field(factory=threading.RLock,
repr=False)
+ # One per event loop: an asyncio.Lock binds to the loop that first waits
on it, and asend() can
+ # be called from several at once (a sync indexed task running
asyncio.run() while the iteration's loop
+ # has calls in flight). Each loop queues its calls on its own lock, so
only one thread of its
+ # default executor waits for the channel; _thread_lock orders the loops. A
loop that is gone
+ # drops its lock with it.
+ _async_locks: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop,
asyncio.Lock] = attrs.field(
Review Comment:
Fixed in e6e806c45e (with the docstring of `asend` updated in c4262724a8).
An `asyncio.Lock` only binds to its loop when a call has to wait for it, so a
lock nobody waited for is still reused across calls and drops with its loop
through the weak key; a lock a call waited for is dropped when the loop's last
call leaves. The bookkeeping moved out of `asend` into `_LoopQueues`, which
owns the weak-keyed map, and `_LoopQueue`, an async context manager per loop
that counts its calls, notes contention and leaves the map when it must. A new
test runs three `asyncio.run()` loops with overlapping calls against a real
`CommsDecoder` on a socket pair; before the fix all three closed loops stayed
alive, now none does.
Since `asend` is on the hot path, measured base vs head (CodSpeed simulation
of the comms benchmark from #74562, two rounds; wall clock for the contention
run):
| | before | after |
|---|---:|---:|
| `asend`, one call at a time | ~448k instructions | ~453k (+1.0%) |
| `asend`, 50 concurrent coroutines | ~18.87M | ~18.91M (+0.4%) |
| `send` (control) | ~114.0k | ~114.0k |
| `asyncio.to_thread` wait next to 50 calls in flight (median) | 0.24 ms |
0.26 ms |
A version that dropped the entry after every call cost +1.3% on the single
call; reusing an uncontended lock brought it back to about +1%.
Drafted-by: Claude Fable 5.1; 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]
