lavina/src/main.rs

mod http;

use std::future::Future;
use std::path::Path;

use clap::Parser;
use figment::providers::Format;
use figment::{providers::Toml, Figment};
use opentelemetry::global::set_text_map_propagator;
use opentelemetry::KeyValue;
use opentelemetry_otlp::WithExportConfig;
use opentelemetry_sdk::propagation::TraceContextPropagator;
use opentelemetry_sdk::trace::{BatchConfig, RandomIdGenerator, Sampler};
use opentelemetry_sdk::{runtime, Resource};
use opentelemetry_semantic_conventions::resource::SERVICE_NAME;
use opentelemetry_semantic_conventions::SCHEMA_URL;
use prometheus::Registry as MetricsRegistry;
use serde::Deserialize;
use tracing_opentelemetry::OpenTelemetryLayer;
use tracing_subscriber::fmt::Subscriber;
use tracing_subscriber::prelude::*;

use lavina_core::prelude::*;
use lavina_core::repo::Storage;
use lavina_core::LavinaCore;

#[derive(Deserialize, Debug)]
struct ServerConfig {
    telemetry: http::ServerConfig,
    irc: projection_irc::ServerConfig,
    xmpp: projection_xmpp::ServerConfig,
    storage: lavina_core::repo::StorageConfig,
    cluster: lavina_core::clustering::ClusterConfig,
    tracing: Option<TracingConfig>,
}

#[derive(Deserialize, Debug)]
struct TracingConfig {
    endpoint: String,
    service_name: String,
}

#[derive(Parser)]
struct CliArgs {
    #[arg(long)]
    config: Box<Path>,
}

fn load_config() -> Result<ServerConfig> {
    let args = CliArgs::parse();
    let raw_config = Figment::from(Toml::file(args.config));
    let config: ServerConfig = raw_config.extract()?;
    Ok(config)
}

#[tokio::main]
async fn main() -> Result<()> {
    let sleep = ctrl_c()?;
    let config = load_config()?;
    set_up_logging(&config.tracing)?;

    tracing::info!("Booting up");
    tracing::info!("Loaded config: {config:?}");
    let ServerConfig {
        telemetry: telemetry_config,
        irc: irc_config,
        xmpp: xmpp_config,
        storage: storage_config,
        cluster: cluster_config,
        tracing: _,
    } = config;
    let mut metrics = MetricsRegistry::new();
    let storage = Storage::open(storage_config).await?;
    let core = LavinaCore::new(&mut metrics, cluster_config, storage).await?;
    let telemetry_terminator = http::launch(telemetry_config, metrics.clone(), core.clone()).await?;
    let irc = projection_irc::launch(irc_config, core.clone(), metrics.clone()).await?;
    let xmpp = projection_xmpp::launch(xmpp_config, core.clone(), metrics.clone()).await?;
    tracing::info!("Started");

    sleep.await;

    tracing::info!("Begin shutdown");
    xmpp.terminate().await?;
    irc.terminate().await?;
    telemetry_terminator.terminate().await?;
    let storage = core.shutdown().await;
    storage.close().await;
    tracing::info!("Shutdown complete");
    Ok(())
}

#[cfg(windows)]
fn ctrl_c() -> Result<impl Future<Output = ()>> {
    use tokio::signal::windows::*;
    let chan = ctrl_c()?;
    async fn recv(mut chan: CtrlC) {
        let _ = chan.recv().await;
    }
    Ok(recv(chan))
}

#[cfg(unix)]
fn ctrl_c() -> Result<impl Future<Output = ()>> {
    use tokio::signal::unix::*;
    let chan = signal(SignalKind::interrupt())?;
    async fn recv(mut chan: Signal) {
        let _ = chan.recv().await;
    }
    Ok(recv(chan))
}

fn set_up_logging(tracing_config: &Option<TracingConfig>) -> Result<()> {
    let subscriber = tracing_subscriber::registry().with(tracing_subscriber::fmt::layer());

    let targets = {
        use std::{env, str::FromStr};
        use tracing_subscriber::filter::Targets;
        match env::var("RUST_LOG") {
            Ok(var) => Targets::from_str(&var)
                .map_err(|e| {
                    eprintln!("Ignoring `RUST_LOG={:?}`: {}", var, e);
                })
                .unwrap_or_default(),
            Err(env::VarError::NotPresent) => Targets::new().with_default(Subscriber::DEFAULT_MAX_LEVEL),
            Err(e) => {
                eprintln!("Ignoring `RUST_LOG`: {}", e);
                Targets::new().with_default(Subscriber::DEFAULT_MAX_LEVEL)
            }
        }
    };
    if let Some(config) = tracing_config {
        let trace_config = opentelemetry_sdk::trace::Config::default()
            .with_sampler(Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(1.0))))
            .with_id_generator(RandomIdGenerator::default())
            .with_resource(Resource::from_schema_url(
                [KeyValue::new(SERVICE_NAME, config.service_name.to_string())],
                SCHEMA_URL,
            ));
        let trace_exporter = opentelemetry_otlp::new_exporter().tonic().with_endpoint(&config.endpoint);
        let tracer = opentelemetry_otlp::new_pipeline()
            .tracing()
            .with_trace_config(trace_config)
            .with_batch_config(BatchConfig::default())
            .with_exporter(trace_exporter)
            .install_batch(runtime::Tokio)?;
        let subscriber = subscriber.with(OpenTelemetryLayer::new(tracer));
        set_text_map_propagator(TraceContextPropagator::new());
        targets.with_subscriber(subscriber).try_init()?;
    } else {
        targets.with_subscriber(subscriber).try_init()?;
    }
    Ok(())
}