main.rs

  1mod api;
  2mod auth;
  3mod db;
  4mod env;
  5mod rpc;
  6
  7#[cfg(test)]
  8mod db_tests;
  9#[cfg(test)]
 10mod integration_tests;
 11
 12use crate::rpc::ResultExt as _;
 13use axum::{body::Body, Router};
 14use collab::{Error, Result};
 15use db::{Db, PostgresDb};
 16use serde::Deserialize;
 17use std::{
 18    net::{SocketAddr, TcpListener},
 19    sync::Arc,
 20    time::Duration,
 21};
 22use tokio::signal;
 23use tracing_log::LogTracer;
 24use tracing_subscriber::{filter::EnvFilter, fmt::format::JsonFields, Layer};
 25use util::ResultExt;
 26
 27#[derive(Default, Deserialize)]
 28pub struct Config {
 29    pub http_port: u16,
 30    pub database_url: String,
 31    pub api_token: String,
 32    pub invite_link_prefix: String,
 33    pub honeycomb_api_key: Option<String>,
 34    pub honeycomb_dataset: Option<String>,
 35    pub live_kit_server: Option<String>,
 36    pub live_kit_key: Option<String>,
 37    pub live_kit_secret: Option<String>,
 38    pub rust_log: Option<String>,
 39    pub log_json: Option<bool>,
 40}
 41
 42pub struct AppState {
 43    db: Arc<dyn Db>,
 44    api_token: String,
 45    invite_link_prefix: String,
 46    live_kit_client: Option<Arc<dyn live_kit_server::api::Client>>,
 47}
 48
 49impl AppState {
 50    async fn new(config: &Config) -> Result<Arc<Self>> {
 51        let db = PostgresDb::new(&config.database_url, 5).await?;
 52        let live_kit_client = if let Some(((server, key), secret)) = config
 53            .live_kit_server
 54            .as_ref()
 55            .zip(config.live_kit_key.as_ref())
 56            .zip(config.live_kit_secret.as_ref())
 57        {
 58            Some(Arc::new(live_kit_server::api::LiveKitClient::new(
 59                server.clone(),
 60                key.clone(),
 61                secret.clone(),
 62            )) as Arc<dyn live_kit_server::api::Client>)
 63        } else {
 64            None
 65        };
 66
 67        let this = Self {
 68            db: Arc::new(db),
 69            live_kit_client,
 70            api_token: config.api_token.clone(),
 71            invite_link_prefix: config.invite_link_prefix.clone(),
 72        };
 73        Ok(Arc::new(this))
 74    }
 75}
 76
 77#[tokio::main]
 78async fn main() -> Result<()> {
 79    if let Err(error) = env::load_dotenv() {
 80        eprintln!(
 81            "error loading .env.toml (this is expected in production): {}",
 82            error
 83        );
 84    }
 85
 86    let config = envy::from_env::<Config>().expect("error loading config");
 87    init_tracing(&config);
 88    let state = AppState::new(&config).await?;
 89
 90    let listener = TcpListener::bind(&format!("0.0.0.0:{}", config.http_port))
 91        .expect("failed to bind TCP listener");
 92    let rpc_server = rpc::Server::new(state.clone(), None);
 93
 94    rpc_server.start_recording_project_activity(Duration::from_secs(5 * 60), rpc::RealExecutor);
 95
 96    let app = Router::<Body>::new()
 97        .merge(api::routes(rpc_server.clone(), state.clone()))
 98        .merge(rpc::routes(rpc_server.clone()));
 99
100    axum::Server::from_tcp(listener)?
101        .serve(app.into_make_service_with_connect_info::<SocketAddr>())
102        .with_graceful_shutdown(graceful_shutdown(rpc_server, state))
103        .await?;
104
105    Ok(())
106}
107
108pub fn init_tracing(config: &Config) -> Option<()> {
109    use std::str::FromStr;
110    use tracing_subscriber::layer::SubscriberExt;
111    let rust_log = config.rust_log.clone()?;
112
113    LogTracer::init().log_err()?;
114
115    let subscriber = tracing_subscriber::Registry::default()
116        .with(if config.log_json.unwrap_or(false) {
117            Box::new(
118                tracing_subscriber::fmt::layer()
119                    .fmt_fields(JsonFields::default())
120                    .event_format(
121                        tracing_subscriber::fmt::format()
122                            .json()
123                            .flatten_event(true)
124                            .with_span_list(true),
125                    ),
126            ) as Box<dyn Layer<_> + Send + Sync>
127        } else {
128            Box::new(
129                tracing_subscriber::fmt::layer()
130                    .event_format(tracing_subscriber::fmt::format().pretty()),
131            )
132        })
133        .with(EnvFilter::from_str(rust_log.as_str()).log_err()?);
134
135    tracing::subscriber::set_global_default(subscriber).unwrap();
136
137    None
138}
139
140async fn graceful_shutdown(rpc_server: Arc<rpc::Server>, state: Arc<AppState>) {
141    let ctrl_c = async {
142        signal::ctrl_c()
143            .await
144            .expect("failed to install Ctrl+C handler");
145    };
146
147    #[cfg(unix)]
148    let terminate = async {
149        signal::unix::signal(signal::unix::SignalKind::terminate())
150            .expect("failed to install signal handler")
151            .recv()
152            .await;
153    };
154
155    #[cfg(not(unix))]
156    let terminate = std::future::pending::<()>();
157
158    tokio::select! {
159        _ = ctrl_c => {},
160        _ = terminate => {},
161    }
162
163    if let Some(live_kit) = state.live_kit_client.as_ref() {
164        let deletions = rpc_server
165            .store()
166            .await
167            .rooms()
168            .values()
169            .map(|room| {
170                let name = room.live_kit_room.clone();
171                async {
172                    live_kit.delete_room(name).await.trace_err();
173                }
174            })
175            .collect::<Vec<_>>();
176
177        tracing::info!("deleting all live-kit rooms");
178        if let Err(_) = tokio::time::timeout(
179            Duration::from_secs(10),
180            futures::future::join_all(deletions),
181        )
182        .await
183        {
184            tracing::error!("timed out waiting for live-kit room deletion");
185        }
186    }
187}