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