Merge branch 'main' into auto-deploy-collab

crates/collab/src/main.rs

@@ -9,6 +9,7 @@ mod db_tests;
 #[cfg(test)]
 mod integration_tests;
 
+use crate::rpc::ResultExt as _;
 use anyhow::anyhow;
 use axum::{routing::get, Router};
 use collab::{Error, Result};
@@ -21,6 +22,7 @@ use std::{
     sync::Arc,
     time::Duration,
 };
+use tokio::signal;
 use tracing_log::LogTracer;
 use tracing_subscriber::{filter::EnvFilter, fmt::format::JsonFields, Layer};
 use util::ResultExt;
@@ -33,6 +35,9 @@ pub struct Config {
     pub database_url: String,
     pub api_token: String,
     pub invite_link_prefix: String,
+    pub live_kit_server: Option<String>,
+    pub live_kit_key: Option<String>,
+    pub live_kit_secret: Option<String>,
     pub rust_log: Option<String>,
     pub log_json: Option<bool>,
 }
@@ -45,9 +50,37 @@ pub struct MigrateConfig {
 
 pub struct AppState {
     db: Arc<dyn Db>,
+    live_kit_client: Option<Arc<dyn live_kit_server::api::Client>>,
     config: Config,
 }
 
+impl AppState {
+    async fn new(config: Config) -> Result<Arc<Self>> {
+        let db = PostgresDb::new(&config.database_url, 5).await?;
+        let live_kit_client = if let Some(((server, key), secret)) = config
+            .live_kit_server
+            .as_ref()
+            .zip(config.live_kit_key.as_ref())
+            .zip(config.live_kit_secret.as_ref())
+        {
+            Some(Arc::new(live_kit_server::api::LiveKitClient::new(
+                server.clone(),
+                key.clone(),
+                secret.clone(),
+            )) as Arc<dyn live_kit_server::api::Client>)
+        } else {
+            None
+        };
+
+        let this = Self {
+            db: Arc::new(db),
+            live_kit_client,
+            config,
+        };
+        Ok(Arc::new(this))
+    }
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
     if let Err(error) = env::load_dotenv() {
@@ -83,14 +116,9 @@ async fn main() -> Result<()> {
         }
         Some("serve") => {
             let config = envy::from_env::<Config>().expect("error loading config");
-            let db = PostgresDb::new(&config.database_url, 5).await?;
-
             init_tracing(&config);
-            let state = Arc::new(AppState {
-                db: Arc::new(db),
-                config,
-            });
 
+            let state = AppState::new(config).await?;
             let listener = TcpListener::bind(&format!("0.0.0.0:{}", state.config.http_port))
                 .expect("failed to bind TCP listener");
 
@@ -98,12 +126,13 @@ async fn main() -> Result<()> {
             rpc_server
                 .start_recording_project_activity(Duration::from_secs(5 * 60), rpc::RealExecutor);
 
-            let app = api::routes(&rpc_server, state.clone())
-                .merge(rpc::routes(rpc_server))
+            let app = api::routes(rpc_server.clone(), state.clone())
+                .merge(rpc::routes(rpc_server.clone()))
                 .merge(Router::new().route("/", get(handle_root)));
 
             axum::Server::from_tcp(listener)?
                 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
+                .with_graceful_shutdown(graceful_shutdown(rpc_server, state))
                 .await?;
         }
         _ => {
@@ -148,3 +177,52 @@ pub fn init_tracing(config: &Config) -> Option<()> {
 
     None
 }
+
+async fn graceful_shutdown(rpc_server: Arc<rpc::Server>, state: Arc<AppState>) {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => {},
+        _ = terminate => {},
+    }
+
+    if let Some(live_kit) = state.live_kit_client.as_ref() {
+        let deletions = rpc_server
+            .store()
+            .await
+            .rooms()
+            .values()
+            .map(|room| {
+                let name = room.live_kit_room.clone();
+                async {
+                    live_kit.delete_room(name).await.trace_err();
+                }
+            })
+            .collect::<Vec<_>>();
+
+        tracing::info!("deleting all live-kit rooms");
+        if let Err(_) = tokio::time::timeout(
+            Duration::from_secs(10),
+            futures::future::join_all(deletions),
+        )
+        .await
+        {
+            tracing::error!("timed out waiting for live-kit room deletion");
+        }
+    }
+}