From deedbe956b979ce70a14eac6d490143caf02d251 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 11:34:08 +0200 Subject: [PATCH 1/8] Add profiling library and settings to ooniprobe --- ooniapi/common/src/common/config.py | 4 ++++ ooniapi/services/ooniprobe/pyproject.toml | 1 + 2 files changed, 5 insertions(+) diff --git a/ooniapi/common/src/common/config.py b/ooniapi/common/src/common/config.py index 57e9e1819..e9a63124b 100644 --- a/ooniapi/common/src/common/config.py +++ b/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 = "" diff --git a/ooniapi/services/ooniprobe/pyproject.toml b/ooniapi/services/ooniprobe/pyproject.toml index ba4c30590..d19f077a8 100644 --- a/ooniapi/services/ooniprobe/pyproject.toml +++ b/ooniapi/services/ooniprobe/pyproject.toml @@ -75,6 +75,7 @@ dependencies = [ "pytest-asyncio", "freezegun", "pytest-docker", + "pyinstrument" ] path = ".venv/" From 737c8087dfa08907beefaff5cf4795ad11888e1f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 11:46:17 +0200 Subject: [PATCH 2/8] Add initial version of profiling middleware --- .../common/src/common/profile_middleware.py | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 ooniapi/common/src/common/profile_middleware.py diff --git a/ooniapi/common/src/common/profile_middleware.py b/ooniapi/common/src/common/profile_middleware.py new file mode 100644 index 000000000..66ab2afd8 --- /dev/null +++ b/ooniapi/common/src/common/profile_middleware.py @@ -0,0 +1,39 @@ +from starlette.middleware.base import BaseHTTPMiddleware +from starlette.requests import Request +from starlette.responses import Response +from pathlib import Path + + +class ProfileMiddleware(BaseHTTPMiddleware): + """ + Profiles a request, generating an html report on disk + """ + + def __init__(self, app, profiling_active : bool, report_path : str): + super().__init__(app) + self.profiling_active = profiling_active + self.report_path = report_path + + async def dispatch(self, request: Request, call_next) -> Response: + + if not self.profiling_active: + return await call_next(request) + + # Pyinstrument is only available on development modes + from pyinstrument import Profiler + + 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) + + return response From f7b678bb046d0cef5b64825143aa52dbc8335cd5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 11:55:23 +0200 Subject: [PATCH 3/8] Add whitelisting to profiling --- ooniapi/common/src/common/profile_middleware.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/ooniapi/common/src/common/profile_middleware.py b/ooniapi/common/src/common/profile_middleware.py index 66ab2afd8..f241169f7 100644 --- a/ooniapi/common/src/common/profile_middleware.py +++ b/ooniapi/common/src/common/profile_middleware.py @@ -9,14 +9,20 @@ class ProfileMiddleware(BaseHTTPMiddleware): Profiles a request, generating an html report on disk """ - def __init__(self, app, profiling_active : bool, report_path : str): + 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: + if not self.profiling_active or not self.should_profile(request): return await call_next(request) # Pyinstrument is only available on development modes @@ -37,3 +43,6 @@ async def dispatch(self, request: Request, call_next) -> Response: f.write(report) return response + + def should_profile(self, request: Request) -> bool: + return request.url.path.startswith(self.whitelist) From aa6012bd8c0d61824b56357fd1d1faecd3a98ffb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 12:09:25 +0200 Subject: [PATCH 4/8] Add profiling to ooniprobe --- ooniapi/common/src/common/profile_middleware.py | 5 +++++ .../services/ooniprobe/src/ooniprobe/main.py | 17 +++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/ooniapi/common/src/common/profile_middleware.py b/ooniapi/common/src/common/profile_middleware.py index f241169f7..9ae4f96bd 100644 --- a/ooniapi/common/src/common/profile_middleware.py +++ b/ooniapi/common/src/common/profile_middleware.py @@ -2,7 +2,9 @@ from starlette.requests import Request from starlette.responses import Response from pathlib import Path +import logging +log = logging.getLogger(__name__) class ProfileMiddleware(BaseHTTPMiddleware): """ @@ -28,6 +30,8 @@ async def dispatch(self, request: Request, call_next) -> Response: # 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) @@ -41,6 +45,7 @@ async def dispatch(self, request: Request, call_next) -> Response: with report_path.open("w") as f: f.write(report) + log.debug(f"Report saved to: {report_path.absolute()}") return response diff --git a/ooniapi/services/ooniprobe/src/ooniprobe/main.py b/ooniapi/services/ooniprobe/src/ooniprobe/main.py index baf4fd970..6309a1062 100644 --- a/ooniapi/services/ooniprobe/src/ooniprobe/main.py +++ b/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, + profile_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}") From 00da33758c7b44e71bb07e2d6a4bd65d10221256 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 12:30:33 +0200 Subject: [PATCH 5/8] Fix bad arg name --- ooniapi/services/ooniprobe/src/ooniprobe/main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ooniapi/services/ooniprobe/src/ooniprobe/main.py b/ooniapi/services/ooniprobe/src/ooniprobe/main.py index 6309a1062..f7e58f570 100644 --- a/ooniapi/services/ooniprobe/src/ooniprobe/main.py +++ b/ooniapi/services/ooniprobe/src/ooniprobe/main.py @@ -91,7 +91,7 @@ def update_geoip_task(): settings = get_settings() app.add_middleware( ProfileMiddleware, - profile_active = settings.profiling_active, + profiling_active = settings.profiling_active, report_path = settings.profiling_report_path, whitelist = ("/api/v1/submit_measurement",) ) From a7b3887dd9e6fa032e7608d338301f034237aa0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 12:30:57 +0200 Subject: [PATCH 6/8] fix variable name bug --- .../ooniprobe/src/ooniprobe/routers/v1/probe_services.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ooniapi/services/ooniprobe/src/ooniprobe/routers/v1/probe_services.py b/ooniapi/services/ooniprobe/src/ooniprobe/routers/v1/probe_services.py index b4535baab..c7a790606 100644 --- a/ooniapi/services/ooniprobe/src/ooniprobe/routers/v1/probe_services.py +++ b/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, From 4ba40021c3b52a7804bb6bbd831d31a85705a5bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 16:39:21 +0200 Subject: [PATCH 7/8] Add tests for profiling middleware --- ooniapi/services/ooniprobe/tests/conftest.py | 28 +++++++++++++++++++- ooniapi/services/ooniprobe/tests/utils.py | 10 +++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/ooniapi/services/ooniprobe/tests/conftest.py b/ooniapi/services/ooniprobe/tests/conftest.py index ded810123..5b205b951 100644 --- a/ooniapi/services/ooniprobe/tests/conftest.py +++ b/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) diff --git a/ooniapi/services/ooniprobe/tests/utils.py b/ooniapi/services/ooniprobe/tests/utils.py index e32d8e7e1..f062afaf4 100644 --- a/ooniapi/services/ooniprobe/tests/utils.py +++ b/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 From 78042327342ee8759d8579a7c80f00ee321e3517 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20D=C3=ADaz?= Date: Wed, 22 Jul 2026 16:39:57 +0200 Subject: [PATCH 8/8] Add profiling middleware tests --- .../tests/test_profiling_middleware.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 ooniapi/services/ooniprobe/tests/test_profiling_middleware.py diff --git a/ooniapi/services/ooniprobe/tests/test_profiling_middleware.py b/ooniapi/services/ooniprobe/tests/test_profiling_middleware.py new file mode 100644 index 000000000..6d0da1660 --- /dev/null +++ b/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()