Profiling middleware

ooniapi/common/src/common/config.py

@@ -106,3 +106,7 @@ class Settings(BaseSettings):
     github_token: str = ""
     origin_repo: str = ""
     push_repo: str = ""
+
+    # Profiling settings
+    profiling_active: bool = False
+    profiling_report_path: str = ""

ooniapi/common/src/common/profile_middleware.py

@@ -0,0 +1,53 @@
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from pathlib import Path
+import logging
+
+log = logging.getLogger(__name__)
+
+class ProfileMiddleware(BaseHTTPMiddleware):
+    """
+    Profiles a request, generating an html report on disk
+    """
+
+    def __init__(self, app, profiling_active : bool, report_path : str, whitelist: tuple[str]):
+        """
+        - profiling_active: whether profiling is enabled
+        - report_path: local disk path where the report is written
+        - whitelist: path prefixes to profile (only matching requests are profiled)
+        """
+        super().__init__(app)
+        self.profiling_active = profiling_active
+        self.report_path = report_path
+        self.whitelist = whitelist
+
+    async def dispatch(self, request: Request, call_next) -> Response:
+
+        if not self.profiling_active or not self.should_profile(request):
+            return await call_next(request)
+
+        # Pyinstrument is only available on development modes
+        from pyinstrument import Profiler
+
+        log.debug(f"Profiling: {request.url.path}")
+
+        profiler = Profiler()
+        profiler.start()
+        response = await call_next(request)
+        profiler.stop()
+
+        # Save report to a file
+        report = profiler.output_html()
+        report_path = Path(self.report_path)
+        report_path.parent.mkdir(exist_ok=True)
+        report_path.touch(exist_ok=True)
+
+        with report_path.open("w") as f:
+            f.write(report)
+            log.debug(f"Report saved to: {report_path.absolute()}")
+
+        return response
+
+    def should_profile(self, request: Request) -> bool:
+        return request.url.path.startswith(self.whitelist)

ooniapi/services/ooniprobe/pyproject.toml

@@ -75,6 +75,7 @@ dependencies = [
   "pytest-asyncio",
   "freezegun",
   "pytest-docker",
+  "pyinstrument"
 ]
 path = ".venv/"
 

ooniapi/services/ooniprobe/src/ooniprobe/main.py

@@ -15,6 +15,7 @@
 
 from . import models
 from .__about__ import VERSION
+from .common.profile_middleware import ProfileMiddleware
 from .common.clickhouse_utils import query_click
 from .common.config import Settings
 from .common.dependencies import ClickhouseDep, SettingsDep, get_settings
@@ -87,6 +88,14 @@ def update_geoip_task():
     allow_headers=["*"],
 )
 
+settings = get_settings()
+app.add_middleware(
+    ProfileMiddleware,
+    profiling_active = settings.profiling_active,
+    report_path = settings.profiling_report_path,
+    whitelist = ("/api/v1/submit_measurement",)
+)
+
 app.include_router(vpn.router, prefix="/api")
 app.include_router(probe_services.router, prefix="/api")
 app.include_router(reports.router)
@@ -186,6 +195,14 @@ async def health(
         "build_label": build_label,
     }
 
+    if settings.profiling_active:
+        try:
+            import pyinstrument  # noqa: F401
+        except ImportError:
+            # In case we set profiling active in a profile that doesn't includes
+            # development tools
+            errors.append("profiling_active_without_pyinstrument")
+
     if len(errors):
         log.error(f"Health check errors detected: {errors}")
 

ooniapi/services/ooniprobe/src/ooniprobe/routers/v1/probe_services.py

@@ -1038,7 +1038,7 @@ async def submit_measurement(
 
     # wasn't possible to send msmnt to fastpath, try to send it to s3
     ts_prefix = now.strftime("%Y%m%d%H")
-    s3_key = f"postcans/{ts_prefix}/{ts_prefix}_{cc}_{tn}/{msmt_uid}.post"
+    s3_key = f"postcans/{ts_prefix}/{ts_prefix}_{cc}_{test_name}/{msmt_uid}.post"
     try:
         await run_in_threadpool(
             request.app.state.s3_client.upload_fileobj,

ooniapi/services/ooniprobe/tests/conftest.py

@@ -19,6 +19,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
+from ooniprobe.common.profile_middleware import ProfileMiddleware
 from ooniprobe.common.clickhouse_utils import insert_click
 from ooniprobe.common.config import Settings
 from ooniprobe.common.dependencies import get_settings
@@ -35,7 +36,7 @@
 from ooniprobe.main import app, lifespan
 from ooniprobe.routers.v1.probe_services import TorTarget
 
-from .utils import setup_user
+from .utils import setup_user, set_middleware_params
 
 
 def make_override_get_settings(**kw):
@@ -461,3 +462,28 @@ async def client_with_two_working_fastpaths(
         test_settings, geoip_db_dir, test_creds, [first_url, second_url]
     ) as (client, mock_fastpath):
         yield client, mock_fastpath, first_url, second_url
+
+
+@pytest_asyncio.fixture(scope='function')
+def profiling_enabled(tmp_path):
+    old = set_middleware_params(app, ProfileMiddleware,
+        profiling_active = True,
+        report_path = str(tmp_path / "report.html"),
+        whitelist = ("/api/v1/manifest",)
+    ) or {}
+
+    yield
+
+    set_middleware_params(app, ProfileMiddleware, **old)
+
+@pytest_asyncio.fixture(scope='function')
+def profiling_disabled(tmp_path):
+    old = set_middleware_params(app, ProfileMiddleware,
+        profiling_active = False,
+        report_path = str(tmp_path / "report.html"),
+        whitelist = ("/api/v1/manifest",)
+    ) or {}
+
+    yield
+
+    set_middleware_params(app, ProfileMiddleware, **old)

ooniapi/services/ooniprobe/tests/test_profiling_middleware.py

@@ -0,0 +1,24 @@
+from .utils import getj
+import pytest
+from pathlib import Path
+
+
+@pytest.mark.asyncio
+async def test_profiling_enabled(client, db, tmp_path, profiling_enabled):
+    report_path: Path = tmp_path / "report.html"
+
+    assert not report_path.exists()
+
+    getj(client, "/api/v1/manifest")
+
+    assert report_path.exists()
+
+@pytest.mark.asyncio
+async def test_profiling_disabled(client, db, tmp_path, profiling_disabled):
+    report_path: Path = tmp_path / "report.html"
+
+    assert not report_path.exists()
+
+    getj(client, "/api/v1/manifest")
+
+    assert not report_path.exists()

ooniapi/services/ooniprobe/tests/utils.py

@@ -55,3 +55,13 @@ def get_msmt_hash(msmt: Dict[str, Any], is_verified: str = "u") -> str:
     payload = copy.deepcopy(msmt)
     payload["is_verified"] = is_verified
     return sha512(ujson.dumps(payload).encode()).hexdigest()[:16]
+
+def set_middleware_params(app, middleware_class, **kwargs):
+    old = None
+    for m in app.user_middleware:
+        if m.cls is middleware_class:
+            old = m.options.copy()
+            m.options.update(kwargs)
+            app.middleware_stack = None  # force Starlette to rebuild on next request
+            break
+    return old