stream-strace.py

"""Demonstrate streaming with strace.

https://github.com/strace/strace

Usage:
    # Start the server in one terminal:
    uv run python tools/demo_streaming.py serve

    # In another terminal, strace the download:
    strace -f -e trace=recvfrom,write -o trace.log uv run tools/demo_streaming.py download

    # You should see interleaved recvfrom and write calls in trace.log (on Linux).
"""

import sys
import tempfile
from http.server import BaseHTTPRequestHandler, HTTPServer

HOST, PORT = "127.0.0.1", 18232
URL = f"http://{HOST}:{PORT}/data"
TOTAL = 10 * 1024 * 1024  # 10 MB
FLUSH_EVERY = 512 * 1024  # server flushes every 512 KB


class Handler(BaseHTTPRequestHandler):
    def do_GET(self):  # noqa: N802
        self.send_response(200)
        self.send_header("Content-Length", str(TOTAL))
        self.end_headers()
        sent = 0
        chunk = b"x" * FLUSH_EVERY
        while sent < TOTAL:
            self.wfile.write(chunk[: min(FLUSH_EVERY, TOTAL - sent)])
            self.wfile.flush()
            sent += FLUSH_EVERY

    def log_message(self, *a):
        print(f"Received request: {a}")  # noqa: T201
        pass


def serve():
    print(f"Serving {TOTAL // 1024 // 1024} MB on {URL}")
    HTTPServer((HOST, PORT), Handler).serve_forever()


def download():
    import requests

    with tempfile.NamedTemporaryFile(mode="wb", delete=True) as tmp:
        r = requests.get(URL, stream=True)
        for chunk in r.iter_content(chunk_size=8192):
            tmp.write(chunk)


if __name__ == "__main__":
    cmd = sys.argv[1] if len(sys.argv) > 1 else "serve"
    if cmd == "serve":
        serve()
    elif cmd == "download":
        download()