Collecting native backtraces during Python frame execution

backtrace.py

import sys
from cffi import FFI
from numba import jit


ffi = FFI()
ffi.cdef("""
int backtrace(void** array, int size);
char** backtrace_symbols(void* const* array, int size);
void backtrace_symbols_fd(void* const* array, int size, int fd);
""")
C = ffi.dlopen(None)


class NativeTrace:
    def __init__(self):
        ptr_size = 8
        stack_entries = 1024
        callstack = bytearray(ptr_size * stack_entries)
        callstack_ptr = ffi.from_buffer(callstack)
        n_frames = C.backtrace(ffi.cast("void**", callstack_ptr), stack_entries)

        self.callstack = callstack
        self.callstack_ptr = callstack_ptr
        self.n_frames = n_frames

    def print_stack(self):
        strs = C.backtrace_symbols(ffi.cast("void**", self.callstack_ptr),
                                   self.n_frames)
        for i in range(self.n_frames):
            print(f"{i:2d}: {ffi.string(strs[i]).decode()}")


def c_walk_stack():
    trace = NativeTrace()
    trace.print_stack()


def walk_stack():
    f = sys._getframe()
    while f is not None:
        print(f)
        f = f.f_back
    print()
    c_walk_stack()


@jit(forceobj=True)
def numba_walk_stack():
    walk_stack()


def main():
    numba_walk_stack()


native_stacks = []
max_stack_depth = 0
max_native_stack_depth = 0

def native_stack_trace(frame, event, arg):
    #print(f"{event}: {frame}")
    stack_depth = len(native_stacks)
    global max_stack_depth, max_native_stack_depth
    max_stack_depth = max(max_stack_depth, stack_depth)
    #print(f"Stack depth: {len(native_stacks)}")
    if event == 'call':
        trace = NativeTrace()
        native_stacks.append((frame, trace))
        max_native_stack_depth = max(max_native_stack_depth, trace.n_frames)
    elif event == 'return':
        old_frame, _ = native_stacks.pop()
        if frame != old_frame:
            raise RuntimeError("Frames not equal between call and return")
    #trace.print_stack()
    #print("---\n")
    return native_stack_trace


if __name__ == '__main__':
    sys.settrace(native_stack_trace)
    main()
    sys.settrace(None)
    print(f"Max stack depth {max_stack_depth}")
    print(f"Max native stack depth {max_native_stack_depth}")