Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Lib/asyncio/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -541,13 +541,18 @@ async def _cancel_and_wait(fut):
cb = functools.partial(_release_waiter, waiter)
fut.add_done_callback(cb)

# gh-157058: awaiting the waiter leaves no edge on fut, add it here
cur_task = current_task()
futures.future_add_to_awaited_by(fut, cur_task)

try:
fut.cancel()
# We cannot wait on *fut* directly to make
# sure _cancel_and_wait itself is reliably cancellable.
await waiter
finally:
fut.remove_done_callback(cb)
futures.future_discard_from_awaited_by(fut, cur_task)


class _AsCompletedIterator:
Expand Down
27 changes: 27 additions & 0 deletions Lib/test/test_asyncio/test_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,6 +146,33 @@ async def main():
'async generator CallStackTestBase.test_stack_async_gen.<locals>.gen()',
stack_for_gen_nested_call[1])

async def test_stack_wait_for_non_positive_timeout(self):
# gh-157058: wait_for(fut, 0) must still record the waiter
cleanup = asyncio.Future()

async def worker():
try:
await asyncio.Future()
finally:
await cleanup

async def probe(t):
await asyncio.wait_for(t, 0)

t = asyncio.ensure_future(worker())
p = asyncio.create_task(probe(t), name='probe')
for _ in range(5):
await asyncio.sleep(0)

stack = capture_test_stack(fut=t)

cleanup.set_result(None)
await asyncio.gather(p, t, return_exceptions=True)

self.assertEqual(stack[0][2], [
['T<probe>', ['a _cancel_and_wait', 'a wait_for', 'a probe'], []],
])

async def test_stack_gather(self):

stack_for_deep = None
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
:func:`asyncio.wait_for` with a non-positive timeout now records the waiter
in the call graph.
Loading