Skip to content

Commit 1dbd492

Browse files
committed
Fix asyncio.shield leaking tasks via await-graph.
There were two places where `cur_task` could live until `inner` ended: 1. `inner._asyncio_awaited_by`, and 2. as a capture in `_clear_awaited_by_callback` in `inner._callbacks`.
1 parent fd9feab commit 1dbd492

3 files changed

Lines changed: 22 additions & 0 deletions

File tree

Lib/asyncio/tasks.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -997,6 +997,9 @@ def _outer_done_callback(outer):
997997
# Keep only one callback to log on cancel
998998
inner.remove_done_callback(_log_on_exception)
999999
inner.add_done_callback(_log_on_exception)
1000+
if cur_task is not None:
1001+
inner.remove_done_callback(_clear_awaited_by_callback)
1002+
futures.future_discard_from_awaited_by(inner, cur_task)
10001003

10011004
if cur_task is not None:
10021005
inner.add_done_callback(_clear_awaited_by_callback)

Lib/test/test_asyncio/test_tasks.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2109,6 +2109,8 @@ def test_shield_cancel_outer(self):
21092109
test_utils.run_briefly(self.loop)
21102110
self.assertTrue(outer.cancelled())
21112111
self.assertEqual(0, 0 if outer._callbacks is None else len(outer._callbacks))
2112+
self.assertEqual(0, 0 if inner._asyncio_awaited_by is None else len(inner._asyncio_awaited_by))
2113+
self.assertTrue({f for f, _ctx in inner._callbacks or []} <= {asyncio.tasks._log_on_exception})
21122114

21132115
def test_shield_cancel_outer_result(self):
21142116
mock_handler = mock.Mock()
@@ -2134,6 +2136,21 @@ def test_shield_cancel_outer_exception(self):
21342136
test_utils.run_briefly(self.loop)
21352137
mock_handler.assert_called_once()
21362138

2139+
def test_shield_cancel_outer_in_task(self):
2140+
inner = self.new_future(self.loop)
2141+
2142+
async def coro():
2143+
outer = asyncio.shield(inner)
2144+
self.assertNotEqual(0, len(inner._callbacks))
2145+
outer.cancel()
2146+
await asyncio.sleep(0)
2147+
self.assertTrue(outer.cancelled())
2148+
2149+
task = self.new_task(self.loop, coro())
2150+
self.loop.run_until_complete(task)
2151+
self.assertEqual(0, 0 if inner._asyncio_awaited_by is None else len(inner._asyncio_awaited_by))
2152+
self.assertTrue({f for f, _ctx in inner._callbacks or []} <= {asyncio.tasks._log_on_exception})
2153+
21372154
def test_shield_duplicate_log_once(self):
21382155
mock_handler = mock.Mock()
21392156
self.loop.set_exception_handler(mock_handler)
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix :func:`asyncio.shield` leaking the calling task via the await-graph and
2+
callbacks when called on a future that never resolves.

0 commit comments

Comments
 (0)