main.rs

  1use anyhow::anyhow;
  2use axum::{routing::get, Router};
  3use collab::{db, env, executor::Executor, AppState, Config, MigrateConfig, Result};
  4use db::Database;
  5use std::{
  6    env::args,
  7    net::{SocketAddr, TcpListener},
  8    path::Path,
  9};
 10use tokio::signal::unix::SignalKind;
 11use tracing_log::LogTracer;
 12use tracing_subscriber::{filter::EnvFilter, fmt::format::JsonFields, Layer};
 13use util::ResultExt;
 14
 15const VERSION: &'static str = env!("CARGO_PKG_VERSION");
 16
 17#[tokio::main]
 18async fn main() -> Result<()> {
 19    if let Err(error) = env::load_dotenv() {
 20        eprintln!(
 21            "error loading .env.toml (this is expected in production): {}",
 22            error
 23        );
 24    }
 25
 26    match args().skip(1).next().as_deref() {
 27        Some("version") => {
 28            println!("collab v{VERSION}");
 29        }
 30        Some("migrate") => {
 31            let config = envy::from_env::<MigrateConfig>().expect("error loading config");
 32            let mut db_options = db::ConnectOptions::new(config.database_url.clone());
 33            db_options.max_connections(5);
 34            let db = Database::new(db_options).await?;
 35
 36            let migrations_path = config
 37                .migrations_path
 38                .as_deref()
 39                .unwrap_or_else(|| Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")));
 40
 41            let migrations = db.migrate(&migrations_path, false).await?;
 42            for (migration, duration) in migrations {
 43                println!(
 44                    "Ran {} {} {:?}",
 45                    migration.version, migration.description, duration
 46                );
 47            }
 48
 49            return Ok(());
 50        }
 51        Some("serve") => {
 52            let config = envy::from_env::<Config>().expect("error loading config");
 53            init_tracing(&config);
 54
 55            let state = AppState::new(config).await?;
 56
 57            let listener = TcpListener::bind(&format!("0.0.0.0:{}", state.config.http_port))
 58                .expect("failed to bind TCP listener");
 59
 60            let rpc_server = collab::rpc::Server::new(state.clone(), Executor::Production);
 61            rpc_server.start().await?;
 62
 63            let app = collab::api::routes(rpc_server.clone(), state.clone())
 64                .merge(collab::rpc::routes(rpc_server.clone()))
 65                .merge(Router::new().route("/", get(handle_root)));
 66
 67            axum::Server::from_tcp(listener)?
 68                .serve(app.into_make_service_with_connect_info::<SocketAddr>())
 69                .with_graceful_shutdown(async move {
 70                    let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
 71                        .expect("failed to listen for interrupt signal");
 72                    let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
 73                        .expect("failed to listen for interrupt signal");
 74                    let sigterm = sigterm.recv();
 75                    let sigint = sigint.recv();
 76                    futures::pin_mut!(sigterm);
 77                    futures::pin_mut!(sigint);
 78                    futures::future::select(sigterm, sigint).await;
 79                    tracing::info!("Received interrupt signal");
 80                    rpc_server.teardown();
 81                })
 82                .await?;
 83        }
 84        _ => {
 85            Err(anyhow!("usage: collab <version | migrate | serve>"))?;
 86        }
 87    }
 88    Ok(())
 89}
 90
 91async fn handle_root() -> String {
 92    format!("collab v{VERSION}")
 93}
 94
 95pub fn init_tracing(config: &Config) -> Option<()> {
 96    use std::str::FromStr;
 97    use tracing_subscriber::layer::SubscriberExt;
 98    let rust_log = config.rust_log.clone()?;
 99
100    LogTracer::init().log_err()?;
101
102    let subscriber = tracing_subscriber::Registry::default()
103        .with(if config.log_json.unwrap_or(false) {
104            Box::new(
105                tracing_subscriber::fmt::layer()
106                    .fmt_fields(JsonFields::default())
107                    .event_format(
108                        tracing_subscriber::fmt::format()
109                            .json()
110                            .flatten_event(true)
111                            .with_span_list(true),
112                    ),
113            ) as Box<dyn Layer<_> + Send + Sync>
114        } else {
115            Box::new(
116                tracing_subscriber::fmt::layer()
117                    .event_format(tracing_subscriber::fmt::format().pretty()),
118            )
119        })
120        .with(EnvFilter::from_str(rust_log.as_str()).log_err()?);
121
122    tracing::subscriber::set_global_default(subscriber).unwrap();
123
124    None
125}