contextvars bpftrace examples

Author: aabmass

uprobe-contextvars/.python-version

@@ -0,0 +1 @@
+3.13

uprobe-contextvars/README.md

@@ -0,0 +1,40 @@
+# Docker compose file for bpftrace and some python webapps
+
+services:
+  bpftrace:
+    build:
+      context: .
+      dockerfile: Dockerfile.bpftrace
+    privileged: true
+    pid: "host"
+    network_mode: "host"
+    volumes:
+      - ./bpftrace:/app/bpftrace
+    working_dir: /app/bpftrace
+    entrypoint: ["/bin/bash", "-c", "tail -f /dev/null"] # Keep container running
+  
+  flask_app:
+    build:
+      context: .
+      dockerfile: Dockerfile.flask
+    network_mode: "host"
+    environment:
+      - FLASK_APP=app.py
+      - FLASK_RUN_HOST=0.0.0.0
+      - FLASK_RUN_PORT=5000
+    volumes:
+      - ./flask_app:/app/flask_app
+    working_dir: /app/flask_app
+    command: ["flask", "run"]
+
+  django_app:
+    build:
+      context: .
+      dockerfile: Dockerfile.django
+    network_mode: "host"
+    environment:
+      - PYTHONUNBUFFERED=1
+    volumes:
+      - ./django_app:/app/django_app
+    working_dir: /app/django_app
+    command: ["python", "manage.py", "runserver", "0.0.0.0:8000"]

uprobe-contextvars/contextvars.bt

@@ -0,0 +1,15 @@
+from typing import Union
+
+from fastapi import FastAPI
+
+app = FastAPI()
+
+
+@app.get("/")
+def read_root():
+    return {"Hello": "World"}
+
+
+@app.get("/items/{item_id}")
+def read_item(item_id: int, q: Union[str, None] = None):
+    return {"item_id": item_id, "q": q}

uprobe-contextvars/docker-compose.yml

@@ -0,0 +1,107 @@
+import asyncio
+import urllib.request
+import urllib.error
+import sys
+
+
+async def fetch_url(url):
+    """Fetches a URL using urllib.request in a non-blocking way."""
+    try:
+        # urlopen is blocking, so run it in a thread pool executor
+        response_obj = await asyncio.to_thread(urllib.request.urlopen, url)
+        with response_obj as response:
+            return response.read().decode(
+                "utf-8", errors="ignore"
+            )  # Use errors='ignore' for robust decoding
+    except urllib.error.URLError as e:
+        return f"Error fetching {url}: {e.reason}"
+    except Exception as e:
+        return f"An unexpected error occurred while fetching {url}: {e}"
+
+
+async def handle_client(reader, writer):
+    addr = writer.get_extra_info("peername")
+    print(f"Connection from {addr}")
+
+    try:
+        request_data = await reader.read(
+            4096
+        )  # Read more data for potentially larger requests
+        if not request_data:
+            print(f"Client {addr} disconnected before sending data.")
+            writer.close()
+            await writer.wait_closed()
+            return
+
+        request_line = request_data.decode().splitlines()[0]
+        print(f"Received request: {request_line}")
+
+        response_body = ""
+        status_code = 200
+        status_message = "OK"
+        content_type = "text/plain"
+
+        # For simplicity, we'll just check if it's a GET request
+        if request_line.startswith("GET"):
+            github_content = await fetch_url("https://github.com")
+
+            # Truncate for brevity, as github.com's homepage can be very large
+            display_content = github_content
+            if len(github_content) > 1000:
+                display_content = github_content[:1000] + "\n\n... (content truncated)"
+
+            response_body = f"Hello from the Asyncio Server!\n\nFetched content from github.com:\n\n{display_content}"
+        else:
+            status_code = 405
+            status_message = "Method Not Allowed"
+            response_body = "Only GET requests are supported."
+
+        response = (
+            f"HTTP/1.1 {status_code} {status_message}\r\n"
+            f"Content-Type: {content_type}\r\n"
+            f"Content-Length: {len(response_body.encode('utf-8'))}\r\n"  # Encode to get byte length
+            f"\r\n"
+            f"{response_body}"
+        )
+        writer.write(response.encode("utf-8"))
+        await writer.drain()
+
+    except Exception as e:
+        print(f"Error handling client {addr}: {e}", file=sys.stderr)
+        # Attempt to send an error response if possible
+        error_body = f"Internal Server Error: {e}"
+        error_response = (
+            f"HTTP/1.1 500 Internal Server Error\r\n"
+            f"Content-Type: text/plain\r\n"
+            f"Content-Length: {len(error_body.encode('utf-8'))}\r\n"
+            f"\r\n"
+            f"{error_body}"
+        )
+        writer.write(error_response.encode("utf-8"))
+        await writer.drain()
+    finally:
+        print(f"Closing connection for {addr}")
+        writer.close()
+        await writer.wait_closed()
+
+
+async def main():
+    host = "127.0.0.1"
+    port = 8080
+
+    server = await asyncio.start_server(handle_client, host, port)
+
+    print(f"Serving on {host}:{port}")
+
+    async with server:
+        await server.serve_forever()
+
+
+if __name__ == "__main__":
+    try:
+        asyncio.run(main())
+    except KeyboardInterrupt:
+        print("Server stopped by user.")
+    except Exception as e:
+        print(f"An unexpected error occurred in main: {e}", file=sys.stderr)
+        sys.exit(1)

uprobe-contextvars/fastapi_main.py

@@ -0,0 +1,26 @@
+import contextvars
+
+foovar = contextvars.ContextVar("foovar")
+root = contextvars.copy_context()
+
+c1 = root.copy()
+c2 = root.copy()
+
+c1.run(foovar.set, "c1")
+c2.run(foovar.set, "c2")
+
+print(f"{c1=} {c1.get(foovar, None)=}")
+print(f"{c2=} {c2.get(foovar, None)=}")
+
+
+def a():
+    foovar.set("foo")
+    print(foovar.get())
+
+
+def b():
+    print(foovar.get(None))
+
+
+c1.run(b)
+c2.run(b)

uprobe-contextvars/main.py

@@ -0,0 +1,33 @@
+/**
+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);
+}
+*/
+
+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);
+}
+
+uprobe:/usr/lib/libpython3.13.so:PyContext_Enter {
+    printf("%s %s\n", func, args);
+}
+
+uprobe:/usr/lib/libpython3.13.so:PyContext_Exit {
+    printf("%s %s\n", func, args);
+}
+
+uprobe:/usr/lib/libpython3.13.so:_PyContext_Enter {
+    printf("%s %s\n", func, args);
+}
+
+uprobe:/usr/lib/libpython3.13.so:_PyContext_Exit {
+    printf("%s %s\n", func, args);
+}

uprobe-contextvars/plain.py

@@ -0,0 +1,7 @@
+[project]
+name = "uprobe-contextvars"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.13"
+dependencies = []

uprobe-contextvars/print_calls.bt

@@ -0,0 +1,24 @@
+import asyncio
+import contextvars
+
+
+foovar = contextvars.ContextVar("foovar")
+
+
+async def print_foo():
+    foovar.set("foo")
+    await asyncio.sleep(0.1)
+    print(foovar.get())
+
+
+async def main():
+    async with asyncio.TaskGroup() as tg:
+        foovar.set("bar")
+        tg.create_task(print_foo())
+        await asyncio.sleep(0.01)
+        print(foovar.get())
+        print()
+
+
+if __name__ == "__main__":
+    loop = asyncio.run(main(), debug=True)