Something went wrong. Try again.
Something went wrong. Try again.
Rust
at batch-writes
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253use anyhow::Result;use sqlx::{sqlite::SqliteConnectOptions, SqlitePool};use std::collections::HashMap;use std::collections::HashSet;use std::str::FromStr;use std::env;use supercell::cache::Cache;use supercell::cache::CacheTask;use supercell::cleanup::CleanTask;use supercell::vmc::VerificationMethodCacheTask;use tokio::net::TcpListener;use tokio::signal;use tokio_util::{sync::CancellationToken, task::TaskTracker};use tracing_subscriber::prelude::*;
use supercell::consumer::ConsumerTask;use supercell::consumer::ConsumerTaskConfig;use supercell::http::context::WebContext;use supercell::http::server::build_router;use supercell::writer::WriterTask;use supercell::writer::WRITE_CHANNEL_CAPACITY;
#[tokio::main]async fn main() -> Result<()> { tracing_subscriber::registry() .with(tracing_subscriber::EnvFilter::new( std::env::var("RUST_LOG").unwrap_or_else(|_| "supercell=debug,info".into()), )) .with(tracing_subscriber::fmt::layer().pretty()) .init();
let version = supercell::config::version()?;
env::args().for_each(|arg| { if arg == "--version" { println!("{}", version); std::process::exit(0); } });
let config = supercell::config::Config::new()?;
let mut client_builder = reqwest::Client::builder(); for ca_certificate in config.certificate_bundles.as_ref() { tracing::info!("Loading CA certificate: {:?}", ca_certificate); let cert = std::fs::read(ca_certificate)?; let cert = reqwest::Certificate::from_pem(&cert)?; client_builder = client_builder.add_root_certificate(cert); }
client_builder = client_builder.user_agent(config.user_agent.clone()); let http_client = client_builder.build()?;
let connect_options = SqliteConnectOptions::from_str(&config.database_url)? .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal) .create_if_missing(true) .synchronous(sqlx::sqlite::SqliteSynchronous::Normal);
let pool = SqlitePool::connect_with(connect_options).await?; sqlx::migrate!().run(&pool).await?;
let feeds: HashMap<String, (Option<String>, HashSet<String>)> = config .feeds .feeds .iter() .map(|feed| (feed.uri.clone(), (feed.deny.clone(), feed.allow.clone()))) .collect();
let all_dids = feeds .iter() .flat_map(|(_, (_, allow))| allow.iter().cloned()) .collect::<HashSet<String>>();
let cache = Cache::new(20);
let web_context = WebContext::new( pool.clone(), config.external_base.as_str(), feeds, cache.clone(), );
let app = build_router(web_context.clone());
let tracker = TaskTracker::new(); let token = CancellationToken::new();
{ let tracker = tracker.clone(); let inner_token = token.clone();
let ctrl_c = async { signal::ctrl_c() .await .expect("failed to install Ctrl+C handler"); };
let terminate = async { signal::unix::signal(signal::unix::SignalKind::terminate()) .expect("failed to install signal handler") .recv() .await; };
tokio::spawn(async move { tokio::select! { () = inner_token.cancelled() => { }, _ = terminate => {}, _ = ctrl_c => {}, }
tracker.close(); inner_token.cancel(); }); }
{ let inner_config = config.clone(); let task_enable = *inner_config.consumer_task_enable.as_ref(); if task_enable { let (write_tx, write_rx) = tokio::sync::mpsc::channel(WRITE_CHANNEL_CAPACITY);
let consumer_task_config = ConsumerTaskConfig { user_agent: inner_config.user_agent.clone(), compression: *inner_config.compression.as_ref(), zstd_dictionary_location: inner_config.zstd_dictionary.clone(), jetstream_hostname: inner_config.jetstream_hostname.clone(), feeds: inner_config.feeds.clone(), collections: inner_config.collections.as_ref().clone(), }; let task = ConsumerTask::new(pool.clone(), consumer_task_config, token.clone(), write_tx)?; let inner_token = token.clone(); tracker.spawn(async move { if let Err(err) = task.run_background().await { tracing::warn!(error = ?err, "consumer task error"); } inner_token.cancel(); });
let writer = WriterTask::new( pool.clone(), inner_config.jetstream_hostname.clone(), write_rx, ); let inner_token = token.clone(); tracker.spawn(async move { if let Err(err) = writer.run_background().await { tracing::warn!(error = ?err, "writer task error"); } inner_token.cancel(); }); } }
{ let inner_config = config.clone(); let task_enable = *inner_config.vmc_task_enable.as_ref(); if task_enable { let task = VerificationMethodCacheTask::new( pool.clone(), http_client, inner_config.plc_hostname.clone(), all_dids, token.clone(), ); task.main().await?; let inner_token = token.clone(); tracker.spawn(async move { if let Err(err) = task.run_background(chrono::Duration::hours(4)).await { tracing::warn!(error = ?err, "consumer task error"); } inner_token.cancel(); }); } }
{ let inner_config = config.clone(); let task_enable = *inner_config.cache_task_enable.as_ref(); if task_enable { let task = CacheTask::new( pool.clone(), cache.clone(), inner_config.clone(), token.clone(), ); task.main().await?; let inner_token = token.clone(); let interval = *inner_config.cache_task_interval.as_ref(); tracker.spawn(async move { if let Err(err) = task.run_background(interval).await { tracing::warn!(error = ?err, "cache task error"); } inner_token.cancel(); }); } }
{ let inner_config = config.clone(); let task_enable = *inner_config.cleanup_task_enable.as_ref(); let max_age = *inner_config.cleanup_task_max_age.as_ref(); if task_enable { let task = CleanTask::new(pool.clone(), max_age, token.clone()); task.main().await?; let inner_token = token.clone(); let interval = *inner_config.cleanup_task_interval.as_ref(); tracker.spawn(async move { if let Err(err) = task.run_background(interval).await { tracing::warn!(error = ?err, "cleanup task error"); } inner_token.cancel(); }); } }
{ let inner_config = config.clone(); let http_port = *inner_config.http_port.as_ref(); let inner_token = token.clone(); tracker.spawn(async move { let listener = TcpListener::bind(&format!("0.0.0.0:{}", http_port)) .await .unwrap();
let shutdown_token = inner_token.clone(); let result = axum::serve(listener, app) .with_graceful_shutdown(async move { tokio::select! { () = shutdown_token.cancelled() => { } } tracing::info!("axum graceful shutdown complete"); }) .await; if let Err(err) = result { tracing::error!("axum task failed: {}", err); }
inner_token.cancel(); }); }
tracker.wait().await;
tracing::info!("closing database connection pool"); pool.close().await;
tracing::info!("shutdown complete");
Ok(())}