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/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,58 @@
+#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);
+}
+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);
+}
+*/