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/common/src/common/profile_middleware.py b/ooniapi/common/src/common/profile_middleware.py new file mode 100644 index 000000000..9ae4f96bd --- /dev/null +++ b/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) 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/" diff --git a/ooniapi/services/ooniprobe/src/ooniprobe/main.py b/ooniapi/services/ooniprobe/src/ooniprobe/main.py index baf4fd970..f7e58f570 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, + 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}") 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, 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/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() 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