Skip to content

Commit c627c2f

Browse files
miss-islingtondeadlovelllasvetlov
authored
[3.15] gh-159168: Fix RecursionError in asyncio pstree on a deep chain of tasks (GH-159170) (#161176)
gh-159168: Fix RecursionError in asyncio pstree on a deep chain of tasks (GH-159170) * gh-159168: Fix RecursionError in asyncio pstree on a deep chain of tasks * Apply suggestions from code review --------- (cherry picked from commit bb9d26d) Co-authored-by: Timofei Ivankov <128279579+deadlovelll@users.noreply.github.com> Co-authored-by: Andrew Svetlov <andrew.svetlov@gmail.com>
1 parent 429a9dc commit c627c2f

3 files changed

Lines changed: 100 additions & 16 deletions

File tree

‎Lib/asyncio/tools.py‎

Lines changed: 26 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -128,14 +128,22 @@ def _find_cycles(graph):
128128
def dfs(v):
129129
color[v] = GREY
130130
path.append(v)
131-
for w in graph.get(v, ()):
132-
if color[w] == WHITE:
133-
dfs(w)
134-
elif color[w] == GREY: # back-edge → cycle!
135-
i = path.index(w)
136-
cycles.append(path[i:] + [w]) # make a copy
137-
color[v] = BLACK
138-
path.pop()
131+
stack = [(v, iter(graph.get(v, ())))]
132+
while stack:
133+
node, children = stack[-1]
134+
for w in children:
135+
if color[w] == WHITE:
136+
color[w] = GREY
137+
path.append(w)
138+
stack.append((w, iter(graph.get(w, ()))))
139+
break
140+
elif color[w] == GREY: # back-edge → cycle!
141+
i = path.index(w)
142+
cycles.append(path[i:] + [w]) # make a copy
143+
else:
144+
color[node] = BLACK
145+
path.pop()
146+
stack.pop()
139147

140148
for v in list(graph):
141149
if color[v] == WHITE:
@@ -167,14 +175,16 @@ def pretty(node):
167175
flag = task_emoji if node[0] == NodeType.TASK else cor_emoji
168176
return f"{flag} {labels[node]}"
169177

170-
def render(node, prefix="", last=True, buf=None):
171-
if buf is None:
172-
buf = []
173-
buf.append(f"{prefix}{'└── ' if last else '├── '}{pretty(node)}")
174-
new_pref = prefix + (" " if last else "│ ")
175-
kids = children.get(node, [])
176-
for i, kid in enumerate(kids):
177-
render(kid, new_pref, i == len(kids) - 1, buf)
178+
def render(root):
179+
buf = []
180+
stack = [(root, "", True)]
181+
while stack:
182+
node, prefix, last = stack.pop()
183+
buf.append(f"{prefix}{'└── ' if last else '├── '}{pretty(node)}")
184+
new_pref = prefix + (" " if last else "│ ")
185+
kids = children.get(node, [])
186+
for i in range(len(kids) - 1, -1, -1):
187+
stack.append((kids[i], new_pref, i == len(kids) - 1))
178188
return buf
179189

180190
return [render(root) for root in _roots(labels, children)]

‎Lib/test/test_asyncio/test_tools.py‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import sys
12
import unittest
23

34
from asyncio import tools
@@ -1465,6 +1466,77 @@ def test_deep_coroutine_chain(self):
14651466
result = tools.build_async_tree(input_)
14661467
self.assertEqual(result, expected)
14671468

1469+
def test_deep_task_chain(self):
1470+
# gh-159168
1471+
n = sys.getrecursionlimit()
1472+
input_ = [
1473+
AwaitedInfo(
1474+
thread_id=1,
1475+
awaited_by=[
1476+
TaskInfo(
1477+
task_id=i,
1478+
task_name=f"Task-{i}",
1479+
coroutine_stack=[],
1480+
awaited_by=[
1481+
CoroInfo(
1482+
call_stack=[FrameInfo("f", "", LocationInfo(0))],
1483+
task_name=i - 1
1484+
)
1485+
] if i > 1 else []
1486+
)
1487+
for i in range(1, n + 1)
1488+
]
1489+
)
1490+
]
1491+
tree = tools.build_async_tree(input_)
1492+
self.assertEqual(len(tree[0]), 2 * n - 1)
1493+
self.assertEqual(tree[0][-1], " " * 8 * (n - 1) + f"└── (T) Task-{n}")
1494+
1495+
def test_cycle_below_task(self):
1496+
# gh-159168
1497+
input_ = [
1498+
AwaitedInfo(
1499+
thread_id=1,
1500+
awaited_by=[
1501+
TaskInfo(
1502+
task_id=1,
1503+
task_name="Task-1",
1504+
coroutine_stack=[],
1505+
awaited_by=[]
1506+
),
1507+
TaskInfo(
1508+
task_id=2,
1509+
task_name="Task-2",
1510+
coroutine_stack=[],
1511+
awaited_by=[
1512+
CoroInfo(
1513+
call_stack=[FrameInfo("f", "", LocationInfo(0))],
1514+
task_name=1
1515+
),
1516+
CoroInfo(
1517+
call_stack=[FrameInfo("f", "", LocationInfo(0))],
1518+
task_name=3
1519+
)
1520+
]
1521+
),
1522+
TaskInfo(
1523+
task_id=3,
1524+
task_name="Task-3",
1525+
coroutine_stack=[],
1526+
awaited_by=[
1527+
CoroInfo(
1528+
call_stack=[FrameInfo("f", "", LocationInfo(0))],
1529+
task_name=2
1530+
)
1531+
]
1532+
)
1533+
]
1534+
)
1535+
]
1536+
with self.assertRaises(tools.CycleFoundException) as ctx:
1537+
tools.build_async_tree(input_)
1538+
self.assertEqual(ctx.exception.cycles, [[2, 3, 2]])
1539+
14681540
def test_multiple_cycles_same_node(self):
14691541
input_ = [
14701542
AwaitedInfo(
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix :exc:`RecursionError` in ``python -m asyncio pstree`` on a deep chain of
2+
awaiting tasks. Patch by Timofei Ivankov.

0 commit comments

Comments
 (0)