Add tool for profiling (#364)

noxfile.py

@@ -2,7 +2,7 @@ import nox
 
 nox.options.stop_on_first_error = True
 
-source_files = ("httpx", "tests", "setup.py", "noxfile.py")
+source_files = ("httpx", "tools", "tests", "setup.py", "noxfile.py")
 
 
 @nox.session(reuse_venv=True)

setup.cfg

@@ -10,8 +10,8 @@ ignore_missing_imports = True
 combine_as_imports = True
 force_grid_wrap = 0
 include_trailing_comma = True
-known_first_party = httpx,tests
-known_third_party = brotli,certifi,chardet,cryptography,h11,h2,hstspreload,nox,pytest,requests,rfc3986,setuptools,trio,trustme,uvicorn
+known_first_party = httpx,httpxprof,tests
+known_third_party = brotli,certifi,chardet,click,cryptography,h11,h2,hstspreload,nox,pytest,requests,rfc3986,setuptools,tqdm,trio,trustme,uvicorn
 line_length = 88
 multi_line_output = 3
 

tools/httpxprof/.gitignore

@@ -0,0 +1 @@
+out/

tools/httpxprof/README.md

@@ -0,0 +1,27 @@
+# httpxprof
+
+A tool for profiling [HTTPX](https://github.com/encode/httpx) using cProfile and [SnakeViz](https://jiffyclub.github.io/snakeviz/).
+
+## Usage
+
+```bash
+# Run one of the scripts:
+httpxprof run async
+
+# View results:
+httpxprof view async
+```
+
+You can ask for `--help` on `httpxprof` and any of the subcommands.
+
+## Installation
+
+```bash
+# From the HTTPX project root directory:
+pip install -e tools/httpxprof
+
+# From this directory:
+pip install -e .
+```
+
+`httpxprof` assumes it can `import httpx`, so you need to have HTTPX installed (either from local or PyPI).

tools/httpxprof/httpxprof/__main__.py

@@ -0,0 +1,5 @@
+import sys
+
+from .cli import cli
+
+sys.exit(cli())

tools/httpxprof/httpxprof/config.py

@@ -0,0 +1,9 @@
+import pathlib
+
+SERVER_HOST = "127.0.0.1"
+SERVER_PORT = 8123
+SERVER_URL = f"http://{SERVER_HOST}:{SERVER_PORT}"
+
+OUTPUT_DIR = pathlib.Path(__file__).parent / "out"
+SCRIPTS_DIR = pathlib.Path(__file__).parent / "scripts"
+assert SCRIPTS_DIR.exists(), SCRIPTS_DIR

tools/httpxprof/httpxprof/main.py

@@ -0,0 +1,45 @@
+import os
+import subprocess
+
+import click
+
+from .config import OUTPUT_DIR, SCRIPTS_DIR, SERVER_HOST, SERVER_PORT
+from .utils import server
+
+SCRIPTS = [
+    filename.rstrip(".py")
+    for filename in os.listdir(SCRIPTS_DIR)
+    if filename != "__init__.py"
+]
+
+
+@click.group()
+def cli() -> None:
+    pass
+
+
+@cli.command()
+@click.argument("script", type=click.Choice(SCRIPTS))
+def run(script: str) -> None:
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+    out = str(OUTPUT_DIR / f"{script}.prof")
+    target = str(SCRIPTS_DIR / f"{script}.py")
+
+    args = ["python", "-m", "cProfile", "-o", out, target]
+
+    with server(host=SERVER_HOST, port=SERVER_PORT):
+        subprocess.run(args)
+
+
+@cli.command()
+@click.argument("script", type=click.Choice(SCRIPTS))
+def view(script: str) -> None:
+    args = ["snakeviz", str(OUTPUT_DIR / f"{script}.prof")]
+    subprocess.run(args)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(cli())

tools/httpxprof/httpxprof/scripts/async.py

@@ -0,0 +1,15 @@
+import asyncio
+
+import tqdm
+
+import httpx
+from httpxprof.config import SERVER_URL
+
+
+async def main() -> None:
+    async with httpx.AsyncClient() as client:
+        for _ in tqdm.tqdm(range(1000)):
+            await client.get(SERVER_URL)
+
+
+asyncio.run(main())

tools/httpxprof/httpxprof/scripts/sync.py

@@ -0,0 +1,13 @@
+import tqdm
+
+import httpx
+from httpxprof.config import SERVER_URL
+
+
+def main() -> None:
+    with httpx.Client() as client:
+        for _ in tqdm.tqdm(range(1000)):
+            client.get(SERVER_URL)
+
+
+main()

tools/httpxprof/httpxprof/utils.py

@@ -0,0 +1,49 @@
+import contextlib
+import multiprocessing
+import time
+import typing
+
+import uvicorn
+
+
+async def app(scope: dict, receive: typing.Callable, send: typing.Callable) -> None:
+    assert scope["type"] == "http"
+    res = b"Hello, world"
+    await send(
+        {
+            "type": "http.response.start",
+            "status": 200,
+            "headers": [
+                [b"content-type", b"text/plain"],
+                [b"content-length", b"%d" % len(res)],
+            ],
+        }
+    )
+    await send({"type": "http.response.body", "body": res})
+
+
+@contextlib.contextmanager
+def server(host: str, port: int) -> typing.Iterator[None]:
+    config = uvicorn.Config(
+        app=app,
+        host=host,
+        port=port,
+        lifespan="off",
+        loop="asyncio",
+        log_level="warning",
+    )
+    server = uvicorn.Server(config)
+
+    proc = multiprocessing.Process(target=server.run)
+    proc.start()
+
+    # Wait a bit for the uvicorn server process to be ready to accept connections.
+    time.sleep(0.2)
+    print(f"Server started at {host}:{port}.")
+
+    try:
+        yield
+    finally:
+        print("Stopping server...")
+        proc.terminate()
+        proc.join()

tools/httpxprof/setup.py

@@ -0,0 +1,20 @@
+import typing
+from pathlib import Path
+
+from setuptools import setup
+
+
+def get_packages(package: str) -> typing.List[str]:
+    return [str(path.parent) for path in Path(package).glob("**/__init__.py")]
+
+
+setup(
+    name="httpxprof",
+    version="0.1",
+    packages=get_packages("httpxprof"),
+    install_requires=["click", "snakeviz", "uvicorn", "tqdm"],
+    entry_points="""
+        [console_scripts]
+        httpxprof=httpxprof.main:cli
+    """,
+)