task-waiter iterative bfs walk

Author: maurycy

Lib/test/test_external_inspection.py

@@ -3809,9 +3809,9 @@ class TestFrameChainLimits(RemoteInspectionTestBase):
     """Frame chain walks abort instead of looping/overflowing on deep chains."""
 
     # Limits plus one, to exceed them (must match MAX_FRAME_CHAIN_DEPTH /
-    # MAX_TASK_WAITER_CHAIN_DEPTH from _remote_debugging.h)
+    # MAX_TASK_WAITER_WALK_TASKS from _remote_debugging.h)
     FRAME_CHAIN_DEPTH = 1024 + 512 + 1
-    TASK_WAITER_CHAIN_DEPTH = 256 + 1
+    TASK_WAITER_WALK_TASKS = 2**16 + 1
 
     def _assert_unwinder_limit_error(self, unwind, expected_substring):
         """Call unwind() until it raises the frame chain limit error.
@@ -3886,7 +3886,7 @@ async def chain(n):
                 task = asyncio.create_task(chain(n - 1))
                 await task
 
-            asyncio.run(chain({self.TASK_WAITER_CHAIN_DEPTH}))
+            asyncio.run(chain({self.TASK_WAITER_WALK_TASKS}))
             """
         with self._target_process(script_body) as (p, client_socket, _):
             _wait_for_signal(client_socket, b"ready")

Modules/_remote_debugging/_remote_debugging.h

@@ -148,7 +148,7 @@ typedef enum _WIN32_THREADSTATE {
 #define MAX_LONG_DIGITS 64  /* Allows values up to ~2^1920 */
 #define MAX_SET_TABLE_SIZE (1 << 20)  /* 1 million entries max for set iteration */
 #define MAX_FRAME_CHAIN_DEPTH (1024 + 512)  /* Iteration bound for frame chain walks */
-#define MAX_TASK_WAITER_CHAIN_DEPTH 256  /* Task waiter walk recursion cap: 256 * SIZEOF_TASK_OBJ ~= 1 MiB */
+#define MAX_TASK_WAITER_WALK_TASKS (1 << 16)  /* Total-task bound for waiter walks */
 
 #ifndef MAX
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
@@ -480,13 +480,6 @@ typedef int (*thread_processor_func)(
     void *context
 );
 
-typedef int (*set_entry_processor_func)(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t key_addr,
-    void *context,
-    size_t depth
-);
-
 typedef int (*interpreter_processor_func)(
     RuntimeOffsets *offsets,
     uintptr_t interpreter_state_addr,
@@ -746,38 +739,13 @@ extern int parse_async_frame_chain(
     uintptr_t running_task_code_obj
 );
 
-/* Set iteration */
-extern int iterate_set_entries(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t set_addr,
-    set_entry_processor_func processor,
-    void *context,
-    size_t depth
-);
-
-/* Task awaited_by processing */
-extern int process_task_awaited_by(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t task_address,
-    set_entry_processor_func processor,
-    void *context,
-    size_t depth
-);
-
 extern int process_single_task_node(
     RemoteUnwinderObject *unwinder,
     uintptr_t task_addr,
     PyObject **task_info,
     PyObject *result
 );
 
-extern int process_task_and_waiters(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t task_addr,
-    PyObject *result,
-    size_t depth
-);
-
 extern int find_running_task_in_thread(
     RemoteUnwinderObject *unwinder,
     uintptr_t thread_state_addr,

Modules/_remote_debugging/asyncio.c

@@ -116,13 +116,11 @@ ensure_async_debug_offsets(RemoteUnwinderObject *unwinder)
  * SET ITERATION FUNCTIONS
  * ============================================================================ */
 
-int
+static int
 iterate_set_entries(
     RemoteUnwinderObject *unwinder,
     uintptr_t set_addr,
-    set_entry_processor_func processor,
-    void *context,
-    size_t depth
+    PyObject *awaited_by
 ) {
     char set_object[SIZEOF_SET_OBJ];
     if (_Py_RemoteDebug_PagedReadRemoteMemory(&unwinder->handle, set_addr,
@@ -154,19 +152,10 @@ iterate_set_entries(
         }
 
         if ((void*)key_addr != NULL) {
-            Py_ssize_t ref_cnt;
-            if (read_Py_ssize_t(unwinder, table_ptr, &ref_cnt) < 0) {
-                set_exception_cause(unwinder, PyExc_RuntimeError, "Failed to read set entry ref count");
+            if (parse_task(unwinder, key_addr, awaited_by) < 0) {
                 return -1;
             }
-
-            if (ref_cnt) {
-                // Process this valid set entry
-                if (processor(unwinder, key_addr, context, depth) < 0) {
-                    return -1;
-                }
-                els++;
-            }
+            els++;
         }
         table_ptr += sizeof(void*) * 2;
         i++;
@@ -526,44 +515,11 @@ parse_task(
  * TASK AWAITED_BY PROCESSING
  * ============================================================================ */
 
-// Forward declaration for mutual recursion
-static int process_waiter_task(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t key_addr,
-    void *context,
-    size_t depth
-);
-
-// Processor function for parsing tasks in sets
-static int
-process_task_parser(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t key_addr,
-    void *context,
-    size_t depth
-) {
-    (void)depth;
-    PyObject *awaited_by = (PyObject *)context;
-    return parse_task(unwinder, key_addr, awaited_by);
-}
-
 static int
 parse_task_awaited_by(
     RemoteUnwinderObject *unwinder,
     uintptr_t task_address,
     PyObject *awaited_by
-) {
-    return process_task_awaited_by(
-        unwinder, task_address, process_task_parser, awaited_by, 0);
-}
-
-int
-process_task_awaited_by(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t task_address,
-    set_entry_processor_func processor,
-    void *context,
-    size_t depth
 ) {
     // Read the entire TaskObj at once
     char task_obj[SIZEOF_TASK_OBJ];
@@ -582,11 +538,10 @@ process_task_awaited_by(
     char awaited_by_is_a_set = GET_MEMBER(char, task_obj, unwinder->async_debug_offsets.asyncio_task_object.task_awaited_by_is_set);
 
     if (awaited_by_is_a_set) {
-        return iterate_set_entries(
-            unwinder, task_ab_addr, processor, context, depth);
+        return iterate_set_entries(unwinder, task_ab_addr, awaited_by);
     } else {
         // Single task waiting
-        return processor(unwinder, task_ab_addr, context, depth);
+        return parse_task(unwinder, task_ab_addr, awaited_by);
     }
 }
 
@@ -681,40 +636,39 @@ process_single_task_node(
     return -1;
 }
 
-int
-process_task_and_waiters(
+static int
+process_task_waiters(
     RemoteUnwinderObject *unwinder,
-    uintptr_t task_addr,
-    PyObject *result,
-    size_t depth
+    PyObject *result
 ) {
-    if (depth >= MAX_TASK_WAITER_CHAIN_DEPTH) {
-        PyErr_SetString(PyExc_RuntimeError,
-            "Too many task waiters (possible infinite loop)");
-        set_exception_cause(unwinder, PyExc_RuntimeError,
-            "Task waiter chain depth limit exceeded");
-        return -1;
-    }
-
-    // First, add this task to the result
-    if (process_single_task_node(unwinder, task_addr, NULL, result) < 0) {
-        return -1;
+    for (Py_ssize_t i = 0; i < PyList_GET_SIZE(result); i++) {
+        PyObject *task_info = PyList_GET_ITEM(result, i);
+        PyObject *waiters = PyStructSequence_GET_ITEM(task_info, 3);
+        for (Py_ssize_t j = 0; j < PyList_GET_SIZE(waiters); j++) {
+            if (PyList_GET_SIZE(result) >= MAX_TASK_WAITER_WALK_TASKS) {
+                PyErr_SetString(PyExc_RuntimeError,
+                    "Too many task waiters (possible infinite loop)");
+                set_exception_cause(unwinder, PyExc_RuntimeError,
+                    "Task waiter walk size limit exceeded");
+                return -1;
+            }
+            PyObject *waiter = PyList_GET_ITEM(waiters, j);
+            PyObject *task_id = PyStructSequence_GET_ITEM(waiter, 1);
+            void *task_ptr = PyLong_AsVoidPtr(task_id);
+            if (task_ptr == NULL && PyErr_Occurred()) {
+                set_exception_cause(unwinder, PyExc_RuntimeError,
+                                    "Failed to parse waiter task ID");
+                return -1;
+            }
+            if (process_single_task_node(
+                    unwinder, (uintptr_t)task_ptr, NULL, result) < 0)
+            {
+                return -1;
+            }
+        }
     }
 
-    // Now find all tasks that are waiting for this task and process them
-    return process_task_awaited_by(
-        unwinder, task_addr, process_waiter_task, result, depth + 1);
-}
-
-// Processor function for task waiters
-static int
-process_waiter_task(
-    RemoteUnwinderObject *unwinder,
-    uintptr_t key_addr,
-    void *context,
-    size_t depth
-) {
-    return process_task_and_waiters(unwinder, key_addr, (PyObject *)context, depth);
+    return 0;
 }
 
 /* ============================================================================
@@ -1017,9 +971,7 @@ process_running_task_chain(
     }
 
     // Now find all tasks that are waiting for this task and process them
-    if (process_task_awaited_by(
-            unwinder, running_task_addr, process_waiter_task, result, 1) < 0)
-    {
+    if (process_task_waiters(unwinder, result) < 0) {
         return -1;
     }