Lib/asyncio/tools.py - external/github.com/python/cpython - Git at Google


            
            """Tools to analyze tasks running in asyncio programs."""


            from collections import defaultdict

            
            from itertools import count

            
            from enum import Enum

            
            import sys

            
            from _remote_debugging import RemoteUnwinder, FrameInfo


            class NodeType(Enum):

            
                COROUTINE = 1

            
                TASK = 2


            class CycleFoundException(Exception):

            
                """Raised when there is a cycle when drawing the call tree."""

            
                def __init__(

            
                        self,

            
                        cycles: list[list[int]],

            
                        id2name: dict[int, str],

            
                    ) -> None:

            
                    super().__init__(cycles, id2name)

            
                    self.cycles = cycles

            
                    self.id2name = id2name


            # ─── indexing helpers ───────────────────────────────────────────

            
            def _format_stack_entry(elem: str|FrameInfo) -> str:

            
                if not isinstance(elem, str):

            
                    if elem.location is None:

            
                        if elem.filename in ("", "~"):

            
                            return f"{elem.funcname}"

            
                        return f"{elem.funcname} {elem.filename}"

            
                    if elem.location.lineno == 0 and elem.filename == "":

            
                        return f"{elem.funcname}"

            
                    else:

            
                        return f"{elem.funcname} {elem.filename}:{elem.location.lineno}"

            
                return elem


            def _index(result):

            
                id2name, awaits, task_stacks = {}, [], {}

            
                for awaited_info in result:

            
                    for task_info in awaited_info.awaited_by:

            
                        task_id = task_info.task_id

            
                        task_name = task_info.task_name

            
                        id2name[task_id] = task_name


                        # Store the internal coroutine stack for this task

            
                        if task_info.coroutine_stack:

            
                            for coro_info in task_info.coroutine_stack:

            
                                call_stack = coro_info.call_stack

            
                                internal_stack = [_format_stack_entry(frame) for frame in call_stack]

            
                                task_stacks[task_id] = internal_stack


                        # Add the awaited_by relationships (external dependencies)

            
                        if task_info.awaited_by:

            
                            for coro_info in task_info.awaited_by:

            
                                call_stack = coro_info.call_stack

            
                                parent_task_id = coro_info.task_name

            
                                stack = [_format_stack_entry(frame) for frame in call_stack]

            
                                awaits.append((parent_task_id, stack, task_id))

            
                return id2name, awaits, task_stacks


            def _build_tree(id2name, awaits, task_stacks):

            
                id2label = {(NodeType.TASK, tid): name for tid, name in id2name.items()}

            
                children = defaultdict(list)

            
                cor_nodes = defaultdict(dict)  # Maps parent -> {frame_name: node_key}

            
                next_cor_id = count(1)


                def get_or_create_cor_node(parent, frame):

            
                    """Get existing coroutine node or create new one under parent"""

            
                    if frame in cor_nodes[parent]:

            
                        return cor_nodes[parent][frame]


                    node_key = (NodeType.COROUTINE, f"c{next(next_cor_id)}")

            
                    id2label[node_key] = frame

            
                    children[parent].append(node_key)

            
                    cor_nodes[parent][frame] = node_key

            
                    return node_key


                # Build task dependency tree with coroutine frames

            
                for parent_id, stack, child_id in awaits:

            
                    cur = (NodeType.TASK, parent_id)

            
                    for frame in reversed(stack):

            
                        cur = get_or_create_cor_node(cur, frame)


                    child_key = (NodeType.TASK, child_id)

            
                    if child_key not in children[cur]:

            
                        children[cur].append(child_key)


                # Add coroutine stacks for leaf tasks

            
                awaiting_tasks = {parent_id for parent_id, _, _ in awaits}

            
                for task_id in id2name:

            
                    if task_id not in awaiting_tasks and task_id in task_stacks:

            
                        cur = (NodeType.TASK, task_id)

            
                        for frame in reversed(task_stacks[task_id]):

            
                            cur = get_or_create_cor_node(cur, frame)


                return id2label, children


            def _roots(id2label, children):

            
                all_children = {c for kids in children.values() for c in kids}

            
                return [n for n in id2label if n not in all_children]


            # ─── detect cycles in the task-to-task graph ───────────────────────

            
            def _task_graph(awaits):

            
                """Return {parent_task_id: {child_task_id, …}, …}."""

            
                g = defaultdict(set)

            
                for parent_id, _stack, child_id in awaits:

            
                    g[parent_id].add(child_id)

            
                return g


            def _find_cycles(graph):

            
                """

            
                Depth-first search for back-edges.


                Returns a list of cycles (each cycle is a list of task-ids) or an

            
                empty list if the graph is acyclic.

            
                """

            
                WHITE, GREY, BLACK = 0, 1, 2

            
                color = defaultdict(lambda: WHITE)

            
                path, cycles = [], []


                def dfs(v):

            
                    color[v] = GREY

            
                    path.append(v)

            
                    for w in graph.get(v, ()):

            
                        if color[w] == WHITE:

            
                            dfs(w)

            
                        elif color[w] == GREY:            # back-edge → cycle!

            
                            i = path.index(w)

            
                            cycles.append(path[i:] + [w])  # make a copy

            
                    color[v] = BLACK

            
                    path.pop()


                for v in list(graph):

            
                    if color[v] == WHITE:

            
                        dfs(v)

            
                return cycles


            # ─── PRINT TREE FUNCTION ───────────────────────────────────────

            
            def get_all_awaited_by(pid):

            
                unwinder = RemoteUnwinder(pid)

            
                return unwinder.get_all_awaited_by()


            def build_async_tree(result, task_emoji="(T)", cor_emoji=""):

            
                """

            
                Build a list of strings for pretty-print an async call tree.


                The call tree is produced by `get_all_async_stacks()`, prefixing tasks

            
                with `task_emoji` and coroutine frames with `cor_emoji`.

            
                """

            
                id2name, awaits, task_stacks = _index(result)

            
                g = _task_graph(awaits)

            
                cycles = _find_cycles(g)

            
                if cycles:

            
                    raise CycleFoundException(cycles, id2name)

            
                labels, children = _build_tree(id2name, awaits, task_stacks)


                def pretty(node):

            
                    flag = task_emoji if node[0] == NodeType.TASK else cor_emoji

            
                    return f"{flag} {labels[node]}"


                def render(node, prefix="", last=True, buf=None):

            
                    if buf is None:

            
                        buf = []

            
                    buf.append(f"{prefix}{'└── ' if last else '├── '}{pretty(node)}")

            
                    new_pref = prefix + ("    " if last else "│   ")

            
                    kids = children.get(node, [])

            
                    for i, kid in enumerate(kids):

            
                        render(kid, new_pref, i == len(kids) - 1, buf)

            
                    return buf


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


            def build_task_table(result):

            
                id2name, _, _ = _index(result)

            
                table = []


                for awaited_info in result:

            
                    thread_id = awaited_info.thread_id

            
                    for task_info in awaited_info.awaited_by:

            
                        # Get task info

            
                        task_id = task_info.task_id

            
                        task_name = task_info.task_name


                        # Build coroutine stack string

            
                        frames = [frame for coro in task_info.coroutine_stack

            
                                 for frame in coro.call_stack]

            
                        coro_stack = " -> ".join(x.funcname for x in frames)


                        # Handle tasks with no awaiters

            
                        if not task_info.awaited_by:

            
                            table.append([thread_id, hex(task_id), task_name, coro_stack,

            
                                        "", "", "0x0"])

            
                            continue


                        # Handle tasks with awaiters

            
                        for coro_info in task_info.awaited_by:

            
                            parent_id = coro_info.task_name

            
                            awaiter_frames = [x.funcname for x in coro_info.call_stack]

            
                            awaiter_chain = " -> ".join(awaiter_frames)

            
                            awaiter_name = id2name.get(parent_id, "Unknown")

            
                            parent_id_str = (hex(parent_id) if isinstance(parent_id, int)

            
                                           else str(parent_id))


                            table.append([thread_id, hex(task_id), task_name, coro_stack,

            
                                        awaiter_chain, awaiter_name, parent_id_str])


                return table


            def _print_cycle_exception(exception: CycleFoundException):

            
                print("ERROR: await-graph contains cycles - cannot print a tree!", file=sys.stderr)

            
                print("", file=sys.stderr)

            
                for c in exception.cycles:

            
                    inames = " → ".join(exception.id2name.get(tid, hex(tid)) for tid in c)

            
                    print(f"cycle: {inames}", file=sys.stderr)


            def exit_with_permission_help_text():

            
                """

            
                Prints a message pointing to platform-specific permission help text and exits the program.

            
                This function is called when a PermissionError is encountered while trying

            
                to attach to a process.

            
                """

            
                print(

            
                    "Error: The specified process cannot be attached to due to insufficient permissions.\n"

            
                    "See the Python documentation for details on required privileges and troubleshooting:\n"

            
                    "https://docs.python.org/3/howto/remote_debugging.html#permission-requirements\n",

            
                    file=sys.stderr,

            
                )

            
                sys.exit(1)


            _TRANSIENT_ERRORS = (RuntimeError, OSError, UnicodeDecodeError, MemoryError)


            def _get_awaited_by_tasks(pid: int, retries: int = 3) -> list:

            
                for attempt in range(retries + 1):

            
                    try:

            
                        return get_all_awaited_by(pid)

            
                    except PermissionError:

            
                        exit_with_permission_help_text()

            
                    except ProcessLookupError:

            
                        print(f"Error: process {pid} not found.", file=sys.stderr)

            
                        sys.exit(1)

            
                    except _TRANSIENT_ERRORS as e:

            
                        if attempt < retries:

            
                            continue

            
                        if isinstance(e, RuntimeError):

            
                            while e.__context__ is not None:

            
                                e = e.__context__

            
                        print(f"Error retrieving tasks: {e}", file=sys.stderr)

            
                        sys.exit(1)


            def display_awaited_by_tasks_table(pid: int, retries: int = 3) -> None:

            
                """Build and print a table of all pending tasks under `pid`."""


                tasks = _get_awaited_by_tasks(pid, retries=retries)

            
                table = build_task_table(tasks)

            
                # Print the table in a simple tabular format

            
                print(

            
                    f"{'tid':<10} {'task id':<20} {'task name':<20} {'coroutine stack':<50} {'awaiter chain':<50} {'awaiter name':<15} {'awaiter id':<15}"

            
                )

            
                print("-" * 180)

            
                for row in table:

            
                    print(f"{row[0]:<10} {row[1]:<20} {row[2]:<20} {row[3]:<50} {row[4]:<50} {row[5]:<15} {row[6]:<15}")


            def display_awaited_by_tasks_tree(pid: int, retries: int = 3) -> None:

            
                """Build and print a tree of all pending tasks under `pid`."""


                tasks = _get_awaited_by_tasks(pid, retries=retries)

            
                try:

            
                    result = build_async_tree(tasks)

            
                except CycleFoundException as e:

            
                    _print_cycle_exception(e)

            
                    sys.exit(1)


                for tree in result:

            
                    print("\n".join(tree))
	"""Tools to analyze tasks running in asyncio programs."""

	from collections import defaultdict
	from itertools import count
	from enum import Enum
	import sys
	from _remote_debugging import RemoteUnwinder, FrameInfo

	class NodeType(Enum):
	COROUTINE = 1
	TASK = 2


	class CycleFoundException(Exception):
	"""Raised when there is a cycle when drawing the call tree."""
	def __init__(
	self,
	cycles: list[list[int]],
	id2name: dict[int, str],
	) -> None:
	super().__init__(cycles, id2name)
	self.cycles = cycles
	self.id2name = id2name



	# ─── indexing helpers ───────────────────────────────────────────
	def _format_stack_entry(elem: str\|FrameInfo) -> str:
	if not isinstance(elem, str):
	if elem.location is None:
	if elem.filename in ("", "~"):
	return f"{elem.funcname}"
	return f"{elem.funcname} {elem.filename}"
	if elem.location.lineno == 0 and elem.filename == "":
	return f"{elem.funcname}"
	else:
	return f"{elem.funcname} {elem.filename}:{elem.location.lineno}"
	return elem


	def _index(result):
	id2name, awaits, task_stacks = {}, [], {}
	for awaited_info in result:
	for task_info in awaited_info.awaited_by:
	task_id = task_info.task_id
	task_name = task_info.task_name
	id2name[task_id] = task_name

	# Store the internal coroutine stack for this task
	if task_info.coroutine_stack:
	for coro_info in task_info.coroutine_stack:
	call_stack = coro_info.call_stack
	internal_stack = [_format_stack_entry(frame) for frame in call_stack]
	task_stacks[task_id] = internal_stack

	# Add the awaited_by relationships (external dependencies)
	if task_info.awaited_by:
	for coro_info in task_info.awaited_by:
	call_stack = coro_info.call_stack
	parent_task_id = coro_info.task_name
	stack = [_format_stack_entry(frame) for frame in call_stack]
	awaits.append((parent_task_id, stack, task_id))
	return id2name, awaits, task_stacks


	def _build_tree(id2name, awaits, task_stacks):
	id2label = {(NodeType.TASK, tid): name for tid, name in id2name.items()}
	children = defaultdict(list)
	cor_nodes = defaultdict(dict) # Maps parent -> {frame_name: node_key}
	next_cor_id = count(1)

	def get_or_create_cor_node(parent, frame):
	"""Get existing coroutine node or create new one under parent"""
	if frame in cor_nodes[parent]:
	return cor_nodes[parent][frame]

	node_key = (NodeType.COROUTINE, f"c{next(next_cor_id)}")
	id2label[node_key] = frame
	children[parent].append(node_key)
	cor_nodes[parent][frame] = node_key
	return node_key

	# Build task dependency tree with coroutine frames
	for parent_id, stack, child_id in awaits:
	cur = (NodeType.TASK, parent_id)
	for frame in reversed(stack):
	cur = get_or_create_cor_node(cur, frame)

	child_key = (NodeType.TASK, child_id)
	if child_key not in children[cur]:
	children[cur].append(child_key)

	# Add coroutine stacks for leaf tasks
	awaiting_tasks = {parent_id for parent_id, _, _ in awaits}
	for task_id in id2name:
	if task_id not in awaiting_tasks and task_id in task_stacks:
	cur = (NodeType.TASK, task_id)
	for frame in reversed(task_stacks[task_id]):
	cur = get_or_create_cor_node(cur, frame)

	return id2label, children


	def _roots(id2label, children):
	all_children = {c for kids in children.values() for c in kids}
	return [n for n in id2label if n not in all_children]

	# ─── detect cycles in the task-to-task graph ───────────────────────
	def _task_graph(awaits):
	"""Return {parent_task_id: {child_task_id, …}, …}."""
	g = defaultdict(set)
	for parent_id, _stack, child_id in awaits:
	g[parent_id].add(child_id)
	return g


	def _find_cycles(graph):
	"""
	Depth-first search for back-edges.

	Returns a list of cycles (each cycle is a list of task-ids) or an
	empty list if the graph is acyclic.
	"""
	WHITE, GREY, BLACK = 0, 1, 2
	color = defaultdict(lambda: WHITE)
	path, cycles = [], []

	def dfs(v):
	color[v] = GREY
	path.append(v)
	for w in graph.get(v, ()):
	if color[w] == WHITE:
	dfs(w)
	elif color[w] == GREY: # back-edge → cycle!
	i = path.index(w)
	cycles.append(path[i:] + [w]) # make a copy
	color[v] = BLACK
	path.pop()

	for v in list(graph):
	if color[v] == WHITE:
	dfs(v)
	return cycles


	# ─── PRINT TREE FUNCTION ───────────────────────────────────────
	def get_all_awaited_by(pid):
	unwinder = RemoteUnwinder(pid)
	return unwinder.get_all_awaited_by()


	def build_async_tree(result, task_emoji="(T)", cor_emoji=""):
	"""
	Build a list of strings for pretty-print an async call tree.

	The call tree is produced by `get_all_async_stacks()`, prefixing tasks
	with `task_emoji` and coroutine frames with `cor_emoji`.
	"""
	id2name, awaits, task_stacks = _index(result)
	g = _task_graph(awaits)
	cycles = _find_cycles(g)
	if cycles:
	raise CycleFoundException(cycles, id2name)
	labels, children = _build_tree(id2name, awaits, task_stacks)

	def pretty(node):
	flag = task_emoji if node[0] == NodeType.TASK else cor_emoji
	return f"{flag} {labels[node]}"

	def render(node, prefix="", last=True, buf=None):
	if buf is None:
	buf = []
	buf.append(f"{prefix}{'└── ' if last else '├── '}{pretty(node)}")
	new_pref = prefix + (" " if last else "│ ")
	kids = children.get(node, [])
	for i, kid in enumerate(kids):
	render(kid, new_pref, i == len(kids) - 1, buf)
	return buf

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


	def build_task_table(result):
	id2name, _, _ = _index(result)
	table = []

	for awaited_info in result:
	thread_id = awaited_info.thread_id
	for task_info in awaited_info.awaited_by:
	# Get task info
	task_id = task_info.task_id
	task_name = task_info.task_name

	# Build coroutine stack string
	frames = [frame for coro in task_info.coroutine_stack
	for frame in coro.call_stack]
	coro_stack = " -> ".join(x.funcname for x in frames)

	# Handle tasks with no awaiters
	if not task_info.awaited_by:
	table.append([thread_id, hex(task_id), task_name, coro_stack,
	"", "", "0x0"])
	continue

	# Handle tasks with awaiters
	for coro_info in task_info.awaited_by:
	parent_id = coro_info.task_name
	awaiter_frames = [x.funcname for x in coro_info.call_stack]
	awaiter_chain = " -> ".join(awaiter_frames)
	awaiter_name = id2name.get(parent_id, "Unknown")
	parent_id_str = (hex(parent_id) if isinstance(parent_id, int)
	else str(parent_id))

	table.append([thread_id, hex(task_id), task_name, coro_stack,
	awaiter_chain, awaiter_name, parent_id_str])

	return table

	def _print_cycle_exception(exception: CycleFoundException):
	print("ERROR: await-graph contains cycles - cannot print a tree!", file=sys.stderr)
	print("", file=sys.stderr)
	for c in exception.cycles:
	inames = " → ".join(exception.id2name.get(tid, hex(tid)) for tid in c)
	print(f"cycle: {inames}", file=sys.stderr)


	def exit_with_permission_help_text():
	"""
	Prints a message pointing to platform-specific permission help text and exits the program.
	This function is called when a PermissionError is encountered while trying
	to attach to a process.
	"""
	print(
	"Error: The specified process cannot be attached to due to insufficient permissions.\n"
	"See the Python documentation for details on required privileges and troubleshooting:\n"
	"https://docs.python.org/3/howto/remote_debugging.html#permission-requirements\n",
	file=sys.stderr,
	)
	sys.exit(1)


	_TRANSIENT_ERRORS = (RuntimeError, OSError, UnicodeDecodeError, MemoryError)


	def _get_awaited_by_tasks(pid: int, retries: int = 3) -> list:
	for attempt in range(retries + 1):
	try:
	return get_all_awaited_by(pid)
	except PermissionError:
	exit_with_permission_help_text()
	except ProcessLookupError:
	print(f"Error: process {pid} not found.", file=sys.stderr)
	sys.exit(1)
	except _TRANSIENT_ERRORS as e:
	if attempt < retries:
	continue
	if isinstance(e, RuntimeError):
	while e.__context__ is not None:
	e = e.__context__
	print(f"Error retrieving tasks: {e}", file=sys.stderr)
	sys.exit(1)


	def display_awaited_by_tasks_table(pid: int, retries: int = 3) -> None:
	"""Build and print a table of all pending tasks under `pid`."""

	tasks = _get_awaited_by_tasks(pid, retries=retries)
	table = build_task_table(tasks)
	# Print the table in a simple tabular format
	print(
	f"{'tid':<10} {'task id':<20} {'task name':<20} {'coroutine stack':<50} {'awaiter chain':<50} {'awaiter name':<15} {'awaiter id':<15}"
	)
	print("-" * 180)
	for row in table:
	print(f"{row[0]:<10} {row[1]:<20} {row[2]:<20} {row[3]:<50} {row[4]:<50} {row[5]:<15} {row[6]:<15}")


	def display_awaited_by_tasks_tree(pid: int, retries: int = 3) -> None:
	"""Build and print a tree of all pending tasks under `pid`."""

	tasks = _get_awaited_by_tasks(pid, retries=retries)
	try:
	result = build_async_tree(tasks)
	except CycleFoundException as e:
	_print_cycle_exception(e)
	sys.exit(1)

	for tree in result:
	print("\n".join(tree))