初始提交

2026-07-24 09:43:10 +08:00
commit b5b4721456
793 changed files with 116305 additions and 0 deletions
@@ -0,0 +1,54 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+
+from starlette.types import Message, Scope
+
+from .ctx import TraceCtx
+
+
+class Span:
+    """
+    整个http生命周期：
+        request(before) --> request(after) --> response(before) --> response(after)
+    """
+
+    def __init__(self, scope: Scope) -> None:
+        self.scope = scope
+
+    async def request_before(self) -> None:
+        """
+        request_before: 处理header信息等, 如记录请求体信息
+        """
+        TraceCtx.set_trace_id()
+        TraceCtx.set_request_id()
+        TraceCtx.set_span_id()
+        TraceCtx.set_request_path(self.scope.get('path', ''))
+        TraceCtx.set_request_method(self.scope.get('method', ''))
+
+    async def request_after(self, message: Message) -> Message:
+        """
+        request_after: 处理请求bytes， 如记录请求参数
+
+        example:
+            message: {'type': 'http.request', 'body': b'{\r\n    "name": "\xe8\x8b\x8f\xe8\x8b\x8f\xe8\x8b\x8f"\r\n}', 'more_body': False}
+        """
+        return message
+
+    async def response(self, message: Message) -> Message:
+        """
+        if message['type'] == "http.response.start":   -----> request-before
+            pass
+        if message['type'] == "http.response.body":    -----> request-after
+            message.get('body', b'')
+            pass
+        """
+        if message['type'] == 'http.response.start':
+            message['headers'].append((b'request-id', TraceCtx.get_request_id().encode()))
+            message['headers'].append((b'trace-id', TraceCtx.get_trace_id().encode()))
+            message['headers'].append((b'span-id', TraceCtx.get_span_id().encode()))
+        return message
+
+
+@asynccontextmanager
+async def get_current_span(scope: Scope) -> AsyncGenerator[Span, None]:
+    yield Span(scope)