Skip to content

Commit 6a93c91

Browse files
[3.15] gh-155143: Fix asyncio.shield() leaking tasks via await-graph (GH-155144) (#157320)
Co-authored-by: Andrew Geng <pteromys@gmail.com>
1 parent 0decd77 commit 6a93c91

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
@@ -1013,6 +1013,9 @@ def _outer_done_callback(outer):
10131013
# Keep only one callback to log on cancel
10141014
inner.remove_done_callback(_log_on_exception)
10151015
inner.add_done_callback(_log_on_exception)
1016+
if cur_task is not None:
1017+
inner.remove_done_callback(_clear_awaited_by_callback)
1018+
futures.future_discard_from_awaited_by(inner, cur_task)
10161019

10171020
if cur_task is not None:
10181021
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
@@ -2147,6 +2147,8 @@ def test_shield_cancel_outer(self):
21472147
test_utils.run_briefly(self.loop)
21482148
self.assertTrue(outer.cancelled())
21492149
self.assertEqual(0, 0 if outer._callbacks is None else len(outer._callbacks))
2150+
self.assertFalse(inner._asyncio_awaited_by)
2151+
self.assertTrue({f for f, _ctx in inner._callbacks or []} <= {asyncio.tasks._log_on_exception})
21502152

21512153
def test_shield_cancel_outer_result(self):
21522154
mock_handler = mock.Mock()
@@ -2172,6 +2174,21 @@ def test_shield_cancel_outer_exception(self):
21722174
test_utils.run_briefly(self.loop)
21732175
mock_handler.assert_called_once()
21742176

2177+
def test_shield_cancel_outer_in_task(self):
2178+
inner = self.new_future(self.loop)
2179+
2180+
async def coro():
2181+
outer = asyncio.shield(inner)
2182+
self.assertNotEqual(0, len(inner._callbacks))
2183+
outer.cancel()
2184+
await asyncio.sleep(0)
2185+
self.assertTrue(outer.cancelled())
2186+
2187+
task = self.new_task(self.loop, coro())
2188+
self.loop.run_until_complete(task)
2189+
self.assertFalse(inner._asyncio_awaited_by)
2190+
self.assertTrue({f for f, _ctx in inner._callbacks or []} <= {asyncio.tasks._log_on_exception})
2191+
21752192
def test_shield_duplicate_log_once(self):
21762193
mock_handler = mock.Mock()
21772194
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)