gh-144881: Add retry logic to asyncio debugging tools

Transient errors can occur when attaching to a process that is actively
using thread delegation (e.g. asyncio.to_thread). Add a retry loop to
_get_awaited_by_tasks for RuntimeError, OSError, UnicodeDecodeError, and
MemoryError, and expose --retries CLI flag on both `ps` and `pstree`
subcommands (default: 3).

Author: johnslavik

Lib/asyncio/__main__.py

@@ -162,17 +162,29 @@ def interrupt(self) -> None:
         "ps", help="Display a table of all pending tasks in a process"
     )
     ps.add_argument("pid", type=int, help="Process ID to inspect")
+    ps.add_argument(
+        "--retries",
+        type=int,
+        default=3,
+        help="Number of retries on transient attach errors (default: 3)",
+    )
     pstree = subparsers.add_parser(
         "pstree", help="Display a tree of all pending tasks in a process"
     )
     pstree.add_argument("pid", type=int, help="Process ID to inspect")
+    pstree.add_argument(
+        "--retries",
+        type=int,
+        default=3,
+        help="Number of retries on transient attach errors (default: 3)",
+    )
     args = parser.parse_args()
     match args.command:
         case "ps":
-            asyncio.tools.display_awaited_by_tasks_table(args.pid)
+            asyncio.tools.display_awaited_by_tasks_table(args.pid, retries=args.retries)
             sys.exit(0)
         case "pstree":
-            asyncio.tools.display_awaited_by_tasks_tree(args.pid)
+            asyncio.tools.display_awaited_by_tasks_tree(args.pid, retries=args.retries)
             sys.exit(0)
         case None:
             pass  # continue to the interactive shell

Lib/asyncio/tools.py

@@ -236,22 +236,37 @@ def exit_with_permission_help_text():
     sys.exit(1)
 
 
-def _get_awaited_by_tasks(pid: int) -> list:
-    try:
-        return get_all_awaited_by(pid)
-    except RuntimeError as e:
-        while e.__context__ is not None:
-            e = e.__context__
-        print(f"Error retrieving tasks: {e}")
-        sys.exit(1)
-    except PermissionError:
-        exit_with_permission_help_text()
+_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:
+                print(
+                    f"Transient error while reading process state "
+                    f"(attempt {attempt + 1}/{retries + 1}), retrying...",
+                    file=sys.stderr,
+                )
+                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) -> None:
+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)
+    tasks = _get_awaited_by_tasks(pid, retries=retries)
     table = build_task_table(tasks)
     # Print the table in a simple tabular format
     print(
@@ -262,10 +277,10 @@ def display_awaited_by_tasks_table(pid: int) -> None:
         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) -> None:
+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)
+    tasks = _get_awaited_by_tasks(pid, retries=retries)
     try:
         result = build_async_tree(tasks)
     except CycleFoundException as e: