Add thread roots

We are seeing a few mysterious cases where there is a strongly-connected
component that's just floating around.

This change now adds the code location to `builtins.frame`s and also
grabs all the frames of all the threads to see what the locals are.

Author: lhchavez

analyze_heap.py

@@ -26,6 +26,7 @@ class HeapObject:
     addr: int
     typename: str
     root: bool
+    thread_root: bool
     children_size_exclusive: int = 0
     # Pointers from this object to other objects.
     referents: set[int] = dataclasses.field(default_factory=set)
@@ -61,6 +62,7 @@ def _scanheap(filename: str, populate_referrers: bool = True) -> dict[int, HeapO
                     typename=typenames[record["objtype_addr"]],
                     size=record["size"],
                     root=record["root"],
+                    thread_root=record.get("thread_root", False),
                     payload=payload,
                 )
             elif record["t"] == "referents":
@@ -109,7 +111,11 @@ def _top(args: argparse.Namespace) -> None:
             print(
                 f"- {obj.size:8d} {obj.addr:016x} {obj.typename} "
                 + (repr(obj.payload) if obj.payload is not None else "")
-                + (" (root)" if obj.root else ""),
+                + (
+                    " (root)"
+                    if obj.root
+                    else (" (thread root)" if obj.thread_root else "")
+                ),
             )
             if args.show_parents:
                 seen = set()
@@ -126,7 +132,11 @@ def _top(args: argparse.Namespace) -> None:
                             else ""
                         )
                         + (" (cycle)" if addr in seen else "")
-                        + (" (root)" if child_obj.root else ""),
+                        + (
+                            " (root)"
+                            if child_obj.root
+                            else (" (thread root)" if child_obj.thread_root else "")
+                        ),
                     )
                     if addr in seen:
                         continue
@@ -235,35 +245,40 @@ def _render_graph(
 
             style = ""
             if obj.addr in exclude_addresses:
-                style = ",style=filled,fillcolor=gray"
+                style += ",style=filled,fillcolor=gray"
             elif len(path) == 1:
-                style = ",style=filled,fillcolor=red"
+                style += ",style=filled,fillcolor=red"
             elif highlight and highlight in obj.typename:
-                style = ",style=filled,fillcolor=yellow"
+                style += ",style=filled,fillcolor=yellow"
             payload = ""
             if obj.payload is not None:
                 if len(payload) <= 32:
                     payload = f"\\n{obj.payload}"
                 else:
                     payload = f"\\n{obj.payload[:31]}…"
+            label = f"0x{obj.addr:x}"
+            if obj.root:
+                label += " (gcroot)"
+            elif obj.thread_root:
+                label += " (thread root)"
             if censor_prefixes:
                 if any(
                     obj.typename.startswith(censored) for censored in censor_prefixes
                 ):
                     output.write(
-                        f'  x{obj.addr:x} [label="0x{obj.addr:x}\\n[omitted]\\n{obj.size}"{style}];\n'.encode(
+                        f'  x{obj.addr:x} [label="{label}\\n[omitted]\\n{obj.size}"{style}];\n'.encode(
                             "utf-8"
                         )
                     )
                 else:
                     output.write(
-                        f'  x{obj.addr:x} [label="0x{obj.addr:x}\\n{obj.typename}\\n{obj.size}"{style}];\n'.encode(
+                        f'  x{obj.addr:x} [label="{label}\\n{obj.typename}\\n{obj.size}"{style}];\n'.encode(
                             "utf-8"
                         )
                     )
             else:
                 output.write(
-                    f'  x{obj.addr:x} [label="0x{obj.addr:x}\\n{obj.typename}\\n{obj.size}{payload}"{style}];\n'.encode(
+                    f'  x{obj.addr:x} [label="{label}\\n{obj.typename}\\n{obj.size}{payload}"{style}];\n'.encode(
                         "utf-8"
                     )
                 )

src/payloads/dump_heap.py

@@ -6,6 +6,7 @@ def __payload_entrypoint(output_path: str) -> None:
     import gc
     import struct
     import sys
+    import threading
     from typing import Any
 
     DEFAULT_MAX_PAYLOAD_SIZE = 128
@@ -110,19 +111,28 @@ def _pack_type(objtype_addr: int, typename: bytes) -> bytes:
         )
 
     def _pack_object(
-        *, root: bool, addr: int, objtype_addr: int, size: int, payload: bytes | None
+        *,
+        root: bool,
+        thread_root: bool,
+        addr: int,
+        objtype_addr: int,
+        size: int,
+        payload: bytes | None,
     ) -> bytes:
         if payload:
             return b"".join(
                 (
-                    _pack_map(6),
+                    _pack_map(7),
                     #
                     _pack_str(b"t"),
                     _pack_str(b"object"),
                     #
                     _pack_str(b"root"),
                     _pack_bool(root),
                     #
+                    _pack_str(b"thread_root"),
+                    _pack_bool(thread_root),
+                    #
                     _pack_str(b"objtype_addr"),
                     _pack_int(objtype_addr),
                     #
@@ -139,14 +149,17 @@ def _pack_object(
         else:
             return b"".join(
                 (
-                    _pack_map(5),
+                    _pack_map(6),
                     #
                     _pack_str(b"t"),
                     _pack_str(b"object"),
                     #
                     _pack_str(b"root"),
                     _pack_bool(root),
                     #
+                    _pack_str(b"thread_root"),
+                    _pack_bool(thread_root),
+                    #
                     _pack_str(b"objtype_addr"),
                     _pack_int(objtype_addr),
                     #
@@ -232,6 +245,8 @@ def _get_payload(obj: Any) -> bytes | None:
                     if source_traceback
                     else repr(obj)
                 )
+            elif inspect.isframe(obj):
+                payload_str = f"{obj.f_code.co_qualname} at {obj.f_code.co_filename}:{obj.f_lineno}"
             elif inspect.ismodule(obj):
                 payload_str = obj.__name__
             elif inspect.isclass(obj):
@@ -256,9 +271,16 @@ def _get_payload(obj: Any) -> bytes | None:
 
     try:
         with open(output_path, "wb") as output_file:
-            queue: list[tuple[Any, int, bool]] = [
-                (o, id(o), True) for o in gc.get_objects()
+            # Add all the gc-known objects into the heap dump.
+            queue: list[tuple[Any, int, bool, bool]] = [
+                (o, id(o), True, False) for o in gc.get_objects()
             ]
+            # Now also add all the thread-local frames, excluding the one from this method.
+            current_thread_id = threading.current_thread().ident
+            for thread_id, frame in sys._current_frames().items():
+                if current_thread_id == thread_id:
+                    continue
+                queue.append((frame, id(frame), False, True))
             x = 0
             totalsize = 0
 
@@ -272,7 +294,7 @@ def _get_payload(obj: Any) -> bytes | None:
                     )
                 x += 1
 
-                obj, addr, root = queue.pop()
+                obj, addr, gcroot, thread_root = queue.pop()
                 if addr in ignored_addrs:
                     # We don't want to track the objects we own.
                     continue
@@ -299,7 +321,8 @@ def _get_payload(obj: Any) -> bytes | None:
                 totalsize += size
                 output_file.write(
                     _pack_object(
-                        root=root,
+                        root=gcroot,
+                        thread_root=thread_root,
                         addr=addr,
                         objtype_addr=objtype_addr,
                         size=size,
@@ -312,7 +335,7 @@ def _get_payload(obj: Any) -> bytes | None:
                 for child_obj in referents:
                     child_addr = id(child_obj)
                     child_addrs.append(child_addr)
-                    queue.append((child_obj, child_addr, False))
+                    queue.append((child_obj, child_addr, False, False))
                 if child_addrs:
                     output_file.write(
                         _pack_referents(addr=addr, child_addrs=child_addrs)