bpftrace proof of concept for contextvars based thread correlation

Author: aabmass

uprobe-contextvars/Dockerfile

@@ -0,0 +1,2 @@
+FROM archlinux:base-20251019.0.436919
+RUN pacman -Syu --noconfirm bpftrace debuginfod python uv

uprobe-contextvars/README.md

@@ -0,0 +1,65 @@
+## Running
+
+```console
+docker compose up --remove-orphans  --build
+```
+
+## Sample output
+
+```console
+bpftrace-1  | Trying to attach probe: uretprobe:/usr/lib/libpython3.13.so:context_run*
+bpftrace-1  | Attaching to 1 functions
+bpftrace-1  | /usr/lib/libpython3.13.so:context_run
+bpftrace-1  | Trying to attach probe: uprobe:/usr/lib/libpython3.13.so:context_run*
+bpftrace-1  | Attaching to 1 functions
+bpftrace-1  | /usr/lib/libpython3.13.so:context_run
+bpftrace-1  | Trying to attach probe: uprobe:/usr/lib/libpython3.13.so:PyContext_Exit
+bpftrace-1  | Trying to attach probe: uprobe:/usr/lib/libpython3.13.so:PyContext_Enter
+bpftrace-1  | Attached 4 probes
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd36158a98, .nargs = 3, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd36188240, .args = 0x7fdd35f37098, .nargs = 2, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062958 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd3722e938, .nargs = 1, .kwnames = 0x0 }
+python-1    | 200
+bpftrace-1  | tid=3062958 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd35f8b740, .args = 0x7fdd35ff7758, .nargs = 3, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd364f40c0, .args = 0x7fdd3600d588, .nargs = 1, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd364320c0, .args = 0x7fdd35f36118, .nargs = 2, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd36188240, .args = 0x7fdd35f36118, .nargs = 2, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd36158a98, .nargs = 3, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd36188240, .args = 0x7fdd35f36118, .nargs = 2, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062958 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd3722e938, .nargs = 1, .kwnames = 0x0 }
+bpftrace-1  | tid=3062958 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd35f8b740, .args = 0x7fdd35ff7758, .nargs = 3, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd364f40c0, .args = 0x7fdd3600d738, .nargs = 1, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd364320c0, .args = 0x7fdd35fbaf58, .nargs = 2, .kwnames = 0x0 }
+bpftrace-1  | tid=3062951 return context_run
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd36188240, .args = 0x7fdd35fbaf58, .nargs = 2, .kwnames = 0x0 }
+python-1    | 200
+```
+
+In this example `3062951` is the main event loop thread and `3062958` is a thread pool worker,
+given work by `asyncio.to_thread()`. When work is run in the thread pool, asyncio wraps that in
+`context_run()` so you can track which thread started the work. The [self
+argument](https://github.com/python/cpython/blob/v3.13.9/Python/context.c#L649) is a pointer to
+the [Context object](https://docs.python.org/3/library/contextvars.html#contextvars.Context).
+
+In the above example, Context `0x7fdd35f88900` moves between the two threads.
+```console
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd36158a98, .nargs = 3, .kwnames = 0x0 }
+...
+bpftrace-1  | tid=3062958 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd3722e938, .nargs = 1, .kwnames = 0x0 }
+...
+bpftrace-1  | tid=3062951 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd36158a98, .nargs = 3, .kwnames = 0x0 }
+...
+bpftrace-1  | tid=3062958 context_run { .self = 0x7fdd35f88900, .args = 0x7fdd3722e938, .nargs = 1, .kwnames = 0x0 }
+```

uprobe-contextvars/docker-compose.yml

@@ -0,0 +1,31 @@
+# Docker compose file for bpftrace and some python webapps
+
+services:
+  bpftrace:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    privileged: true
+    pid: "host"
+    network_mode: "host"
+    volumes:
+      - .:/app
+      - /sys/:/sys/
+    working_dir: /app
+    environment:
+      - DEBUGINFOD_URLS=https://debuginfod.archlinux.org
+    entrypoint: ["bpftrace", "-v", "print_calls.bt"]
+    
+
+  python:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    environment:
+      - PYTHONUNBUFFERED=1
+    volumes:
+      - .:/app
+    working_dir: /app
+    entrypoint: ["uv", "run", "--script", "plain.py"]
+    depends_on:
+      - bpftrace

uprobe-contextvars/plain.py

@@ -0,0 +1,27 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "requests",
+# ]
+# ///
+
+import asyncio
+
+import requests
+
+
+async def make_request() -> None:
+    # Runs the blocking IO in a thread pool
+    status_code = await asyncio.to_thread(
+        lambda: requests.get("https://github.com").status_code
+    )
+    print(status_code)
+
+
+async def main() -> None:
+    while True:
+        await make_request()
+        await asyncio.sleep(2)
+
+
+asyncio.run(main())

uprobe-contextvars/print_calls.bt

@@ -0,0 +1,61 @@
+#define YELLOW "\033[33m"
+#define NC "\033[0m" // No Color (resets the terminal color)
+
+// These are the official C API for enter/exiting a context, but is not called by the cpython
+// internals. Something like uvloop might use this.
+uprobe:/usr/lib/libpython3.13.so:PyContext_Enter {
+    printf("%stid=%s%s %s %s\n", YELLOW, tid, NC, func, args);
+}
+
+uprobe:/usr/lib/libpython3.13.so:PyContext_Exit {
+    printf("%stid=%s%s %s %s\n", YELLOW, tid, NC, func, args);
+}
+
+// This is the C implementation of contextvars.run() python function
+// https://github.com/python/cpython/blob/v3.13.9/Python/context.c#L648-L650. This plus the
+// PyContext_Enter/PyContext_Exit is probably the most reliable.
+uprobe:/usr/lib/libpython3.13.so:context_run* {
+    printf("%stid=%s%s %s %s\n", YELLOW, tid, NC, func, args);
+}
+uretprobe:/usr/lib/libpython3.13.so:context_run* {
+    printf("%stid=%s%s return %s\n", YELLOW, tid, NC, func);
+}
+
+// -----------
+
+// These might be useful for tracking parentage of new Context objects
+/*
+uretprobe:/usr/lib/libpython3.13.so:PyContext_CopyCurrent {
+    printf("%s retval = %#x\n", func, retval);
+}
+uprobe:/usr/lib/libpython3.13.so:PyContext_Copy {
+    printf("%s %s\n", func, args);
+}
+uretprobe:/usr/lib/libpython3.13.so:PyContext_Copy {
+    printf("%s retval = %#x\n", func, retval);
+}
+*/
+
+// These might be useful for watching the trace context for OTel python instrumented code
+/*
+uprobe:/usr/lib/libpython3.13.so:PyContextVar_Get {
+    printf("%s %s\n", func, args);
+}
+uprobe:/usr/lib/libpython3.13.so:PyContextVar_Set {
+    printf("%s %s\n", func, args);
+}
+*/
+
+// These are called by PyContext_Enter/PyContext_Exit and context_run internally. It seems to
+// get inlined in many production builds of python (e.g. on Debian and uv managed python), so
+// might not be reliable.
+//
+// See https://github.com/python/cpython/blob/v3.13.9/Python/context.c#L112-L113
+/*
+uprobe:/usr/lib/libpython3.13.so:_PyContext_Enter {
+    printf("tid=%s %s %s\n", tid, func, args);
+}
+uprobe:/usr/lib/libpython3.13.so:_PyContext_Exit {
+    printf("tid=%s %s %s\n", tid, func, args);
+}
+*/