life-echo/api/app/core/telemetry.py

"""
OpenTelemetry 初始化：traces / metrics / logs 导出至 OTLP Collector。

在 ``setup_logging()`` 之后、FastAPI / Celery 应用创建前调用 ``setup_telemetry(service_name=...)``。
``OTEL_ENABLED=false`` 时无操作，便于测试与无 Collector 环境。
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

from opentelemetry import metrics, trace
from opentelemetry._logs import set_logger_provider
from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter
from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.celery import CeleryInstrumentor
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.instrumentation.logging import LoggingInstrumentor
from opentelemetry.instrumentation.redis import RedisInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.trace.sampling import ParentBasedTraceIdRatio

from app.core.config import settings

if TYPE_CHECKING:
    from fastapi import FastAPI

_initialized = False
_otel_logging_handler: LoggingHandler | None = None


def _build_resource(service_name: str) -> Resource:
    return Resource.create(
        {
            "service.name": service_name,
            "deployment.environment": settings.app_environment,
            "service.version": "0.2.0",
        }
    )


def _build_sampler():
    from opentelemetry.sdk.trace.sampling import (
        ALWAYS_OFF,
        ALWAYS_ON,
        TraceIdRatioBased,
    )

    name = (settings.otel_traces_sampler or "always_on").strip().lower()
    arg = settings.otel_traces_sampler_arg
    if name in ("always_on", "alwayson"):
        return ALWAYS_ON
    if name in ("always_off", "alwaysoff"):
        return ALWAYS_OFF
    ratio = 0.1 if arg is None else arg
    if name == "traceidratio":
        return TraceIdRatioBased(ratio)
    return ParentBasedTraceIdRatio(ratio)


def setup_telemetry(*, service_name: str) -> None:
    """配置 OTLP exporter 与自动 instrumentation（幂等）。"""
    global _initialized, _otel_logging_handler
    if _initialized or not settings.otel_enabled:
        return

    endpoint = settings.otel_exporter_otlp_endpoint.rstrip("/")
    insecure = settings.otel_exporter_otlp_insecure

    resource = _build_resource(service_name)

    span_exporter = OTLPSpanExporter(endpoint=endpoint, insecure=insecure)
    tracer_provider = TracerProvider(resource=resource, sampler=_build_sampler())
    tracer_provider.add_span_processor(BatchSpanProcessor(span_exporter))
    trace.set_tracer_provider(tracer_provider)

    metric_exporter = OTLPMetricExporter(endpoint=endpoint, insecure=insecure)
    metric_reader = PeriodicExportingMetricReader(
        metric_exporter,
        export_interval_millis=settings.otel_metric_export_interval_ms,
    )
    meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
    metrics.set_meter_provider(meter_provider)

    log_exporter = OTLPLogExporter(endpoint=endpoint, insecure=insecure)
    log_provider = LoggerProvider(resource=resource)
    log_provider.add_log_record_processor(BatchLogRecordProcessor(log_exporter))
    set_logger_provider(log_provider)

    LoggingInstrumentor().instrument(set_logging_format=True)
    _otel_logging_handler = LoggingHandler(
        level=logging.NOTSET,
        logger_provider=log_provider,
    )
    logging.getLogger().addHandler(_otel_logging_handler)

    HTTPXClientInstrumentor().instrument()
    RedisInstrumentor().instrument()
    SQLAlchemyInstrumentor().instrument()

    _initialized = True


def instrument_fastapi_app(app: FastAPI) -> None:
    if not settings.otel_enabled:
        return
    FastAPIInstrumentor.instrument_app(
        app,
        excluded_urls="/health",
    )


def instrument_celery() -> None:
    if not settings.otel_enabled:
        return
    CeleryInstrumentor().instrument()


def get_tracer(name: str):
    return trace.get_tracer(name)


def get_meter(name: str):
    return metrics.get_meter(name)


def current_trace_context() -> dict[str, str]:
    """返回当前 span 的 trace_id / span_id（十六进制），无活跃 span 时为空 dict。"""
    span = trace.get_current_span()
    ctx = span.get_span_context()
    if not ctx.is_valid:
        return {}
    return {
        "trace_id": format(ctx.trace_id, "032x"),
        "span_id": format(ctx.span_id, "016x"),
    }