pict-rs-aggregator/src/main.rs

use actix_web::{App, HttpServer};
use awc::Client;
#[cfg(feature = "console")]
use console_subscriber::ConsoleLayer;
use opentelemetry::{
    sdk::{propagation::TraceContextPropagator, Resource},
    KeyValue,
};
use opentelemetry_otlp::WithExportConfig;
use std::time::Duration;
use structopt::StructOpt;
use tracing::subscriber::set_global_default;
use tracing_actix_web::TracingLogger;
use tracing_awc::Tracing;
use tracing_error::ErrorLayer;
use tracing_log::LogTracer;
use tracing_subscriber::{
    filter::Targets, fmt::format::FmtSpan, layer::SubscriberExt, Layer, Registry,
};
use url::Url;

#[actix_web::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let config = pict_rs_aggregator::Config::from_args();

    init_logger(config.opentelemetry_url())?;

    let mut db_path = config.db_path().to_owned();
    db_path.push("sled");
    db_path.push("db-0-34");
    let db = sled::open(db_path)?;
    let bind_address = config.bind_address();
    let state = pict_rs_aggregator::state(config, "", db)?;

    HttpServer::new(move || {
        let client = Client::builder()
            .wrap(Tracing)
            .timeout(Duration::from_secs(30))
            .add_default_header(("User-Agent", "pict_rs_aggregator-v0.1.0"))
            .finish();

        App::new()
            .wrap(TracingLogger::default())
            .configure(|cfg| pict_rs_aggregator::configure(cfg, state.clone(), client))
    })
    .bind(bind_address)?
    .run()
    .await?;

    Ok(())
}

fn init_logger(opentelemetry_url: Option<&Url>) -> Result<(), Box<dyn std::error::Error>> {
    LogTracer::init()?;

    opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());

    let targets: Targets = std::env::var("RUST_LOG")
        .unwrap_or_else(|_| "info".into())
        .parse()?;

    let format_layer = tracing_subscriber::fmt::layer()
        .with_span_events(FmtSpan::NEW | FmtSpan::CLOSE)
        .with_filter(targets.clone());

    #[cfg(feature = "console")]
    let console_layer = ConsoleLayer::builder()
        .with_default_env()
        .server_addr(([0, 0, 0, 0], 6669))
        .event_buffer_capacity(1024 * 1024)
        .spawn();

    let subscriber = Registry::default()
        .with(format_layer)
        .with(ErrorLayer::default());

    #[cfg(feature = "console")]
    let subscriber = subscriber.with(console_layer);

    if let Some(url) = opentelemetry_url {
        let tracer =
            opentelemetry_otlp::new_pipeline()
                .tracing()
                .with_trace_config(opentelemetry::sdk::trace::config().with_resource(
                    Resource::new(vec![KeyValue::new("service.name", "pict-rs-aggregator")]),
                ))
                .with_exporter(
                    opentelemetry_otlp::new_exporter()
                        .tonic()
                        .with_endpoint(url.as_str()),
                )
                .install_batch(opentelemetry::runtime::Tokio)?;

        let otel_layer = tracing_opentelemetry::layer()
            .with_tracer(tracer)
            .with_filter(targets);

        let subscriber = subscriber.with(otel_layer);

        set_global_default(subscriber)?;
    } else {
        set_global_default(subscriber)?;
    }

    Ok(())
}