parakeet / consumer / src / main.rs

use deadpool_postgres::Runtime;
use did_resolver::{Resolver, ResolverOpts};
use eyre::OptionExt;
use metrics_exporter_prometheus::PrometheusBuilder;
use std::sync::Arc;
use tokio::signal::ctrl_c;
use tokio_postgres::NoTls;

mod backfill;
mod cmd;
mod config;
mod db;
mod firehose;
mod indexer;
mod instrumentation;
mod label_indexer;
mod utils;

#[tokio::main]
async fn main() -> eyre::Result<()> {
    PrometheusBuilder::new().install()?;

    let cli = cmd::parse();
    let conf = config::load_config()?;

    instrumentation::init_instruments(&conf.instruments);
    let user_agent = build_ua(&conf.ua_contact);

    let pool = conf.database.create_pool(Some(Runtime::Tokio1), NoTls)?;

    let (redis_conn, redis_fut) = redis::Client::open(conf.redis_uri)?
        .create_multiplexed_tokio_connection()
        .await?;
    tokio::spawn(async {
        redis_fut.await;
        Ok::<_, eyre::Report>(())
    });

    let resolver = Arc::new(Resolver::new(ResolverOpts {
        plc_directory: conf.plc_directory,
        user_agent: Some(user_agent.clone()),
        ..Default::default()
    })?);

    let index_client = parakeet_index::Client::connect(conf.index_uri).await?;

    let tracker = tokio_util::task::TaskTracker::new();
    let (stop_tx, stop) = tokio::sync::watch::channel(false);

    let resume = (cli.labels || cli.indexer)
        .then::<Result<_, eyre::Report>, _>(|| {
            let resume_path = conf.resume_path.ok_or_eyre(
                "Config item resume_path must be specified when using --indexer or --labels",
            )?;

            let db = sled::open(resume_path)?;

            Ok(db)
        })
        .transpose()?;

    if cli.labels {
        let resume = resume.clone().unwrap();

        let label_mgr =
            label_indexer::LabelServiceManager::new(pool.clone(), resume, user_agent.clone())
                .await?;

        if let Some(label_source) = conf.label_source {
            tracker.spawn(label_mgr.run(label_source, stop.clone()));
        }
    }

    if cli.backfill {
        let bf_cfg = conf
            .backfill
            .ok_or_eyre("Config item [backfill] must be specified when using --backfill")?;

        let backfiller = backfill::BackfillManager::new(
            pool.clone(),
            redis_conn.clone(),
            resolver.clone(),
            (!bf_cfg.skip_aggregation).then_some(index_client.clone()),
            bf_cfg,
        )
        .await?;

        tracker.spawn(backfiller.run(stop.clone()));
    }

    if cli.indexer {
        let resume = resume.clone().unwrap();

        let indexer_cfg = conf
            .indexer
            .ok_or_eyre("Config item [indexer] must be specified when using --indexer")?;

        let (idxc_tx, idxc_rx) = tokio::sync::mpsc::channel(128);

        let start_seq = resume
            .get("firehose")?
            .and_then(utils::u64_from_ivec)
            .or(indexer_cfg.start_commit_seq);

        if let Some(start_seq) = start_seq {
            tracing::info!("starting firehose consumer from {start_seq}");
        }

        let relay_firehose = firehose::FirehoseConsumer::new_relay(
            &indexer_cfg.relay_source,
            start_seq,
            &user_agent,
        )
        .await?;

        let indexer_opts = indexer::RelayIndexerOpts {
            history_mode: indexer_cfg.history_mode,
            skip_handle_validation: indexer_cfg.skip_handle_validation,
            request_backfill: indexer_cfg.request_backfill,
        };

        let relay_indexer = indexer::RelayIndexer::new(
            pool.clone(),
            redis_conn.clone(),
            idxc_tx,
            resolver.clone(),
            relay_firehose,
            resume,
            indexer_opts,
        )
        .await?;

        tracker.spawn(relay_indexer.run(indexer_cfg.workers, stop));
        tracker.spawn(index_transport(index_client, idxc_rx));
    }

    tokio::spawn(async move {
        let _ = ctrl_c().await;
        tracing::info!("stopping consumer");
        stop_tx.send(true).unwrap();
    });

    tracker.close();
    tracker.wait().await;

    Ok(())
}

async fn index_transport(
    mut idxc: parakeet_index::Client,
    rx: tokio::sync::mpsc::Receiver<parakeet_index::AggregateDeltaReq>,
) -> eyre::Result<()> {
    use tokio_stream::wrappers::ReceiverStream;

    idxc.submit_aggregate_delta_stream(ReceiverStream::new(rx))
        .await?;

    Ok(())
}

fn build_ua(contact: &Option<String>) -> String {
    let mut ua = format!("Parakeet {}", env!("CARGO_PKG_VERSION"));

    if let Some(contact) = contact {
        ua += &format!(" ({contact})");
    }

    ua
}