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 epoch = state
 61                .db
 62                .create_server(&state.config.zed_environment)
 63                .await?;
 64            let rpc_server = collab::rpc::Server::new(epoch, state.clone(), Executor::Production);
 65            rpc_server.start().await?;
 66
 67            let app = collab::api::routes(rpc_server.clone(), state.clone())
 68                .merge(collab::rpc::routes(rpc_server.clone()))
 69                .merge(Router::new().route("/", get(handle_root)));
 70
 71            axum::Server::from_tcp(listener)?
 72                .serve(app.into_make_service_with_connect_info::<SocketAddr>())
 73                .with_graceful_shutdown(async move {
 74                    let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
 75                        .expect("failed to listen for interrupt signal");
 76                    let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
 77                        .expect("failed to listen for interrupt signal");
 78                    let sigterm = sigterm.recv();
 79                    let sigint = sigint.recv();
 80                    futures::pin_mut!(sigterm);
 81                    futures::pin_mut!(sigint);
 82                    futures::future::select(sigterm, sigint).await;
 83                    tracing::info!("Received interrupt signal");
 84                    rpc_server.teardown();
 85                })
 86                .await?;
 87        }
 88        _ => {
 89            Err(anyhow!("usage: collab <version | migrate | serve>"))?;
 90        }
 91    }
 92    Ok(())
 93}
 94
 95async fn handle_root() -> String {
 96    format!("collab v{VERSION}")
 97}
 98
 99pub fn init_tracing(config: &Config) -> Option<()> {
100    use std::str::FromStr;
101    use tracing_subscriber::layer::SubscriberExt;
102    let rust_log = config.rust_log.clone()?;
103
104    LogTracer::init().log_err()?;
105
106    let subscriber = tracing_subscriber::Registry::default()
107        .with(if config.log_json.unwrap_or(false) {
108            Box::new(
109                tracing_subscriber::fmt::layer()
110                    .fmt_fields(JsonFields::default())
111                    .event_format(
112                        tracing_subscriber::fmt::format()
113                            .json()
114                            .flatten_event(true)
115                            .with_span_list(true),
116                    ),
117            ) as Box<dyn Layer<_> + Send + Sync>
118        } else {
119            Box::new(
120                tracing_subscriber::fmt::layer()
121                    .event_format(tracing_subscriber::fmt::format().pretty()),
122            )
123        })
124        .with(EnvFilter::from_str(rust_log.as_str()).log_err()?);
125
126    tracing::subscriber::set_global_default(subscriber).unwrap();
127
128    None
129}