Spaces:
Running
Running
File size: 4,644 Bytes
a55ea57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | 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
|