ImageApi/src/otel.rs

use actix_web::HttpRequest;
use actix_web::http::header::HeaderMap;
use opentelemetry::global::{BoxedSpan, BoxedTracer};
use opentelemetry::propagation::TextMapPropagator;
use opentelemetry::trace::{Span, Status, Tracer};
use opentelemetry::{Context, KeyValue, global};
use opentelemetry_appender_log::OpenTelemetryLogBridge;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::Resource;
use opentelemetry_sdk::logs::{BatchLogProcessor, SdkLoggerProvider};
use opentelemetry_sdk::propagation::TraceContextPropagator;

pub fn global_tracer() -> BoxedTracer {
    global::tracer("image-server")
}

#[allow(dead_code)]
pub fn init_tracing() {
    let resources = Resource::builder()
        .with_attributes([
            KeyValue::new("service.name", "image-server"),
            KeyValue::new("service.version", env!("CARGO_PKG_VERSION")),
        ])
        .build();

    let span_exporter = opentelemetry_otlp::SpanExporter::builder()
        .with_tonic()
        .with_endpoint(std::env::var("OTLP_OTLS_ENDPOINT").unwrap())
        .build()
        .unwrap();

    let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder()
        .with_batch_exporter(span_exporter)
        .with_resource(resources)
        .build();

    global::set_tracer_provider(tracer_provider);
}

#[allow(dead_code)]
pub fn init_logs() {
    let otlp_exporter = opentelemetry_otlp::LogExporter::builder()
        .with_tonic()
        .with_endpoint(std::env::var("OTLP_OTLS_ENDPOINT").unwrap())
        .build()
        .unwrap();

    let exporter = opentelemetry_stdout::LogExporter::default();

    let resources = Resource::builder()
        .with_attributes([
            KeyValue::new("service.name", "image-server"),
            KeyValue::new("service.version", env!("CARGO_PKG_VERSION")),
        ])
        .build();

    let log_provider = SdkLoggerProvider::builder()
        .with_log_processor(BatchLogProcessor::builder(exporter).build())
        .with_log_processor(BatchLogProcessor::builder(otlp_exporter).build())
        .with_resource(resources)
        .build();

    let otel_log_appender = OpenTelemetryLogBridge::new(&log_provider);
    log::set_boxed_logger(Box::new(otel_log_appender)).expect("Unable to set boxed logger");
    //TODO: Still set this with the env? Ideally we still have a clean/simple local logger for local dev
    log::set_max_level(log::LevelFilter::Info);
}

struct HeaderExtractor<'a>(&'a HeaderMap);

impl<'a> opentelemetry::propagation::Extractor for HeaderExtractor<'a> {
    fn get(&self, key: &str) -> Option<&str> {
        self.0.get(key).and_then(|v| v.to_str().ok())
    }

    fn keys(&self) -> Vec<&str> {
        self.0.keys().map(|k| k.as_str()).collect()
    }
}
pub fn extract_context_from_request(req: &HttpRequest) -> Context {
    let propagator = TraceContextPropagator::new();
    propagator.extract(&HeaderExtractor(req.headers()))
}

pub fn trace_db_call<F, O>(
    context: &Context,
    query_type: &str,
    operation: &str,
    func: F,
) -> anyhow::Result<O>
where
    F: FnOnce(&mut BoxedSpan) -> anyhow::Result<O>,
{
    let tracer = global::tracer("db");
    let mut span = tracer
        .span_builder(format!("db.{}.{}", query_type, operation))
        .with_attributes(vec![
            KeyValue::new("db.query_type", query_type.to_string().clone()),
            KeyValue::new("db.operation", operation.to_string().clone()),
        ])
        .start_with_context(&tracer, context);

    let result = func(&mut span);
    match &result {
        Ok(_) => {
            span.set_status(Status::Ok);
        }
        Err(e) => span.set_status(Status::error(e.to_string())),
    }

    result
}