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 anyhow::anyhow;
14use axum::{routing::get, Router};
15use collab::{Error, Result};
16use db::{Db, PostgresDb};
17use serde::Deserialize;
18use std::{
19 env::args,
20 net::{SocketAddr, TcpListener},
21 path::PathBuf,
22 sync::Arc,
23 time::Duration,
24};
25use tokio::signal;
26use tracing_log::LogTracer;
27use tracing_subscriber::{filter::EnvFilter, fmt::format::JsonFields, Layer};
28use util::ResultExt;
29
30const VERSION: &'static str = env!("CARGO_PKG_VERSION");
31
32#[derive(Default, Deserialize)]
33pub struct Config {
34 pub http_port: u16,
35 pub database_url: String,
36 pub api_token: String,
37 pub invite_link_prefix: String,
38 pub live_kit_server: Option<String>,
39 pub live_kit_key: Option<String>,
40 pub live_kit_secret: Option<String>,
41 pub rust_log: Option<String>,
42 pub log_json: Option<bool>,
43}
44
45#[derive(Default, Deserialize)]
46pub struct MigrateConfig {
47 pub database_url: String,
48 pub migrations_path: Option<PathBuf>,
49}
50
51pub struct AppState {
52 db: Arc<dyn Db>,
53 live_kit_client: Option<Arc<dyn live_kit_server::api::Client>>,
54 config: Config,
55}
56
57impl AppState {
58 async fn new(config: Config) -> Result<Arc<Self>> {
59 let db = PostgresDb::new(&config.database_url, 5).await?;
60 let live_kit_client = if let Some(((server, key), secret)) = config
61 .live_kit_server
62 .as_ref()
63 .zip(config.live_kit_key.as_ref())
64 .zip(config.live_kit_secret.as_ref())
65 {
66 Some(Arc::new(live_kit_server::api::LiveKitClient::new(
67 server.clone(),
68 key.clone(),
69 secret.clone(),
70 )) as Arc<dyn live_kit_server::api::Client>)
71 } else {
72 None
73 };
74
75 let this = Self {
76 db: Arc::new(db),
77 live_kit_client,
78 config,
79 };
80 Ok(Arc::new(this))
81 }
82}
83
84#[tokio::main]
85async fn main() -> Result<()> {
86 if let Err(error) = env::load_dotenv() {
87 eprintln!(
88 "error loading .env.toml (this is expected in production): {}",
89 error
90 );
91 }
92
93 match args().skip(1).next().as_deref() {
94 Some("version") => {
95 println!("collab v{VERSION}");
96 }
97 Some("migrate") => {
98 let config = envy::from_env::<MigrateConfig>().expect("error loading config");
99 let db = PostgresDb::new(&config.database_url, 5).await?;
100
101 let migrations_path = config
102 .migrations_path
103 .as_deref()
104 .or(db::DEFAULT_MIGRATIONS_PATH.map(|s| s.as_ref()))
105 .ok_or_else(|| anyhow!("missing MIGRATIONS_PATH environment variable"))?;
106
107 let migrations = db.migrate(&migrations_path, false).await?;
108 for (migration, duration) in migrations {
109 println!(
110 "Ran {} {} {:?}",
111 migration.version, migration.description, duration
112 );
113 }
114
115 return Ok(());
116 }
117 Some("serve") => {
118 let config = envy::from_env::<Config>().expect("error loading config");
119 init_tracing(&config);
120
121 let state = AppState::new(config).await?;
122 let listener = TcpListener::bind(&format!("0.0.0.0:{}", state.config.http_port))
123 .expect("failed to bind TCP listener");
124
125 let rpc_server = rpc::Server::new(state.clone(), None);
126 rpc_server
127 .start_recording_project_activity(Duration::from_secs(5 * 60), rpc::RealExecutor);
128
129 let app = api::routes(rpc_server.clone(), state.clone())
130 .merge(rpc::routes(rpc_server.clone()))
131 .merge(Router::new().route("/", get(handle_root)));
132
133 axum::Server::from_tcp(listener)?
134 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
135 .with_graceful_shutdown(graceful_shutdown(rpc_server, state))
136 .await?;
137 }
138 _ => {
139 Err(anyhow!("usage: collab <version | migrate | serve>"))?;
140 }
141 }
142 Ok(())
143}
144
145async fn handle_root() -> String {
146 format!("collab v{VERSION}")
147}
148
149pub fn init_tracing(config: &Config) -> Option<()> {
150 use std::str::FromStr;
151 use tracing_subscriber::layer::SubscriberExt;
152 let rust_log = config.rust_log.clone()?;
153
154 LogTracer::init().log_err()?;
155
156 let subscriber = tracing_subscriber::Registry::default()
157 .with(if config.log_json.unwrap_or(false) {
158 Box::new(
159 tracing_subscriber::fmt::layer()
160 .fmt_fields(JsonFields::default())
161 .event_format(
162 tracing_subscriber::fmt::format()
163 .json()
164 .flatten_event(true)
165 .with_span_list(true),
166 ),
167 ) as Box<dyn Layer<_> + Send + Sync>
168 } else {
169 Box::new(
170 tracing_subscriber::fmt::layer()
171 .event_format(tracing_subscriber::fmt::format().pretty()),
172 )
173 })
174 .with(EnvFilter::from_str(rust_log.as_str()).log_err()?);
175
176 tracing::subscriber::set_global_default(subscriber).unwrap();
177
178 None
179}
180
181async fn graceful_shutdown(rpc_server: Arc<rpc::Server>, state: Arc<AppState>) {
182 let ctrl_c = async {
183 signal::ctrl_c()
184 .await
185 .expect("failed to install Ctrl+C handler");
186 };
187
188 #[cfg(unix)]
189 let terminate = async {
190 signal::unix::signal(signal::unix::SignalKind::terminate())
191 .expect("failed to install signal handler")
192 .recv()
193 .await;
194 };
195
196 #[cfg(not(unix))]
197 let terminate = std::future::pending::<()>();
198
199 tokio::select! {
200 _ = ctrl_c => {},
201 _ = terminate => {},
202 }
203
204 if let Some(live_kit) = state.live_kit_client.as_ref() {
205 let deletions = rpc_server
206 .store()
207 .await
208 .rooms()
209 .values()
210 .map(|room| {
211 let name = room.live_kit_room.clone();
212 async {
213 live_kit.delete_room(name).await.trace_err();
214 }
215 })
216 .collect::<Vec<_>>();
217
218 tracing::info!("deleting all live-kit rooms");
219 if let Err(_) = tokio::time::timeout(
220 Duration::from_secs(10),
221 futures::future::join_all(deletions),
222 )
223 .await
224 {
225 tracing::error!("timed out waiting for live-kit room deletion");
226 }
227 }
228}