app/observability/bootstrap.py

"""TracerProvider / MeterProvider / LoggerProvider wiring and lifecycle.

This module owns the OpenTelemetry SDK objects and the process-global
observability state. It is deliberately the *only* place that constructs
providers, processors, readers, and OTLP/HTTP exporters, so the operation
boundary (:mod:`app.observability.operation`) can stay a thin consumer of
whatever this module has stood up.

Three modes (from :mod:`app.observability.config`):

* ``disabled`` -- no providers, no store; the boundary degrades to a plain
  timing no-op.
* ``local`` -- app-owned in-process span/log processors plus an in-memory
  metric reader and a :class:`LocalObservationStore`. No network traffic.
* ``full`` -- everything ``local`` records semantically, plus batch OTLP/HTTP
  trace, metric, and log exporters pointed at SigNoz, with bounded queues and
  short export timeouts.

Providers are held on module state and handed directly to the boundary and to
the FastAPI/logging instrumentors; the global OTel providers are intentionally
*not* overwritten, which keeps chromadb's own telemetry and the test suite from
fighting over a singleton.
"""

from __future__ import annotations

import logging
import threading
from dataclasses import dataclass, field
from typing import Any

from opentelemetry.instrumentation.logging.handler import LoggingHandler
from opentelemetry.metrics import Histogram, Meter
from opentelemetry.sdk._logs import LoggerProvider
from opentelemetry.sdk._logs.export import (
    BatchLogRecordProcessor,
    InMemoryLogRecordExporter,
    SimpleLogRecordProcessor,
)
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import (
    InMemoryMetricReader,
    PeriodicExportingMetricReader,
)
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor, SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import Tracer

from app.observability.config import ObservabilityConfig, get_observability_config
from app.observability.contracts import TelemetryMode
from app.observability.local_store import LocalObservationStore
from app.observability.logging import LOGGER_NAME as OBSERVABILITY_LOGGER_NAME

logger = logging.getLogger(__name__)

SERVICE_NAME = "research-buddy-backend"

# One duration histogram carries every operation's latency, keyed by the
# metric-safe "operation" dimension. SigNoz derives the named canonical
# measures (vector_search_ms, retrieval_ms ...) by filtering that dimension,
# which keeps metric cardinality bounded to the allow-listed dimensions.
DURATION_HISTOGRAM_NAME = "operation_duration_ms"

# Bounded batch/export settings for full mode. Short timeouts guarantee a
# telemetry backend outage cannot stall the app on flush/shutdown.
_BATCH_MAX_QUEUE = 2048
_BATCH_MAX_EXPORT = 512
_BATCH_SCHEDULE_MS = 5000
_EXPORT_TIMEOUT_MS = 3000
_METRIC_INTERVAL_MS = 10000


@dataclass
class _State:
    config: ObservabilityConfig
    mode: TelemetryMode
    tracer: Tracer | None = None
    meter: Meter | None = None
    duration_histogram: Histogram | None = None
    stage_duration_histogram: Histogram | None = None
    rag_measurement_histogram: Histogram | None = None
    local_store: LocalObservationStore | None = None
    tracer_provider: TracerProvider | None = None
    meter_provider: MeterProvider | None = None
    logger_provider: LoggerProvider | None = None
    span_exporter: InMemorySpanExporter | None = None
    log_exporter: InMemoryLogRecordExporter | None = None
    metric_reader: Any = None
    log_handler: LoggingHandler | None = None
    instrumented_app: Any = None
    suppressed_failures: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock)

    def note_suppressed_failure(self) -> None:
        with self._lock:
            self.suppressed_failures += 1


def _disabled_state() -> _State:
    return _State(config=get_observability_config(), mode="disabled")


_state: _State = _disabled_state()


def get_state() -> _State:
    """Return the current observability state (a disabled default until init)."""
    return _state


# --- Failure accounting -------------------------------------------------


def note_suppressed_failure() -> None:
    _state.note_suppressed_failure()


def get_suppressed_failure_count() -> int:
    return _state.suppressed_failures


def reset_suppressed_failures() -> None:
    _state.suppressed_failures = 0


# --- Test / Control-Room inspection helpers -------------------------------------------------


def get_captured_spans() -> list[Any]:
    """Local mode only: finished spans held by the in-memory exporter."""
    if _state.span_exporter is None:
        return []
    return list(_state.span_exporter.get_finished_spans())


def get_captured_logs() -> list[Any]:
    """Local mode only: log records that actually reached the OTel log pipeline.

    Only records routed through the scoped ``app.observability`` handler land
    here, which lets tests prove ordinary ``app.*`` logs are NOT exported.
    """
    if _state.log_exporter is None:
        return []
    return list(_state.log_exporter.get_finished_logs())


def collect_histogram_points(name: str = DURATION_HISTOGRAM_NAME) -> list[Any]:
    """Local mode only: histogram data points collected from the in-memory reader."""
    reader = _state.metric_reader
    if not isinstance(reader, InMemoryMetricReader):
        return []
    data = reader.get_metrics_data()
    if data is None:
        return []
    points: list[Any] = []
    for resource_metrics in data.resource_metrics:
        for scope_metrics in resource_metrics.scope_metrics:
            for metric in scope_metrics.metrics:
                if metric.name == name:
                    points.extend(metric.data.data_points)
    return points


# --- Provider construction -------------------------------------------------


def _build_resource() -> Resource:
    return Resource.create({"service.name": SERVICE_NAME})


def _build_tracer_provider(resource: Resource, mode: TelemetryMode) -> tuple[TracerProvider, InMemorySpanExporter | None]:
    provider = TracerProvider(resource=resource)
    span_exporter: InMemorySpanExporter | None = None
    if mode == "local":
        # App-owned, in-process: SimpleSpanProcessor flushes on span end so the
        # captured spans are immediately inspectable without network export.
        span_exporter = InMemorySpanExporter()
        provider.add_span_processor(SimpleSpanProcessor(span_exporter))
    elif mode == "full":
        from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter

        exporter = OTLPSpanExporter(
            endpoint=f"{_state_endpoint()}/v1/traces",
            timeout=_EXPORT_TIMEOUT_MS / 1000.0,
        )
        provider.add_span_processor(
            BatchSpanProcessor(
                exporter,
                max_queue_size=_BATCH_MAX_QUEUE,
                max_export_batch_size=_BATCH_MAX_EXPORT,
                schedule_delay_millis=_BATCH_SCHEDULE_MS,
                export_timeout_millis=_EXPORT_TIMEOUT_MS,
            )
        )
    return provider, span_exporter


def _build_meter_provider(resource: Resource, mode: TelemetryMode) -> tuple[MeterProvider, Any]:
    if mode == "full":
        from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter

        reader: Any = PeriodicExportingMetricReader(
            OTLPMetricExporter(
                endpoint=f"{_state_endpoint()}/v1/metrics",
                timeout=_EXPORT_TIMEOUT_MS / 1000.0,
            ),
            export_interval_millis=_METRIC_INTERVAL_MS,
            export_timeout_millis=_EXPORT_TIMEOUT_MS,
        )
    else:
        reader = InMemoryMetricReader()
    provider = MeterProvider(resource=resource, metric_readers=[reader])
    return provider, reader


def _build_logger_provider(
    resource: Resource, mode: TelemetryMode
) -> tuple[LoggerProvider, InMemoryLogRecordExporter | None]:
    provider = LoggerProvider(resource=resource)
    log_exporter: InMemoryLogRecordExporter | None = None
    if mode == "local":
        log_exporter = InMemoryLogRecordExporter()
        provider.add_log_record_processor(SimpleLogRecordProcessor(log_exporter))
    elif mode == "full":
        from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter

        provider.add_log_record_processor(
            BatchLogRecordProcessor(
                OTLPLogExporter(
                    endpoint=f"{_state_endpoint()}/v1/logs",
                    timeout=_EXPORT_TIMEOUT_MS / 1000.0,
                ),
                max_queue_size=_BATCH_MAX_QUEUE,
                max_export_batch_size=_BATCH_MAX_EXPORT,
                schedule_delay_millis=_BATCH_SCHEDULE_MS,
                export_timeout_millis=_EXPORT_TIMEOUT_MS,
            )
        )
    return provider, log_exporter


def _state_endpoint() -> str:
    return _state.config.otlp_endpoint.rstrip("/")


# --- Lifecycle -------------------------------------------------


def initialize_observability(app: Any = None, config: ObservabilityConfig | None = None) -> None:
    """Stand up (or tear down and replace) the process observability state.

    Called once from FastAPI's ``lifespan``. Safe to call repeatedly; a prior
    state is shut down first. Any failure to build the SDK objects is logged
    and degrades the process to a disabled state -- never raises into startup.
    """
    global _state
    shutdown_observability()

    cfg = config or get_observability_config()

    if cfg.mode == "disabled":
        _state = _State(config=cfg, mode="disabled")
        logger.info("observability disabled")
        return

    try:
        _state = _State(config=cfg, mode=cfg.mode)
        _state.config = cfg
        resource = _build_resource()

        tracer_provider, span_exporter = _build_tracer_provider(resource, cfg.mode)
        meter_provider, metric_reader = _build_meter_provider(resource, cfg.mode)
        logger_provider, log_exporter = _build_logger_provider(resource, cfg.mode)

        meter = meter_provider.get_meter("app.observability")
        histogram = meter.create_histogram(
            DURATION_HISTOGRAM_NAME,
            unit="ms",
            description="Duration of one instrumented backend operation.",
        )
        stage_histogram = meter.create_histogram(
            "rag_stage_duration_ms",
            unit="ms",
            description="Duration of one bounded internal RAG stage.",
        )
        measurement_histogram = meter.create_histogram(
            "rag_measurement",
            unit="{item}",
            description="Per-operation bounded RAG counts, sizes, and indicators.",
        )

        _state.tracer_provider = tracer_provider
        _state.tracer = tracer_provider.get_tracer("app.observability")
        _state.meter_provider = meter_provider
        _state.meter = meter
        _state.duration_histogram = histogram
        _state.stage_duration_histogram = stage_histogram
        _state.rag_measurement_histogram = measurement_histogram
        _state.logger_provider = logger_provider
        _state.span_exporter = span_exporter
        _state.log_exporter = log_exporter
        _state.metric_reader = metric_reader
        _state.local_store = LocalObservationStore(cfg.artifact_root)

        _install_log_correlation(logger_provider)
        _instrument_fastapi(app, tracer_provider)

        logger.warning("observability initialized mode=%s endpoint=%s", cfg.mode, cfg.otlp_endpoint)
    except Exception:
        logger.exception("observability initialization failed; degrading to disabled")
        _state = _State(config=cfg, mode="disabled")


def _install_log_correlation(logger_provider: LoggerProvider) -> None:
    """Attach an OTel logging handler + trace-id injection, best-effort."""
    try:
        from opentelemetry.instrumentation.logging import LoggingInstrumentor

        if not LoggingInstrumentor().is_instrumented_by_opentelemetry:
            LoggingInstrumentor().instrument(set_logging_format=False)
    except Exception:
        logger.debug("logging instrumentor unavailable", exc_info=True)

    try:
        # Scope the OTel log handler to the observability module's own logger
        # only -- NOT the whole "app" tree. Only emit_correlated_log's records
        # (which are path-stripped/sanitized) may reach the OTLP log pipeline.
        # Attaching to "app" would ship every app.* log body (Cognee/main.py
        # error strings routinely embed filesystem paths) unsanitized to
        # /v1/logs, violating the never-export-paths constraint.
        handler = LoggingHandler(level=logging.INFO, logger_provider=logger_provider)
        logging.getLogger(OBSERVABILITY_LOGGER_NAME).addHandler(handler)
        _state.log_handler = handler
    except Exception:
        logger.debug("could not attach OTel logging handler", exc_info=True)


def _instrument_fastapi(app: Any, tracer_provider: TracerProvider) -> None:
    if app is None:
        return
    try:
        from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

        FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
        _state.instrumented_app = app
    except Exception:
        logger.debug("could not instrument FastAPI app", exc_info=True)


def shutdown_observability(timeout_seconds: float = 5.0) -> None:
    """Flush and shut down providers within a hard wall-clock deadline.

    Every step is individually guarded so a hung/failed exporter can never
    block or crash process teardown. Resets state to a disabled default.
    """
    global _state
    state = _state
    if state.mode == "disabled" and state.tracer_provider is None:
        _state = _State(config=state.config, mode="disabled")
        return

    deadline_ms = max(1, int(timeout_seconds * 1000))

    if state.instrumented_app is not None:
        try:
            from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

            FastAPIInstrumentor.uninstrument_app(state.instrumented_app)
        except Exception:
            logger.debug("could not uninstrument FastAPI app", exc_info=True)

    if state.log_handler is not None:
        try:
            logging.getLogger(OBSERVABILITY_LOGGER_NAME).removeHandler(state.log_handler)
        except Exception:
            logger.debug("could not remove OTel log handler", exc_info=True)

    for provider in (state.tracer_provider, state.meter_provider, state.logger_provider):
        if provider is None:
            continue
        _safe_provider_shutdown(provider, deadline_ms)

    _state = _State(config=state.config, mode="disabled")


def _safe_provider_shutdown(provider: Any, deadline_ms: int) -> None:
    force_flush = getattr(provider, "force_flush", None)
    if callable(force_flush):
        try:
            force_flush(deadline_ms)  # type: ignore[misc]
        except TypeError:
            try:
                force_flush()
            except Exception:
                logger.debug("force_flush failed", exc_info=True)
        except Exception:
            logger.debug("force_flush failed", exc_info=True)
    try:
        provider.shutdown()
    except Exception:
        logger.debug("provider shutdown failed", exc_info=True)