Combine python and native unwinder into single loop

Python, especially pytorch programs can exhaust the tail call limit
by switching from python to native unwinders more than 29 times.
This happens because of eval/delegation patterns where one python
frame will be decorated with a couple native frames.

In order to unwind these stack successfully fold the native unwinder
into the python unwinder so at each frame a python or native frame
can be unwound.

Replace the separate walk_python_stack inner loop and outer
transition loop with a single switch-in-loop structure using
step_python and step_native helper functions. This reduces
tail call usage from one per batch to one per loop budget
exhaustion (PYTHON_NATIVE_LOOP_ITERS=9 iterations).

Move native unwinder map externs (exe_id_to_*_stack_deltas,
stack_delta_page_to_info, unwind_info_array) out of the
TESTING_COREDUMP guard in extmaps.h so python_tracer.ebpf.c
can include native_stack_trace.h.

Python loop iters is now a ro_vars entry so it can be set low by
default and jacked up with debug_prints are disabled which allows for
much bigger stacks.

Author: gnurizen

support/ebpf/extmaps.h

@@ -16,13 +16,7 @@ extern struct interpreter_offsets_t interpreter_offsets;
 extern struct trace_events_t trace_events;
 extern struct go_labels_procs_t go_labels_procs;
 
-#if defined(TESTING_COREDUMP)
-
-// References to maps in alphabetical order that
-// are needed only for testing.
-
-extern struct apm_int_procs_t apm_int_procs;
-extern struct beam_procs_t beam_procs;
+// These are needed by both native and hybrid python unwinder.
 extern struct exe_id_to_8_stack_deltas_t exe_id_to_8_stack_deltas;
 extern struct exe_id_to_9_stack_deltas_t exe_id_to_9_stack_deltas;
 extern struct exe_id_to_10_stack_deltas_t exe_id_to_10_stack_deltas;
@@ -39,14 +33,21 @@ extern struct exe_id_to_20_stack_deltas_t exe_id_to_20_stack_deltas;
 extern struct exe_id_to_21_stack_deltas_t exe_id_to_21_stack_deltas;
 extern struct exe_id_to_22_stack_deltas_t exe_id_to_22_stack_deltas;
 extern struct exe_id_to_23_stack_deltas_t exe_id_to_23_stack_deltas;
+extern struct stack_delta_page_to_info_t stack_delta_page_to_info;
+extern struct unwind_info_array_t unwind_info_array;
+
+#if defined(TESTING_COREDUMP)
+
+// References to maps in alphabetical order that
+// are needed only for testing.
+extern struct apm_int_procs_t apm_int_procs;
+extern struct beam_procs_t beam_procs;
 extern struct hotspot_procs_t hotspot_procs;
 extern struct dotnet_procs_t dotnet_procs;
 extern struct perl_procs_t perl_procs;
 extern struct php_procs_t php_procs;
 extern struct py_procs_t py_procs;
 extern struct ruby_procs_t ruby_procs;
-extern struct stack_delta_page_to_info_t stack_delta_page_to_info;
-extern struct unwind_info_array_t unwind_info_array;
 extern struct v8_procs_t v8_procs;
 #endif // TESTING_COREDUMP
 

support/ebpf/native_stack_trace.h

@@ -1,6 +1,10 @@
 #ifndef OPTI_NATIVE_STACK_TRACE_H
 #define OPTI_NATIVE_STACK_TRACE_H
 
+#include "bpfdefs.h"
+#include "extmaps.h"
+#include "tracemgmt.h"
+
 // Unwind info value for invalid stack delta
 #define STACK_DELTA_INVALID (STACK_DELTA_COMMAND_FLAG | UNWIND_COMMAND_INVALID)
 #define STACK_DELTA_STOP    (STACK_DELTA_COMMAND_FLAG | UNWIND_COMMAND_STOP)

support/ebpf/python_tracer.ebpf.c

@@ -2,15 +2,11 @@
 
 #include "bpfdefs.h"
 #include "errors.h"
+#include "native_stack_trace.h"
 #include "tracemgmt.h"
 #include "tsd.h"
 #include "types.h"
 
-// The number of Python frames to unwind per frame-unwinding eBPF program. If
-// we start running out of instructions in the walk_python_stack program, one
-// option is to adjust this number downwards.
-#define FRAMES_PER_WALK_PYTHON_STACK 12
-
 // Forward declaration to avoid warnings like
 // "declaration of 'struct pt_regs' will not be visible outside of this function [-Wvisibility]".
 struct pt_regs;
@@ -141,8 +137,10 @@ static EBPF_INLINE ErrorCode process_python_frame(
   }
 
   // Read PyCodeObject
-  if (bpf_probe_read_user(pss->code, sizeof(pss->code), py_codeobject)) {
-    DEBUG_PRINT("Failed to read PyCodeObject at 0x%lx", (unsigned long)(py_codeobject));
+  long pycode_err = bpf_probe_read_user(pss->code, sizeof(pss->code), py_codeobject);
+  if (pycode_err) {
+    DEBUG_PRINT(
+      "Failed to read PyCodeObject at 0x%lx err=%ld", (unsigned long)(py_codeobject), pycode_err);
     increment_metric(metricID_UnwindPythonErrBadCodeObjectArgCountAddr);
     return ERR_PYTHON_BAD_CODE_OBJECT_ADDR;
   }
@@ -169,39 +167,6 @@ static EBPF_INLINE ErrorCode process_python_frame(
   return ERR_OK;
 }
 
-static EBPF_INLINE ErrorCode
-walk_python_stack(PerCPURecord *record, const PyProcInfo *pyinfo, int *unwinder)
-{
-  void *py_frame  = record->pythonUnwindState.py_frame;
-  ErrorCode error = ERR_OK;
-  *unwinder       = PROG_UNWIND_STOP;
-
-  for (u32 i = 0; i < FRAMES_PER_WALK_PYTHON_STACK; ++i) {
-    bool continue_with_next;
-    error = process_python_frame(record, pyinfo, &py_frame, &continue_with_next);
-    if (error) {
-      goto stop;
-    }
-    if (continue_with_next) {
-      *unwinder = get_next_unwinder_after_interpreter();
-      goto stop;
-    }
-    if (!py_frame) {
-      goto stop;
-    }
-  }
-
-  *unwinder = PROG_UNWIND_PYTHON;
-
-stop:
-  // Set up the state for the next invocation of this unwinding program.
-  if (error || !py_frame) {
-    unwinder_mark_done(record, PROG_UNWIND_PYTHON);
-  }
-  record->pythonUnwindState.py_frame = py_frame;
-  return error;
-}
-
 // get_PyThreadState retrieves the PyThreadState* for the current thread.
 //
 // Python 3.12 and earlier set the thread_state using pthread_setspecific with the key
@@ -288,6 +253,60 @@ static EBPF_INLINE ErrorCode get_PyFrame(const PyProcInfo *pyinfo, void **frame)
   return ERR_OK;
 }
 
+// Number of loop iterations in unwind_python. Each iteration handles either
+// one Python frame or one native frame depending on the current unwinder state.
+// This is a RODATA variable so the host agent can tune it based on whether
+// debug output is enabled (which affects the verifier instruction budget).
+BPF_RODATA_VAR(u32, python_native_loop_iters, 6)
+
+// step_python processes one Python frame and updates *unwinder to indicate
+// what should happen next
+static EBPF_INLINE ErrorCode
+step_python(PerCPURecord *record, const PyProcInfo *pyinfo, void **py_frame, int *unwinder)
+{
+  bool continue_with_next;
+  ErrorCode error = process_python_frame(record, pyinfo, py_frame, &continue_with_next);
+  if (error) {
+    *unwinder = PROG_UNWIND_STOP;
+    return error;
+  }
+  if (continue_with_next) {
+    *unwinder = get_next_unwinder_after_interpreter();
+  } else if (!*py_frame) {
+    *unwinder = PROG_UNWIND_STOP;
+  } else {
+    *unwinder = PROG_UNWIND_PYTHON;
+  }
+  return ERR_OK;
+}
+
+// step_native processes one native frame at an interpreter boundary and
+// updates *unwinder
+static EBPF_INLINE ErrorCode step_native(PerCPURecord *record, int *unwinder)
+{
+  Trace *trace = &record->trace;
+  *unwinder    = PROG_UNWIND_STOP;
+
+  increment_metric(metricID_UnwindNativeAttempts);
+  ErrorCode error = push_native(
+    &record->state,
+    trace,
+    record->state.text_section_id,
+    record->state.text_section_offset,
+    record->state.return_address);
+  if (error) {
+    return error;
+  }
+
+  bool stop;
+  error = unwind_one_frame(record, &stop);
+  if (error || stop) {
+    return error;
+  }
+
+  return get_next_unwinder_after_native_frame(record, unwinder);
+}
+
 // unwind_python is the entry point for tracing when invoked from the native tracer
 // or interpreter dispatcher. It does not reset the trace object and will append the
 // Python stack frames to the trace object for the current CPU.
@@ -298,7 +317,7 @@ static EBPF_INLINE int unwind_python(struct pt_regs *ctx)
     return -1;
 
   ErrorCode error = ERR_OK;
-  int unwinder    = get_next_unwinder_after_interpreter();
+  int unwinder    = PROG_UNWIND_PYTHON;
   Trace *trace    = &record->trace;
   u32 pid         = trace->pid;
 
@@ -327,7 +346,26 @@ static EBPF_INLINE int unwind_python(struct pt_regs *ctx)
     goto exit;
   }
 
-  error = walk_python_stack(record, pyinfo, &unwinder);
+  {
+    void *py_frame = record->pythonUnwindState.py_frame;
+
+    for (u32 t = 0; t < python_native_loop_iters; t++) {
+      switch (unwinder) {
+      case PROG_UNWIND_PYTHON: error = step_python(record, pyinfo, &py_frame, &unwinder); break;
+      case PROG_UNWIND_NATIVE: error = step_native(record, &unwinder); break;
+      default: goto done;
+      }
+      if (error) {
+        goto done;
+      }
+    }
+
+  done:
+    if (error || !py_frame) {
+      unwinder_mark_done(record, PROG_UNWIND_PYTHON);
+    }
+    record->pythonUnwindState.py_frame = py_frame;
+  }
 
 exit:
   record->state.unwind_error = error;

support/ebpf/tracer.ebpf.amd64

@@ -302,6 +302,13 @@ func loadRodataVars(coll *cebpf.CollectionSpec, kmod *kallsyms.Module, cfg *Conf
 		if err := coll.Variables["with_debug_output"].Set(uint32(1)); err != nil {
 			return fmt.Errorf("failed to set debug output: %v", err)
 		}
+	} else {
+		// Without debug output the verifier skips DEBUG_PRINT branches,
+		// leaving enough instruction budget to increase the Python
+		// unwinder loop iterations (default 7 -> 12).
+		if err := coll.Variables["python_native_loop_iters"].Set(uint32(12)); err != nil {
+			return fmt.Errorf("failed to set python_native_loop_iters: %v", err)
+		}
 	}
 
 	if err := coll.Variables["off_cpu_threshold"].Set(cfg.OffCPUThreshold); err != nil {