ljwx-health/ljwx-bigscreen-bak/observability.py

"""
可观测性集成模块
集成 Prometheus metrics 和 OpenTelemetry tracing
"""
import os
from prometheus_client import Counter, Histogram, Gauge, generate_latest, CONTENT_TYPE_LATEST
from prometheus_flask_exporter import PrometheusMetrics
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource, SERVICE_NAME
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.instrumentation.redis import RedisInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor


# ============================================================================
# Prometheus 自定义指标
# ============================================================================

# 健康数据查询计数
health_query_counter = Counter(
    'bigscreen_health_query_total',
    'Total health data query count',
    ['endpoint', 'status']
)

# Celery 任务执行时长
celery_task_duration = Histogram(
    'bigscreen_celery_task_duration_seconds',
    'Celery task execution duration in seconds',
    ['task_name', 'status']
)

# 活跃用户数
active_users_gauge = Gauge(
    'bigscreen_active_users',
    'Number of active users on bigscreen'
)

# WebSocket 连接数
websocket_connections = Gauge(
    'bigscreen_websocket_connections',
    'Number of active WebSocket connections'
)

# 数据库查询时长
db_query_duration = Histogram(
    'bigscreen_db_query_duration_seconds',
    'Database query duration in seconds',
    ['operation', 'table']
)

# Redis 操作时长
redis_operation_duration = Histogram(
    'bigscreen_redis_operation_duration_seconds',
    'Redis operation duration in seconds',
    ['operation', 'command']
)

# API 错误计数
api_error_counter = Counter(
    'bigscreen_api_error_total',
    'Total API error count',
    ['endpoint', 'error_type']
)

# 数据刷新延迟
data_refresh_latency = Histogram(
    'bigscreen_data_refresh_latency_seconds',
    'Data refresh latency in seconds',
    ['data_type']
)


# ============================================================================
# OpenTelemetry 配置
# ============================================================================

def init_tracing(app):
    """
    初始化 OpenTelemetry 分布式链路追踪

    :param app: Flask app 实例
    """
    # 获取 OTLP 配置
    otlp_endpoint = os.getenv('OTEL_EXPORTER_OTLP_ENDPOINT', 'http://localhost:4317')
    service_name = os.getenv('OTEL_SERVICE_NAME', 'ljwx-bigscreen')
    sampling_rate = float(os.getenv('OTEL_TRACES_SAMPLER_ARG', '0.1'))

    # 创建 Resource（服务标识）
    resource = Resource(attributes={
        SERVICE_NAME: service_name,
        'service.namespace': 'ljwx-health',
        'service.instance.id': os.getenv('HOSTNAME', 'localhost'),
        'deployment.environment': os.getenv('ENV', 'production')
    })

    # 配置 TracerProvider
    tracer_provider = TracerProvider(resource=resource)

    # 配置 OTLP Exporter
    otlp_exporter = OTLPSpanExporter(
        endpoint=otlp_endpoint,
        insecure=True  # 在生产环境中应使用 TLS
    )

    # 添加 BatchSpanProcessor
    tracer_provider.add_span_processor(BatchSpanProcessor(otlp_exporter))

    # 设置全局 TracerProvider
    trace.set_tracer_provider(tracer_provider)

    # 自动埋点：Flask
    FlaskInstrumentor().instrument_app(app)

    # 自动埋点：Requests（HTTP 客户端）
    RequestsInstrumentor().instrument()

    app.logger.info(
        f'OpenTelemetry tracing initialized: '
        f'service={service_name}, endpoint={otlp_endpoint}, '
        f'sampling_rate={sampling_rate}'
    )


def init_tracing_for_db(engine):
    """
    初始化数据库查询追踪

    :param engine: SQLAlchemy engine
    """
    SQLAlchemyInstrumentor().instrument(engine=engine)


def init_tracing_for_redis():
    """
    初始化 Redis 操作追踪
    """
    RedisInstrumentor().instrument()


# ============================================================================
# Prometheus Metrics 初始化
# ============================================================================

def init_metrics(app):
    """
    初始化 Prometheus metrics 端点

    :param app: Flask app 实例
    :return: PrometheusMetrics 实例
    """
    # 创建 PrometheusMetrics 实例
    metrics = PrometheusMetrics(app)

    # 注册默认指标（HTTP 请求时长、请求计数等）
    metrics.info('ljwx_bigscreen_info', 'LJWX BigScreen application info', version='1.0.0')

    # 自定义 /metrics 端点
    @app.route('/metrics')
    def metrics_endpoint():
        """Prometheus metrics 端点"""
        return generate_latest(), 200, {'Content-Type': CONTENT_TYPE_LATEST}

    app.logger.info('Prometheus metrics endpoint initialized at /metrics')

    return metrics


# ============================================================================
# 便捷函数
# ============================================================================

def record_health_query(endpoint, status='success'):
    """记录健康数据查询"""
    health_query_counter.labels(endpoint=endpoint, status=status).inc()


def record_celery_task(task_name, duration, status='success'):
    """记录 Celery 任务执行"""
    celery_task_duration.labels(task_name=task_name, status=status).observe(duration)


def set_active_users(count):
    """设置活跃用户数"""
    active_users_gauge.set(count)


def set_websocket_connections(count):
    """设置 WebSocket 连接数"""
    websocket_connections.set(count)


def record_db_query(operation, table, duration):
    """记录数据库查询时长"""
    db_query_duration.labels(operation=operation, table=table).observe(duration)


def record_redis_operation(operation, command, duration):
    """记录 Redis 操作时长"""
    redis_operation_duration.labels(operation=operation, command=command).observe(duration)


def record_api_error(endpoint, error_type):
    """记录 API 错误"""
    api_error_counter.labels(endpoint=endpoint, error_type=error_type).inc()


def record_data_refresh_latency(data_type, latency):
    """记录数据刷新延迟"""
    data_refresh_latency.labels(data_type=data_type).observe(latency)


# ============================================================================
# Context Manager for Timing
# ============================================================================

class TimingContext:
    """计时上下文管理器"""

    def __init__(self, histogram, **labels):
        self.histogram = histogram
        self.labels = labels
        self.start_time = None

    def __enter__(self):
        import time
        self.start_time = time.time()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        import time
        duration = time.time() - self.start_time
        self.histogram.labels(**self.labels).observe(duration)


# 使用示例：
# with TimingContext(db_query_duration, operation='select', table='user_health'):
#     # 执行数据库查询
#     pass