Spaces:
Running
Running
| import os | |
| import time | |
| from uuid import uuid4 | |
| from starlette.types import ASGIApp, Message, Receive, Scope, Send | |
| from src.lib.tracing.events import TraceEventType, TraceSeverity | |
| from src.lib.tracing.tracker import tracer | |
| class TraceMiddleware: | |
| def __init__( | |
| self, | |
| app: ASGIApp, | |
| service_name: str, | |
| slow_request_ms: float | None = None, | |
| ): | |
| self.app = app | |
| self.service_name = service_name | |
| self.slow_request_ms = ( | |
| slow_request_ms | |
| if slow_request_ms is not None | |
| else float(os.getenv("TRACE_SLOW_REQUEST_MS", "1000")) | |
| ) | |
| async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: | |
| if scope["type"] not in ("http", "websocket"): | |
| await self.app(scope, receive, send) | |
| return | |
| started_at = time.perf_counter() | |
| status_code: int | None = None | |
| tenant_id = self._tenant_id(scope) | |
| trace_id = self._header(scope, "x-trace-id") or str(uuid4()) | |
| tokens = tracer.set_context(trace_id=trace_id, tenant_id=tenant_id) | |
| async def send_wrapper(message: Message) -> None: | |
| nonlocal status_code | |
| if message["type"] == "http.response.start": | |
| status_code = message.get("status") | |
| await send(message) | |
| try: | |
| await self.app(scope, receive, send_wrapper) | |
| except Exception as exc: | |
| duration_ms = (time.perf_counter() - started_at) * 1000 | |
| tracer.exception( | |
| name=self._operation(scope), | |
| exception=exc, | |
| service=self.service_name, | |
| operation=self._operation(scope), | |
| duration_ms=duration_ms, | |
| status_code=status_code, | |
| threshold_ms=self.slow_request_ms, | |
| metadata=self._scope_metadata(scope, status_code), | |
| ) | |
| tracer.reset_context(tokens) | |
| raise | |
| duration_ms = (time.perf_counter() - started_at) * 1000 | |
| operation = self._operation(scope) | |
| if status_code is not None and status_code >= 500: | |
| tracer.error( | |
| name=operation, | |
| service=self.service_name, | |
| operation=operation, | |
| message=f"{self.service_name} returned HTTP {status_code}", | |
| severity=TraceSeverity.ERROR, | |
| duration_ms=duration_ms, | |
| status_code=status_code, | |
| threshold_ms=self.slow_request_ms, | |
| metadata=self._scope_metadata(scope, status_code), | |
| ) | |
| if self.slow_request_ms is not None and duration_ms >= self.slow_request_ms: | |
| tracer.publish( | |
| TraceEventType.SLA_BREACH, | |
| name=operation, | |
| severity=TraceSeverity.WARNING, | |
| service=self.service_name, | |
| operation=operation, | |
| duration_ms=duration_ms, | |
| threshold_ms=self.slow_request_ms, | |
| status_code=status_code, | |
| message=f"{operation} exceeded SLA threshold", | |
| metadata=self._scope_metadata(scope, status_code), | |
| ) | |
| tracer.reset_context(tokens) | |
| def _operation(self, scope: Scope) -> str: | |
| method = scope.get("method", scope["type"]) | |
| return f"{method} {scope.get('path', '')}".strip() | |
| def _scope_metadata(self, scope: Scope, status_code: int | None) -> dict[str, str]: | |
| metadata = { | |
| "service": self.service_name, | |
| "protocol": scope["type"], | |
| "path": str(scope.get("path", "")), | |
| } | |
| if status_code is not None: | |
| metadata["status_code"] = str(status_code) | |
| return metadata | |
| def _tenant_id(self, scope: Scope) -> str | None: | |
| return ( | |
| self._header(scope, "x-brain-id") | |
| or self._query_param(scope, "brain_id") | |
| or self._header(scope, "x-tenant-id") | |
| ) | |
| def _header(self, scope: Scope, name: str) -> str | None: | |
| encoded_name = name.lower().encode() | |
| for header_name, value in scope.get("headers", []): | |
| if header_name.lower() == encoded_name: | |
| return value.decode(errors="replace").rstrip() or None | |
| return None | |
| def _query_param(self, scope: Scope, name: str) -> str | None: | |
| raw_query = scope.get("query_string", b"").decode(errors="replace") | |
| for item in raw_query.split("&"): | |
| key, _, value = item.partition("=") | |
| if key == name and value: | |
| return value.rstrip() | |
| return None | |