节点在线、应用在线、配置在线使用令牌查询
大石头 authored at 2021-12-16 19:49:30
11.09 KiB
Stardust
"""Web 入口自动埋点

提供 WSGI / ASGI 中间件,并在安装时自动识别已经导入的 Flask、Django、Starlette:

- Flask:包装 ``Flask.wsgi_app``
- Django:包装 ``BaseHandler.get_response``
- Starlette:包装 ``Starlette.__call__``

若框架在安装之后才导入,请显式使用 ``StardustWSGIMiddleware`` 或
``StardustASGIMiddleware``。同一请求只产生一个入口 Span(通过上下文标记去重)。
"""

import sys
import threading
from contextvars import ContextVar
from typing import Any, Callable, Optional

from ..trace_context import find_traceparent
from . import get_tracer, is_suppressed

_ENTRY_ENVIRON_KEY = "stardust.entry_span"
_ENTRY_SCOPE_KEY = "stardust.entry_span"

_lock = threading.RLock()
_entry_active = ContextVar("stardust_entry_active", default=0)
_patches = []


def patch() -> None:
    """安装 Web 入口埋点(仅针对已导入的框架,幂等)"""
    with _lock:
        for installer in (_patch_flask, _patch_django, _patch_starlette):
            try:
                installer()
            except Exception:
                pass


def unpatch() -> None:
    """卸载 Web 入口埋点,恢复框架原始方法"""
    with _lock:
        for owner, attr, original in reversed(_patches):
            try:
                if getattr(getattr(owner, attr, None), "__stardust_instrumented__", False):
                    setattr(owner, attr, original)
            except Exception:
                pass
        del _patches[:]


def _status_code(value: Any) -> int:
    if isinstance(value, int):
        return value
    try:
        return int(str(value).split(" ", 1)[0])
    except Exception:
        return 0


def _wsgi_url(environ: Any) -> str:
    try:
        scheme = environ.get("wsgi.url_scheme", "http") or "http"
        host = environ.get("HTTP_HOST", "") or ""
        if not host:
            host = environ.get("SERVER_NAME", "") or ""
            port = environ.get("SERVER_PORT", "") or ""
            if port and not (
                (scheme == "http" and port == "80")
                or (scheme == "https" and port == "443")
            ):
                host = "%s:%s" % (host, port)
        path = environ.get("PATH_INFO", "") or "/"
        query = environ.get("QUERY_STRING", "") or ""
        return "%s://%s%s%s" % (scheme, host, path, ("?" + query) if query else "")
    except Exception:
        return environ.get("PATH_INFO", "/") or "/"


def _asgi_url(scope: Any) -> str:
    try:
        scheme = scope.get("scheme", "http") or "http"
        path = scope.get("path", "") or "/"
        query = scope.get("query_string", b"") or b""
        if isinstance(query, bytes):
            query = query.decode("latin-1")
        host = ""
        for key, value in scope.get("headers") or []:
            name = key.decode("latin-1") if isinstance(key, bytes) else str(key)
            if name.lower() == "host":
                host = value.decode("latin-1") if isinstance(value, bytes) else str(value)
                break
        if not host:
            server = scope.get("server") or ("", 0)
            if server and server[0]:
                host = server[0]
                if len(server) > 1 and server[1] and not (
                    (scheme == "http" and server[1] == 80)
                    or (scheme == "https" and server[1] == 443)
                ):
                    host = "%s:%s" % (host, server[1])
        return "%s://%s%s%s" % (scheme, host, path, ("?" + query) if query else "")
    except Exception:
        return scope.get("path", "/") or "/"


def _new_server_span(tracer: Any, method: str, path: str, full_url: str, trace_parent: Any) -> Any:
    span = tracer.new_span(
        "%s %s" % (method, path),
        parent_id=trace_parent.span_id if trace_parent is not None else "",
        trace_id=trace_parent.trace_id if trace_parent is not None else None,
    )
    span.tag = "%s %s" % (method, full_url)
    return span


class StardustWSGIMiddleware:
    """WSGI 入口中间件

    Args:
        app: WSGI 应用
        tracer: 追踪器,缺省使用全局安装的追踪器
    """

    def __init__(self, app: Callable, tracer: Optional[Any] = None):
        self.app = app
        self.tracer = tracer

    def __call__(self, environ: Any, start_response: Callable) -> Any:
        tracer = self.tracer or get_tracer()
        if tracer is None or is_suppressed() or _entry_active.get() > 0:
            return self.app(environ, start_response)
        return _invoke_wsgi(
            tracer, lambda env, sr: self.app(env, sr), environ, start_response
        )


def _invoke_wsgi(
    tracer: Any,
    app_call: Callable,
    environ: Any,
    start_response: Callable,
) -> Any:
    if environ.get(_ENTRY_ENVIRON_KEY):
        return app_call(environ, start_response)

    method = environ.get("REQUEST_METHOD", "GET") or "GET"
    path = environ.get("PATH_INFO", "") or "/"
    full_url = _wsgi_url(environ)
    trace_parent = find_traceparent(environ)

    try:
        span = _new_server_span(tracer, method, path, full_url, trace_parent)
    except Exception:
        return app_call(environ, start_response)

    status_holder = {}

    def wrapped_start_response(status, headers, exc_info=None):
        status_holder["status"] = status
        return start_response(status, headers, exc_info)

    previous = environ.get(_ENTRY_ENVIRON_KEY)
    environ[_ENTRY_ENVIRON_KEY] = True
    token = _entry_active.set(_entry_active.get() + 1)
    try:
        with span:
            result = app_call(environ, wrapped_start_response)
            code = _status_code(status_holder.get("status"))
            if code >= 500:
                span.set_error("HTTP %d" % code)
            return result
    finally:
        _entry_active.reset(token)
        if previous is None:
            environ.pop(_ENTRY_ENVIRON_KEY, None)
        else:
            environ[_ENTRY_ENVIRON_KEY] = previous


class StardustASGIMiddleware:
    """ASGI 入口中间件

    Args:
        app: ASGI 应用
        tracer: 追踪器,缺省使用全局安装的追踪器
    """

    def __init__(self, app: Callable, tracer: Optional[Any] = None):
        self.app = app
        self.tracer = tracer

    async def __call__(self, scope: Any, receive: Callable, send: Callable) -> Any:
        tracer = self.tracer or get_tracer()
        if (
            tracer is None
            or is_suppressed()
            or scope.get("type") != "http"
            or _entry_active.get() > 0
        ):
            return await self.app(scope, receive, send)
        return await _invoke_asgi(
            tracer, lambda s: self.app(scope, receive, s), scope, send
        )


async def _invoke_asgi(
    tracer: Any,
    app_call: Callable,
    scope: Any,
    send: Callable,
) -> Any:
    if scope.get(_ENTRY_SCOPE_KEY):
        return await app_call(send)

    method = scope.get("method", "GET") or "GET"
    path = scope.get("path", "") or "/"
    full_url = _asgi_url(scope)
    trace_parent = find_traceparent(scope.get("headers"))

    try:
        span = _new_server_span(tracer, method, path, full_url, trace_parent)
    except Exception:
        return await app_call(send)

    status_holder = {}

    async def wrapped_send(message):
        try:
            if message.get("type") == "http.response.start":
                status_holder["status"] = message.get("status")
        except Exception:
            pass
        await send(message)

    previous = scope.get(_ENTRY_SCOPE_KEY)
    scope[_ENTRY_SCOPE_KEY] = True
    token = _entry_active.set(_entry_active.get() + 1)
    try:
        with span:
            result = await app_call(wrapped_send)
            code = _status_code(status_holder.get("status"))
            if code >= 500:
                span.set_error("HTTP %d" % code)
            return result
    finally:
        _entry_active.reset(token)
        if previous is None:
            scope.pop(_ENTRY_SCOPE_KEY, None)
        else:
            scope[_ENTRY_SCOPE_KEY] = previous


def _patch_flask() -> None:
    flask = sys.modules.get("flask")
    if flask is None:
        return
    app_cls = getattr(flask, "Flask", None)
    if app_cls is None:
        return

    original = app_cls.wsgi_app
    if getattr(original, "__stardust_instrumented__", False):
        return

    def wsgi_app(self, environ, start_response):
        tracer = get_tracer()
        if tracer is None or is_suppressed() or _entry_active.get() > 0:
            return original(self, environ, start_response)
        return _invoke_wsgi(
            tracer, lambda env, sr: original(self, env, sr), environ, start_response
        )

    wsgi_app.__stardust_instrumented__ = True
    wsgi_app.__wrapped__ = original
    app_cls.wsgi_app = wsgi_app
    _patches.append((app_cls, "wsgi_app", original))


def _patch_django() -> None:
    if "django" not in sys.modules:
        return

    try:
        from django.core.handlers.base import BaseHandler
    except Exception:
        return

    original = BaseHandler.get_response
    if getattr(original, "__stardust_instrumented__", False):
        return

    def get_response(self, request):
        tracer = get_tracer()
        if tracer is None or is_suppressed() or _entry_active.get() > 0:
            return original(self, request)

        method = getattr(request, "method", "GET") or "GET"
        path = getattr(request, "path", "") or "/"
        meta = getattr(request, "META", None) or {}
        trace_parent = find_traceparent(meta)
        try:
            full_url = request.build_absolute_uri()
        except Exception:
            full_url = path

        try:
            span = _new_server_span(tracer, method, path, full_url, trace_parent)
        except Exception:
            return original(self, request)

        token = _entry_active.set(_entry_active.get() + 1)
        try:
            with span:
                response = original(self, request)
                code = _status_code(getattr(response, "status_code", 0))
                if code >= 500:
                    span.set_error("HTTP %d" % code)
                return response
        finally:
            _entry_active.reset(token)

    get_response.__stardust_instrumented__ = True
    get_response.__wrapped__ = original
    BaseHandler.get_response = get_response
    _patches.append((BaseHandler, "get_response", original))


def _patch_starlette() -> None:
    module = sys.modules.get("starlette.applications")
    if module is None:
        return
    app_cls = getattr(module, "Starlette", None)
    if app_cls is None:
        return

    original = app_cls.__call__
    if getattr(original, "__stardust_instrumented__", False):
        return

    async def __call__(self, scope, receive, send):
        tracer = get_tracer()
        if (
            tracer is None
            or is_suppressed()
            or scope.get("type") != "http"
            or _entry_active.get() > 0
        ):
            return await original(self, scope, receive, send)
        return await _invoke_asgi(
            tracer, lambda s: original(self, scope, receive, s), scope, send
        )

    __call__.__stardust_instrumented__ = True
    __call__.__wrapped__ = original
    app_cls.__call__ = __call__
    _patches.append((app_cls, "__call__", original))