use std::path::PathBuf; use clap::{Parser, Subcommand}; use sea_orm::{ConnectOptions, Database}; use signal_hook::consts::signal::*; use signal_hook_tokio::Signals; use futures::stream::StreamExt; use upub::{context, ext::LoggableError}; #[cfg(feature = "cli")] use upub_cli as cli; #[cfg(feature = "migrate")] use upub_migrations as migrations; #[cfg(feature = "serve")] use upub_routes as routes; #[cfg(feature = "worker")] use upub_worker as worker; #[derive(Parser)] /// all names were taken struct Args { #[clap(subcommand)] /// command to run command: Mode, /// path to config file, leave empty to not use any #[arg(short, long)] config: Option, #[arg(long = "db")] /// database connection uri, overrides config value database: Option, #[arg(long)] /// instance base domain, for AP ids, overrides config value domain: Option, #[arg(long, default_value_t=false)] /// run with debug level tracing debug: bool, #[arg(long)] /// force set number of worker threads for async runtime, defaults to number of cores threads: Option, } #[derive(Clone, Subcommand)] enum Mode { /// print current or default configuration Config, #[cfg(feature = "migrate")] /// apply database migrations Migrate, #[cfg(feature = "cli")] /// run maintenance CLI tasks Cli { #[clap(subcommand)] /// task to run command: cli::CliCommand, }, #[cfg(all(feature = "serve", feature = "worker"))] /// start both api routes and background workers Monolith { #[arg(short, long, default_value="127.0.0.1:3000")] /// addr to bind and serve onto bind: String, #[arg(short, long, default_value_t = 4)] /// how many concurrent jobs to process with this worker tasks: usize, #[arg(short, long, default_value_t = 20)] /// interval for polling new tasks poll: u64, }, #[cfg(feature = "serve")] /// start api routes server Serve { #[arg(short, long, default_value="127.0.0.1:3000")] /// addr to bind and serve onto bind: String, }, #[cfg(feature = "worker")] /// start background job worker Work { /// only run tasks of this type, run all if not given filter: Filter, /// how many concurrent jobs to process with this worker #[arg(short, long, default_value_t = 4)] tasks: usize, #[arg(short, long, default_value_t = 20)] /// interval for polling new tasks poll: u64, }, } fn main() { let args = Args::parse(); tracing_subscriber::fmt() .compact() .with_max_level(if args.debug { tracing::Level::DEBUG } else { tracing::Level::INFO }) .init(); let config = upub::Config::load(args.config.as_ref()); if matches!(args.command, Mode::Config) { println!("{}", toml::to_string_pretty(&config).expect("failed serializing config")); return; } let mut runtime = tokio::runtime::Builder::new_multi_thread(); if let Some(threads) = args.threads { runtime.worker_threads(threads); } runtime .enable_io() .enable_time() .thread_name("upub-worker") .build() .expect("failed creating tokio async runtime") .block_on(async { init(args, config).await }) } async fn init(args: Args, config: upub::Config) { let database = args.database.unwrap_or(config.datasource.connection_string.clone()); let domain = args.domain.unwrap_or(config.instance.domain.clone()); // TODO can i do connectoptions.into() or .connect() and skip these ugly bindings? let mut opts = ConnectOptions::new(&database); opts .sqlx_logging(true) .sqlx_logging_level(tracing::log::LevelFilter::Debug) .max_connections(config.datasource.max_connections) .min_connections(config.datasource.min_connections) .acquire_timeout(std::time::Duration::from_secs(config.datasource.acquire_timeout_seconds)) .connect_timeout(std::time::Duration::from_secs(config.datasource.connect_timeout_seconds)) .sqlx_slow_statements_logging_settings( if config.datasource.slow_query_warn_enable { tracing::log::LevelFilter::Warn } else { tracing::log::LevelFilter::Debug }, std::time::Duration::from_secs(config.datasource.slow_query_warn_seconds) ); let db = Database::connect(opts) .await.expect("error connecting to db"); #[cfg(feature = "migrate")] if matches!(args.command, Mode::Migrate) { use migrations::MigratorTrait; migrations::Migrator::up(&db, None) .await .expect("error applying migrations"); return; } let (tx_wake, rx_wake) = tokio::sync::mpsc::unbounded_channel(); let wake = WakeToken(rx_wake); let ctx = upub::Context::new(db, domain, config.clone(), Some(Box::new(WakerToken(tx_wake)))) .await.expect("failed creating server context"); #[cfg(feature = "cli")] if let Mode::Cli { command } = args.command { cli::run(ctx, command) .await.expect("failed running cli task"); return; } // register signal handler only for long-lasting modes, such as server or worker let (tx, rx) = tokio::sync::watch::channel(false); let signals = Signals::new([SIGTERM, SIGINT]).expect("failed registering signal handler"); let handle = signals.handle(); let signals_task = tokio::spawn(handle_signals(signals, tx)); let stop = CancellationToken(rx); match args.command { #[cfg(feature = "serve")] Mode::Serve { bind } => routes::serve(ctx, bind, stop) .await.expect("failed serving api routes"), #[cfg(feature = "worker")] Mode::Work { filter, tasks, poll } => worker::spawn(ctx, tasks, poll, filter.into(), stop, wake) .await.expect("failed running worker"), #[cfg(all(feature = "serve", feature = "worker"))] Mode::Monolith { bind, tasks, poll } => { worker::spawn(ctx.clone(), tasks, poll, None, stop.clone(), wake); routes::serve(ctx, bind, stop) .await.expect("failed serving api routes"); }, Mode::Config => unreachable!(), #[cfg(feature = "migrate")] Mode::Migrate => unreachable!(), #[cfg(feature = "cli")] Mode::Cli { .. } => unreachable!(), } handle.close(); signals_task.await.expect("failed joining signal handler task"); } struct WakerToken(tokio::sync::mpsc::UnboundedSender<()>); impl context::WakerToken for WakerToken { fn wake(&self) { self.0.send(()).warn_failed("failed waking up workers"); } } struct WakeToken(tokio::sync::mpsc::UnboundedReceiver<()>); impl worker::WakeToken for WakeToken { async fn wait(&mut self) { let _ = self.0.recv().await; } } #[derive(Clone)] struct CancellationToken(tokio::sync::watch::Receiver); impl worker::StopToken for CancellationToken { fn stop(&self) -> bool { *self.0.borrow() } } impl routes::ShutdownToken for CancellationToken { async fn event(mut self) { self.0.changed().await.warn_failed("cancellation token channel closed, stopping..."); } } async fn handle_signals( mut signals: signal_hook_tokio::Signals, tx: tokio::sync::watch::Sender, ) { while let Some(signal) = signals.next().await { match signal { SIGTERM | SIGINT => { tracing::info!("received stop signal, closing tasks"); tx.send(true).info_failed("error sending stop signal to tasks") }, _ => unreachable!(), } } } #[derive(Debug, Clone, clap::ValueEnum)] enum Filter { All, Delivery, Inbound, Outbound, } impl From for Option { fn from(value: Filter) -> Self { match value { Filter::All => None, Filter::Delivery => Some(upub::model::job::JobType::Delivery), Filter::Inbound => Some(upub::model::job::JobType::Inbound), Filter::Outbound => Some(upub::model::job::JobType::Outbound), } } }