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