Skip to content

Commit eb458af

Browse files
committed
gh-157058: Fix missing awaited-by edge in wait_for(fut, 0)
1 parent 878b5e2 commit eb458af

3 files changed

Lines changed: 36 additions & 0 deletions

File tree

Lib/asyncio/tasks.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -541,13 +541,20 @@ async def _cancel_and_wait(fut):
541541
cb = functools.partial(_release_waiter, waiter)
542542
fut.add_done_callback(cb)
543543

544+
# gh-157058: awaiting the waiter leaves no edge on fut, add it here
545+
cur_task = current_task()
546+
if cur_task is not None:
547+
futures.future_add_to_awaited_by(fut, cur_task)
548+
544549
try:
545550
fut.cancel()
546551
# We cannot wait on *fut* directly to make
547552
# sure _cancel_and_wait itself is reliably cancellable.
548553
await waiter
549554
finally:
550555
fut.remove_done_callback(cb)
556+
if cur_task is not None:
557+
futures.future_discard_from_awaited_by(fut, cur_task)
551558

552559

553560
class _AsCompletedIterator:

Lib/test/test_asyncio/test_graph.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,33 @@ class FakeCoro:
173173

174174
self.assertEqual(len(result.call_stack), 2)
175175

176+
async def test_stack_wait_for_non_positive_timeout(self):
177+
# gh-157058: wait_for(fut, 0) must still record the waiter
178+
cleanup = asyncio.Future()
179+
180+
async def worker():
181+
try:
182+
await asyncio.Future()
183+
finally:
184+
await cleanup
185+
186+
async def probe(t):
187+
await asyncio.wait_for(t, 0)
188+
189+
t = asyncio.ensure_future(worker())
190+
p = asyncio.create_task(probe(t), name='probe')
191+
for _ in range(5):
192+
await asyncio.sleep(0)
193+
194+
stack = capture_test_stack(fut=t)
195+
196+
cleanup.set_result(None)
197+
await asyncio.gather(p, t, return_exceptions=True)
198+
199+
self.assertEqual(stack[0][2], [
200+
['T<probe>', ['a _cancel_and_wait', 'a wait_for', 'a probe'], []],
201+
])
202+
176203
async def test_stack_gather(self):
177204

178205
stack_for_deep = None
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
:func:`asyncio.wait_for` with a non-positive timeout now records the waiter
2+
in the call graph.

0 commit comments

Comments
 (0)