perf: add native json request logging fast path

Signed-off-by: lucarlig <luca.carlig@ibm.com>

Author: lucarlig

mcpgateway/middleware/request_logging_middleware.py

@@ -289,6 +289,24 @@ def mask_sensitive_headers(headers):
     return masked_headers
 
 
+def _mask_json_payload_for_logging(payload: bytes, max_depth: int = 10) -> str:
+    """Mask a JSON payload for logging, preferring the native bytes fast path."""
+    if getattr(settings, "experimental_rust_request_logging_masking_enabled", False) is True:
+        rust_module = _load_rust_request_logging_module()
+        if rust_module is not None and hasattr(rust_module, "mask_sensitive_json_bytes"):
+            try:
+                masked_payload = rust_module.mask_sensitive_json_bytes(payload, max_depth)
+                if isinstance(masked_payload, bytes):
+                    return masked_payload.decode("utf-8", errors="ignore")
+                return str(masked_payload)
+            except Exception:
+                pass
+
+    json_payload = orjson.loads(payload)
+    payload_to_log = mask_sensitive_data(json_payload, max_depth)
+    return orjson.dumps(payload_to_log).decode()
+
+
 def _load_rust_request_logging_module():
     """Load the experimental Rust masking native extension on demand."""
     global _RUST_REQUEST_LOGGING_IMPORT_FAILED, _RUST_REQUEST_LOGGING_MODULE
@@ -630,16 +648,13 @@ async def dispatch(self, request: Request, call_next: Callable):
                 truncated = False
                 body_to_log = body
 
-            payload = body_to_log.decode("utf-8", errors="ignore").strip()
+            payload = body_to_log.strip()
             if payload:
                 try:
-                    json_payload = orjson.loads(payload)
-                    payload_to_log = mask_sensitive_data(json_payload)
-                    # Use orjson without indent for performance (compact output)
-                    payload_str = orjson.dumps(payload_to_log).decode()
+                    payload_str = _mask_json_payload_for_logging(payload)
                 except orjson.JSONDecodeError:
                     # For non-JSON payloads, still mask potential sensitive data
-                    payload_str = payload
+                    payload_str = payload.decode("utf-8", errors="ignore")
                     for sensitive_key in SENSITIVE_KEYS:
                         if sensitive_key in payload_str.lower():
                             payload_str = "<contains sensitive data - masked>"

tests/performance/test_request_logging_masking_native_extension_benchmark.py

@@ -8,17 +8,125 @@
 import statistics
 import subprocess
 import time
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Callable
 
 # First-Party
 from mcpgateway.config import settings
+from mcpgateway.middleware import request_logging_middleware
 from mcpgateway.middleware.request_logging_middleware import mask_sensitive_data, mask_sensitive_headers
 
 REPO_ROOT = Path(__file__).resolve().parents[2]
 NATIVE_EXTENSION_MANIFEST = REPO_ROOT / "tools_rust" / "request_logging_masking_native_extension" / "Cargo.toml"
 
 
+@dataclass(frozen=True)
+class BenchmarkScenario:
+    name: str
+    python_fn: Callable[[Any], Any]
+    public_native_fn: Callable[[Any], Any]
+    direct_native_fn: Callable[[Any], Any]
+    payload: Any
+    iterations: int
+
+
+def _build_scenarios(native_extension: Any) -> list[BenchmarkScenario]:
+    def python_data(payload: Any) -> Any:
+        settings.experimental_rust_request_logging_masking_enabled = False
+        return mask_sensitive_data(payload, 12)
+
+    def public_native_data(payload: Any) -> Any:
+        settings.experimental_rust_request_logging_masking_enabled = True
+        return mask_sensitive_data(payload, 12)
+
+    def direct_native_data(payload: Any) -> Any:
+        return native_extension.mask_sensitive_data(payload, 12)
+
+    def python_headers(payload: Any) -> Any:
+        settings.experimental_rust_request_logging_masking_enabled = False
+        return mask_sensitive_headers(payload)
+
+    def public_native_headers(payload: Any) -> Any:
+        settings.experimental_rust_request_logging_masking_enabled = True
+        return mask_sensitive_headers(payload)
+
+    def direct_native_headers(payload: Any) -> Any:
+        return native_extension.mask_sensitive_headers(payload)
+
+    return [
+        BenchmarkScenario(
+            name="nested_payload_masking",
+            python_fn=python_data,
+            public_native_fn=public_native_data,
+            direct_native_fn=direct_native_data,
+            payload={
+                "events": [
+                    {
+                        "actor": {"userName": f"user-{index}", "sessionToken": f"token-{index}", "sessionCount": index},
+                        "request": {
+                            "clientSecret": f"secret-{index}",
+                            "payload": {"safeField": "value" * 8, "authDevice": f"device-{index}", "auth_count": index},
+                        },
+                    }
+                    for index in range(1024)
+                ]
+            },
+            iterations=120,
+        ),
+        BenchmarkScenario(
+            name="headers_masking",
+            python_fn=python_headers,
+            public_native_fn=public_native_headers,
+            direct_native_fn=direct_native_headers,
+            payload={
+                **{f"X-Custom-{index}": f"value-{index}" for index in range(512)},
+                **{f"X-Api-Key-{index}": f"secret-{index}" for index in range(256)},
+                "Cookie": "; ".join([f"jwt_token_{index}=abc{index}" for index in range(128)]),
+            },
+            iterations=300,
+        ),
+    ]
+
+
+def test_build_scenarios_exposes_public_and_direct_native_paths(monkeypatch):
+    native_extension = type(
+        "NativeExtension",
+        (),
+        {
+            "mask_sensitive_data": staticmethod(lambda payload, max_depth: {"path": "direct", "payload": payload, "max_depth": max_depth}),
+            "mask_sensitive_headers": staticmethod(lambda payload: {"path": "direct", "payload": payload}),
+        },
+    )()
+    observed_flags: list[bool] = []
+
+    def fake_mask_sensitive_data(payload: Any, max_depth: int = 10) -> Any:
+        observed_flags.append(settings.experimental_rust_request_logging_masking_enabled)
+        return {"path": "public" if settings.experimental_rust_request_logging_masking_enabled else "python", "max_depth": max_depth}
+
+    def fake_mask_sensitive_headers(payload: Any) -> Any:
+        observed_flags.append(settings.experimental_rust_request_logging_masking_enabled)
+        return {"path": "public" if settings.experimental_rust_request_logging_masking_enabled else "python"}
+
+    monkeypatch.setattr("tests.performance.test_request_logging_masking_native_extension_benchmark.mask_sensitive_data", fake_mask_sensitive_data)
+    monkeypatch.setattr("tests.performance.test_request_logging_masking_native_extension_benchmark.mask_sensitive_headers", fake_mask_sensitive_headers)
+
+    scenarios = {scenario.name: scenario for scenario in _build_scenarios(native_extension)}
+
+    public_data = scenarios["nested_payload_masking"].public_native_fn({"k": "v"})
+    direct_data = scenarios["nested_payload_masking"].direct_native_fn({"k": "v"})
+    python_data = scenarios["nested_payload_masking"].python_fn({"k": "v"})
+    public_headers = scenarios["headers_masking"].public_native_fn({"Authorization": "Bearer x"})
+    direct_headers = scenarios["headers_masking"].direct_native_fn({"Authorization": "Bearer x"})
+
+    assert public_data["path"] == "public"
+    assert direct_data["path"] == "direct"
+    assert python_data["path"] == "python"
+    assert public_headers["path"] == "public"
+    assert direct_headers["path"] == "direct"
+    assert observed_flags == [True, False, True]
+
+
 def _ensure_native_extension_installed() -> Any:
     subprocess.run(["uv", "run", "maturin", "develop", "--release", "--manifest-path", str(NATIVE_EXTENSION_MANIFEST)], check=True, cwd=REPO_ROOT)
     return importlib.import_module("request_logging_masking_native_extension")
@@ -45,74 +153,44 @@ def _assert_parity(python_fn: Callable[[Any], Any], rust_fn: Callable[[Any], Any
             raise AssertionError(f"Parity mismatch for payload {payload!r}: python={python_result!r} rust={rust_result!r}")
 
 
-def main() -> None:
-    native_extension = _ensure_native_extension_installed()
-    settings.experimental_rust_request_logging_masking_enabled = False
+def _prepare_public_native_path(native_extension: Any) -> None:
+    settings.experimental_rust_request_logging_masking_enabled = True
+    request_logging_middleware._RUST_REQUEST_LOGGING_MODULE = native_extension
+    request_logging_middleware._RUST_REQUEST_LOGGING_IMPORT_FAILED = False
 
-    def python_data(payload: Any) -> Any:
-        return mask_sensitive_data(payload, 12)
 
-    def rust_data(payload: Any) -> Any:
-        return native_extension.mask_sensitive_data(payload, 12)
+def main() -> None:
+    native_extension = _ensure_native_extension_installed()
+    _prepare_public_native_path(native_extension)
 
-    python_headers = mask_sensitive_headers
-    rust_headers = native_extension.mask_sensitive_headers
+    scenarios = _build_scenarios(native_extension)
 
     _assert_parity(
-        python_data,
-        rust_data,
+        scenarios[0].python_fn,
+        scenarios[0].direct_native_fn,
         [
             {"password": "secret", "nested": {"authToken": "abc", "ok": "value"}},
             {"token_count": 3, "tokenizer": "ok", "privateKey": "secret"},
             [{"jwt_token": "abc"}, {"normal": "value"}],
         ],
     )
     _assert_parity(
-        python_headers,
-        rust_headers,
+        scenarios[1].python_fn,
+        scenarios[1].direct_native_fn,
         [
             {"Authorization": "Bearer abc", "Cookie": "jwt_token=abc; theme=dark", "X-Trace-Id": "123"},
             {"X-Auth-Count": "5", "X-Api-Key": "secret"},
         ],
     )
 
-    scenarios = [
-        (
-            "nested_payload_masking",
-            python_data,
-            rust_data,
-            {
-                "events": [
-                    {
-                        "actor": {"userName": f"user-{index}", "sessionToken": f"token-{index}", "sessionCount": index},
-                        "request": {
-                            "clientSecret": f"secret-{index}",
-                            "payload": {"safeField": "value" * 8, "authDevice": f"device-{index}", "auth_count": index},
-                        },
-                    }
-                    for index in range(1024)
-                ]
-            },
-            120,
-        ),
-        (
-            "headers_masking",
-            python_headers,
-            rust_headers,
-            {
-                **{f"X-Custom-{index}": f"value-{index}" for index in range(512)},
-                **{f"X-Api-Key-{index}": f"secret-{index}" for index in range(256)},
-                "Cookie": "; ".join([f"jwt_token_{index}=abc{index}" for index in range(128)]),
-            },
-            300,
-        ),
-    ]
-
-    for name, python_fn, rust_fn, payload, iterations in scenarios:
-        print(f"\n{name} ({iterations} iterations)")
-        python_median, _ = _measure("python", python_fn, payload, iterations)
-        rust_median, _ = _measure("rust", rust_fn, payload, iterations)
-        print(f"speedup={python_median / rust_median:.2f}x")
+    for scenario in scenarios:
+        print(f"\n{scenario.name} ({scenario.iterations} iterations)")
+        python_median, _ = _measure("python", scenario.python_fn, scenario.payload, scenario.iterations)
+        public_native_median, _ = _measure("public_native", scenario.public_native_fn, scenario.payload, scenario.iterations)
+        direct_native_median, _ = _measure("direct_native", scenario.direct_native_fn, scenario.payload, scenario.iterations)
+        print(f"public_speedup={python_median / public_native_median:.2f}x")
+        print(f"direct_speedup={python_median / direct_native_median:.2f}x")
+        print(f"public_overhead={(public_native_median / direct_native_median - 1) * 100:.1f}%")
 
 
 if __name__ == "__main__":

tests/unit/mcpgateway/middleware/test_request_logging_middleware.py

@@ -264,6 +264,25 @@ async def test_dispatch_logs_json_body(dummy_logger, mock_structured_logger, dum
     assert "******" in dummy_logger.logged[0][1]
 
 
+@pytest.mark.asyncio
+async def test_dispatch_uses_native_json_bytes_fast_path_when_available(dummy_logger, mock_structured_logger, dummy_call_next, monkeypatch):
+    native_extension = MagicMock()
+    native_extension.mask_sensitive_json_bytes.return_value = b'{"password":"******","data":"ok"}'  # pragma: allowlist secret
+    monkeypatch.setattr("mcpgateway.middleware.request_logging_middleware.settings.experimental_rust_request_logging_masking_enabled", True, raising=False)
+
+    middleware = RequestLoggingMiddleware(app=None, enable_gateway_logging=False, log_detailed_requests=True)
+    body = orjson.dumps({"password": "123", "data": "ok"})
+    request = make_request(body=body, headers={"Authorization": "Bearer abc"})
+
+    with patch("mcpgateway.middleware.request_logging_middleware._load_rust_request_logging_module", return_value=native_extension):
+        response = await middleware.dispatch(request, dummy_call_next)
+
+    assert response.status_code == 200
+    native_extension.mask_sensitive_json_bytes.assert_called_once_with(body, 10)
+    native_extension.mask_sensitive_data.assert_not_called()
+    assert any('"password":"******"' in msg for _, msg in dummy_logger.logged)  # pragma: allowlist secret
+
+
 @pytest.mark.asyncio
 async def test_dispatch_falls_back_to_python_masking_when_native_extension_import_fails(dummy_logger, mock_structured_logger, dummy_call_next, monkeypatch):
     monkeypatch.setattr("mcpgateway.middleware.request_logging_middleware.settings.experimental_rust_request_logging_masking_enabled", True, raising=False)