Something went wrong. Try again.
lightweight `com.atproto.sync.listReposByCollection`
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460use std::net::SocketAddr;use std::path::PathBuf;use std::time::Duration;
use clap::Parser;use tokio::task::JoinSet;use tokio_util::sync::CancellationToken;use tracing::{info, warn};
use lightrail::error::{Error, Result};use lightrail::identity;use lightrail::storage;use lightrail::sync::discovery_queue::DiscoveryQueue;use lightrail::sync::{self, backfill, firehose, resync};use lightrail::util::TokenExt;
#[derive(Parser, Debug)]#[command(name = "lightrail", about = "listReposByCollection indexing service")]struct Args { /// ATProto relay or PDS host to subscribe to (e.g. bsky.network). #[arg(long, env = "LIGHTRAIL_SUBSCRIBE")] subscribe: jacquard_common::url::Url,
/// Path to the fjall database directory. #[arg(long, env = "LIGHTRAIL_DB_PATH", default_value = "lightrail.db")] db_path: PathBuf,
/// TCP address for the XRPC API server. #[arg(long, env = "LIGHTRAIL_LISTEN", default_value = "0.0.0.0:2511")] listen: SocketAddr,
/// PLC directory URL for did:plc resolution. /// Defaults to https://plc.directory. When used together with /// --slingshot-url, acts as a fallback if the slingshot resolver fails. #[arg(long, env = "LIGHTRAIL_PLC_URL")] plc_url: Option<jacquard_common::url::Url>,
/// Slingshot URL for DID resolution. /// Slingshot covers did:plc (primary) and did:web (https well-known, then /// slingshot mini-doc). Omit the URL to use https://slingshot.microcosm.blue. #[arg( long, env = "LIGHTRAIL_SLINGSHOT_URL", num_args = 0..=1, default_missing_value = "https://slingshot.microcosm.blue" )] slingshot_url: Option<jacquard_common::url::Url>,
/// Max identities kept in in-process identity cache. #[arg(long, env = "LIGHTRAIL_IDENT_CACHE_SIZE", default_value_t = 2_000_000)] ident_cache_size: u64,
/// Global identity resolution rate limit (requests/sec). /// Backfill tasks wait for tokens; firehose/resync proceed freely but /// subtract from the budget so backfill slows when the firehose is busy. #[arg(long, env = "LIGHTRAIL_IDENTITY_RESOLUTION_QPS")] identity_resolution_qps: Option<std::num::NonZeroU32>,
/// Maximum concurrent firehose commit worker tasks. #[arg(long, env = "LIGHTRAIL_MAX_FIREHOSE_WORKERS", default_value_t = 6)] max_firehose_workers: usize,
/// Maximum concurrent resync worker tasks. #[arg(long, env = "LIGHTRAIL_MAX_RESYNC_WORKERS", default_value_t = 16)] max_resync_workers: usize,
/// How often to flush the firehose cursor watermark to storage, in seconds. #[arg(long, env = "LIGHTRAIL_CURSOR_SAVE_INTERVAL", default_value_t = 1)] cursor_save_interval_secs: u64,
/// HTTP timeout for describeRepo + getLatestCommit during resync, in seconds. #[arg( long, env = "LIGHTRAIL_DESCRIBE_REPO_FETCH_TIMEOUT", default_value_t = 30 )] describe_repo_fetch_timeout_secs: u64,
/// HTTP timeout for getRepo (full CAR download) during resync, in seconds. #[arg(long, env = "LIGHTRAIL_GET_REPO_FETCH_TIMEOUT", default_value_t = 300)] get_repo_fetch_timeout_secs: u64,
/// TCP address for the Prometheus metrics HTTP endpoint. /// If not set, metrics are not exported. #[arg(long, env = "LIGHTRAIL_METRICS_LISTEN", num_args = 0..=1, default_missing_value = "0.0.0.0:6789")] metrics_listen: Option<SocketAddr>,
/// Admin password for privileged API endpoints. #[arg(long, env = "LIGHTRAIL_ADMIN_PASSWORD")] admin_password: Option<String>,
/// Enable deep crawl: discover PDS hosts via listHosts and crawl each one's repos. #[arg(long, action, env = "LIGHTRAIL_DEEP_CRAWL")] deep_crawl: bool,
/// Heavy mode: always fetch the full repo CAR via getRepo for resync, /// skipping the cheaper describeRepo fast path. #[arg(long, action, env = "LIGHTRAIL_HEAVY")] heavy: bool,
/// Per-PDS HTTP rate limit for crawl/resync requests, in requests per second. #[arg(long, env = "LIGHTRAIL_CRAWL_QPS", default_value_t = std::num::NonZeroU32::new(10).unwrap())] crawl_qps: std::num::NonZeroU32,
/// fjall block cache size in MiB. #[arg(long, env = "LIGHTRAIL_FJALL_CACHE_MB", default_value_t = 256)] fjall_cache_mb: u64,
/// Number of fjall background worker threads (flush + compaction). /// Defaults to fjall's own heuristic (min(CPU cores, 4)). #[arg(long, env = "LIGHTRAIL_FJALL_WORKER_THREADS")] fjall_worker_threads: Option<usize>,
/// Max concurrent per-PDS listRepos workers during deep crawl. #[arg( long, env = "LIGHTRAIL_MAX_DEEP_CRAWL_WORKERS", requires("deep_crawl"), default_value_t = 4 )] max_deep_crawl_workers: usize,}
fn main() { rustls::crypto::aws_lc_rs::default_provider() .install_default() .expect("failed to install rustls crypto provider");
tracing_subscriber::fmt() .with_env_filter(tracing_subscriber::EnvFilter::from_default_env()) .init();
let rt = tokio::runtime::Builder::new_multi_thread() .enable_all() .build() .expect("failed to build tokio runtime");
let result = rt.block_on(run());
// Force-shutdown the runtime after a bounded wait. Without this, a // `spawn_blocking` task genuinely stuck in fjall (e.g. after Poisoned) // holds a blocking-pool thread, and since those threads are non-daemon // they'd prevent process exit. `shutdown_timeout` detaches any remaining // tasks after the deadline; the explicit `process::exit` below then // guarantees we don't wait for detached blocking threads either. rt.shutdown_timeout(Duration::from_secs(10));
match result { Ok(()) => std::process::exit(0), Err(e) => { eprintln!("fatal: {e}"); std::process::exit(1); } }}
async fn run() -> Result<()> { let args = Args::parse();
let subscribe_host = args .subscribe .host() .map(|h| h.to_owned()) .ok_or(crate::Error::Other( "could not get host from --upstream".to_string(), ))?;
let token = CancellationToken::new();
let slingshot_url = args.slingshot_url; let plc_url = args.plc_url; let ident_cache_size = args.ident_cache_size; let resolver = std::sync::Arc::new(identity::build_resolver( slingshot_url, plc_url, ident_cache_size, token.clone(), args.identity_resolution_qps, ));
if let Some(addr) = args.metrics_listen { install_metrics(addr)?; }
let db = storage::open( &args.db_path, args.fjall_cache_mb, args.fjall_worker_threads, )?; let client = lightrail::http::build_client(args.crawl_qps);
let dispatcher_state: resync::dispatcher::DispatcherState = std::sync::Arc::new( std::sync::Mutex::new(resync::dispatcher::DispatcherSnapshot::default()), );
let discovery_queue = std::sync::Arc::new(DiscoveryQueue::new(8192, 32));
let mut tasks: JoinSet<Result<()>> = JoinSet::new();
tasks.spawn({ let token = token.clone(); let db = db.clone(); let host = subscribe_host.clone(); let resolver = resolver.clone(); let client = client.clone(); async move { let mut sub = firehose::Subscriber::new( host, db, resolver, args.max_firehose_workers, Duration::from_secs(args.cursor_save_interval_secs), client, ); sub.run(token) .await .inspect(|_| info!("firehose subscriber done.")) .inspect_err(|e| warn!(error = %e, "firehose exited")) } });
tasks.spawn({ let token = token.clone(); let db = db.clone(); let client = client.clone(); let host = subscribe_host.clone(); let resolver = resolver.clone(); let discovery_queue = discovery_queue.clone(); async move { match backfill::run( host, db, client, token.clone(), resolver, backfill::BackfillMode::Relay, discovery_queue, ) .await { Ok(true) => { info!("backfill complete; idling task"); token.cancelled().await; Ok(()) } Ok(false) => { warn!("backfill ended without finishing, exiting"); Ok(()) } Err(e) => { warn!(error = %e, "backfill errored, exiting"); Err(e) } } } });
tasks.spawn({ let token = token.clone(); let db = db.clone(); let client = client.clone(); let resolver = resolver.clone(); let dispatcher_state = dispatcher_state.clone(); let discovery_queue = discovery_queue.clone(); async move { resync::dispatcher::run(resync::DispatcherConfig { resolver, db, client, max_concurrent: args.max_resync_workers, describe_timeout: Duration::from_secs(args.describe_repo_fetch_timeout_secs), get_repo_timeout: Duration::from_secs(args.get_repo_fetch_timeout_secs), token, force_get_repo: args.heavy, state: dispatcher_state, discovery_queue, }) .await .inspect(|_| info!("resync done.")) .inspect_err(|e| warn!(error = %e, "resync exited")) } });
tasks.spawn({ let token = token.clone(); let db = db.clone(); let addr = args.listen; let dispatcher_state = dispatcher_state.clone(); let client = client.clone(); let admin_config = args .admin_password .map(|pw| lightrail::server::AdminConfig { subscribe_host: subscribe_host.clone(), admin_password: pw, }); async move { lightrail::server::serve( addr, db, token, admin_config, Some(dispatcher_state), Some(client), ) .await .inspect(|_| info!("server done.")) .inspect_err(|e| warn!(error = %e, "server exited")) } });
tasks.spawn({ let db = db.clone(); let token = token.clone(); async move { let mut interval = tokio::time::interval(Duration::from_secs(60)); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); while token.run(interval.tick()).await.is_some() { // Flush the journal BufWriter to the OS. We use // manual_journal_persist so individual writes skip this; // doing it here batches many writes into one flush. if let Err(e) = db.persist_journal() { warn!(error = %e, "failed to persist journal buffer"); } if let Err(e) = storage::meta::save(&db) { warn!(error = %e, "failed to periodically save meta stats"); } // Emit fjall storage gauges. let ss = db.storage_stats(); metrics::gauge!("lightrail_db_disk_bytes", "keyspace" => "total") .set(ss.disk_bytes as f64); metrics::gauge!("lightrail_db_disk_bytes", "keyspace" => "default") .set(ss.default_ks_disk_bytes as f64); metrics::gauge!("lightrail_db_disk_bytes", "keyspace" => "index") .set(ss.index_ks_disk_bytes as f64); metrics::gauge!("lightrail_db_journal_count").set(ss.journal_count as f64); metrics::gauge!("lightrail_db_active_compactions") .set(ss.active_compactions as f64); metrics::gauge!("lightrail_db_compactions_completed") .set(ss.compactions_completed as f64); metrics::gauge!("lightrail_db_time_compacting_seconds") .set(ss.time_compacting.as_secs_f64()); } info!("meta stats done."); Ok(()) } });
if args.deep_crawl { tasks.spawn({ let token = token.clone(); let db = db.clone(); let client = client.clone(); let host = subscribe_host.clone(); let resolver = resolver.clone(); let discovery_queue = discovery_queue.clone(); async move { sync::deep_crawl::run( host, db, client, args.max_deep_crawl_workers, token, resolver, discovery_queue, ) .await .inspect(|_| info!("deep crawl done.")) .inspect_err(|e| warn!(error = %e, "deep crawl exited")) } }); }
// Wait for a shutdown trigger: ctrl-c or any task exiting (including via // panic, surfaced as a JoinError). let first = tokio::select! { _ = tokio::signal::ctrl_c() => { eprintln!("Shutting down..."); None } r = tasks.join_next() => { eprintln!("=== a task exited with: {r:?} ==="); r }, }; token.cancel();
// Drain remaining tasks and surface the first error encountered. let mut error = first.and_then(into_error); while let Some(r) = tasks.join_next().await { if error.is_none() { error = into_error(r); } }
if let Err(e) = storage::meta::save(&db) { warn!(error = %e, "failed to save meta stats on shutdown"); }
error.map_or(Ok(()), Err)}
/// Flatten a task join result into an optional error./// Panics (JoinError) are treated as errors.////// If the error indicates an unrecoverable database state/// ([`Error::is_db_fatal`]), this immediately force-exits the process rather/// than returning. Graceful shutdown isn't safe in that state because other/// tasks may be stuck in blocking fjall calls that will never return.fn into_error(r: std::result::Result<Result<()>, tokio::task::JoinError>) -> Option<Error> { let err = match r { Ok(Ok(())) => return None, Ok(Err(e)) => e, Err(e) => Error::TaskPanic(e), }; if err.is_db_fatal() { eprintln!("FATAL: database poisoned, force-exiting: {err}"); std::process::exit(2); } Some(err)}
fn install_metrics(addr: SocketAddr) -> Result<()> { use metrics_exporter_prometheus::Matcher; metrics_exporter_prometheus::PrometheusBuilder::new() .with_http_listener(addr) // commit car payload sizes.. probably needs more small buckets .set_buckets_for_metric( Matcher::Full("lightrail_commit_car_bytes".to_string()), &[ 1_024., 4_096., 16_384., 65_536., 262_144., 1_048_576., 4_194_304., ], ) // ops per commit: vibes .and_then(|b| { b.set_buckets_for_metric( Matcher::Full("lightrail_commit_ops".to_string()), &[1., 2., 5., 10., 25., 50., 100., 200.], ) }) // retry delay seconds: matches the backoff_secs() ladder // (except for NotFound backoffs) .and_then(|b| { b.set_buckets_for_metric( Matcher::Full("lightrail_resync_retry_delay_seconds".to_string()), &[60., 120., 240., 480., 960., 1_920., 3_600.], ) }) // XRPC request latency: sub-ms cache hits up to multi-second scans .and_then(|b| { b.set_buckets_for_metric( Matcher::Full("lightrail_http_server_request_duration_seconds".to_string()), &[0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1., 5.], ) }) .and_then(|b| b.install()) .map_err(|e| Error::Other(format!("failed to install metrics exporter: {e}")))?; info!(%addr, "metrics exporter listening"); Ok(())}