1use anyhow::anyhow;
2use axum::{
3 extract::MatchedPath,
4 http::{Request, Response},
5 routing::get,
6 Extension, Router,
7};
8use collab::llm::db::LlmDatabase;
9use collab::migrations::run_database_migrations;
10use collab::{api::billing::poll_stripe_events_periodically, llm::LlmState, ServiceMode};
11use collab::{
12 api::fetch_extensions_from_blob_store_periodically, db, env, executor::Executor,
13 rpc::ResultExt, AppState, Config, RateLimiter, Result,
14};
15use db::Database;
16use std::{
17 env::args,
18 net::{SocketAddr, TcpListener},
19 path::Path,
20 sync::Arc,
21 time::Duration,
22};
23#[cfg(unix)]
24use tokio::signal::unix::SignalKind;
25use tower_http::trace::TraceLayer;
26use tracing_subscriber::{
27 filter::EnvFilter, fmt::format::JsonFields, util::SubscriberInitExt, Layer,
28};
29use util::ResultExt as _;
30
31const VERSION: &str = env!("CARGO_PKG_VERSION");
32const REVISION: Option<&'static str> = option_env!("GITHUB_SHA");
33
34#[tokio::main]
35async fn main() -> Result<()> {
36 if let Err(error) = env::load_dotenv() {
37 eprintln!(
38 "error loading .env.toml (this is expected in production): {}",
39 error
40 );
41 }
42
43 let mut args = args().skip(1);
44 match args.next().as_deref() {
45 Some("version") => {
46 println!("collab v{} ({})", VERSION, REVISION.unwrap_or("unknown"));
47 }
48 Some("migrate") => {
49 let config = envy::from_env::<Config>().expect("error loading config");
50 setup_app_database(&config).await?;
51 }
52 Some("seed") => {
53 let config = envy::from_env::<Config>().expect("error loading config");
54 let db_options = db::ConnectOptions::new(config.database_url.clone());
55 let mut db = Database::new(db_options, Executor::Production).await?;
56 db.initialize_notification_kinds().await?;
57
58 collab::seed::seed(&config, &db, true).await?;
59 }
60 Some("serve") => {
61 let mode = match args.next().as_deref() {
62 Some("collab") => ServiceMode::Collab,
63 Some("api") => ServiceMode::Api,
64 Some("llm") => ServiceMode::Llm,
65 Some("all") => ServiceMode::All,
66 _ => {
67 return Err(anyhow!(
68 "usage: collab <version | migrate | seed | serve <api|collab|llm|all>>"
69 ))?;
70 }
71 };
72
73 let config = envy::from_env::<Config>().expect("error loading config");
74 init_tracing(&config);
75 let mut app = Router::new()
76 .route("/", get(handle_root))
77 .route("/healthz", get(handle_liveness_probe))
78 .layer(Extension(mode));
79
80 let listener = TcpListener::bind(&format!("0.0.0.0:{}", config.http_port))
81 .expect("failed to bind TCP listener");
82
83 let mut on_shutdown = None;
84
85 if mode.is_llm() {
86 setup_llm_database(&config).await?;
87
88 let state = LlmState::new(config.clone(), Executor::Production).await?;
89
90 app = app
91 .merge(collab::llm::routes())
92 .layer(Extension(state.clone()));
93 }
94
95 if mode.is_collab() || mode.is_api() {
96 setup_app_database(&config).await?;
97
98 let state = AppState::new(config, Executor::Production).await?;
99
100 if mode.is_collab() {
101 state.db.purge_old_embeddings().await.trace_err();
102 RateLimiter::save_periodically(
103 state.rate_limiter.clone(),
104 state.executor.clone(),
105 );
106
107 let epoch = state
108 .db
109 .create_server(&state.config.zed_environment)
110 .await?;
111 let rpc_server = collab::rpc::Server::new(epoch, state.clone());
112 rpc_server.start().await?;
113
114 app = app
115 .merge(collab::api::routes(rpc_server.clone()))
116 .merge(collab::rpc::routes(rpc_server.clone()));
117
118 on_shutdown = Some(Box::new(move || rpc_server.teardown()));
119 }
120
121 if mode.is_api() {
122 poll_stripe_events_periodically(state.clone());
123 fetch_extensions_from_blob_store_periodically(state.clone());
124
125 app = app
126 .merge(collab::api::events::router())
127 .merge(collab::api::extensions::router())
128 }
129
130 app = app.layer(Extension(state.clone()));
131 }
132
133 app = app.layer(
134 TraceLayer::new_for_http()
135 .make_span_with(|request: &Request<_>| {
136 let matched_path = request
137 .extensions()
138 .get::<MatchedPath>()
139 .map(MatchedPath::as_str);
140
141 tracing::info_span!(
142 "http_request",
143 method = ?request.method(),
144 matched_path,
145 )
146 })
147 .on_response(
148 |response: &Response<_>, latency: Duration, _: &tracing::Span| {
149 let duration_ms = latency.as_micros() as f64 / 1000.;
150 tracing::info!(
151 duration_ms,
152 status = response.status().as_u16(),
153 "finished processing request"
154 );
155 },
156 ),
157 );
158
159 #[cfg(unix)]
160 let signal = async move {
161 let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
162 .expect("failed to listen for interrupt signal");
163 let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
164 .expect("failed to listen for interrupt signal");
165 let sigterm = sigterm.recv();
166 let sigint = sigint.recv();
167 futures::pin_mut!(sigterm, sigint);
168 futures::future::select(sigterm, sigint).await;
169 };
170
171 #[cfg(windows)]
172 let signal = async move {
173 // todo(windows):
174 // `ctrl_close` does not work well, because tokio's signal handler always returns soon,
175 // but system terminates the application soon after returning CTRL+CLOSE handler.
176 // So we should implement blocking handler to treat CTRL+CLOSE signal.
177 let mut ctrl_break = tokio::signal::windows::ctrl_break()
178 .expect("failed to listen for interrupt signal");
179 let mut ctrl_c = tokio::signal::windows::ctrl_c()
180 .expect("failed to listen for interrupt signal");
181 let ctrl_break = ctrl_break.recv();
182 let ctrl_c = ctrl_c.recv();
183 futures::pin_mut!(ctrl_break, ctrl_c);
184 futures::future::select(ctrl_break, ctrl_c).await;
185 };
186
187 axum::Server::from_tcp(listener)
188 .map_err(|e| anyhow!(e))?
189 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
190 .with_graceful_shutdown(async move {
191 signal.await;
192 tracing::info!("Received interrupt signal");
193
194 if let Some(on_shutdown) = on_shutdown {
195 on_shutdown();
196 }
197 })
198 .await
199 .map_err(|e| anyhow!(e))?;
200 }
201 _ => {
202 Err(anyhow!(
203 "usage: collab <version | migrate | seed | serve <api|collab|llm|all>>"
204 ))?;
205 }
206 }
207 Ok(())
208}
209
210async fn setup_app_database(config: &Config) -> Result<()> {
211 let db_options = db::ConnectOptions::new(config.database_url.clone());
212 let mut db = Database::new(db_options, Executor::Production).await?;
213
214 let migrations_path = config.migrations_path.as_deref().unwrap_or_else(|| {
215 #[cfg(feature = "sqlite")]
216 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations.sqlite");
217 #[cfg(not(feature = "sqlite"))]
218 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
219
220 Path::new(default_migrations)
221 });
222
223 let migrations = run_database_migrations(db.options(), migrations_path).await?;
224 for (migration, duration) in migrations {
225 log::info!(
226 "Migrated {} {} {:?}",
227 migration.version,
228 migration.description,
229 duration
230 );
231 }
232
233 db.initialize_notification_kinds().await?;
234
235 if config.seed_path.is_some() {
236 collab::seed::seed(&config, &db, false).await?;
237 }
238
239 Ok(())
240}
241
242async fn setup_llm_database(config: &Config) -> Result<()> {
243 // TODO: This is temporary until we have the LLM database stood up.
244 if !config.is_development() {
245 return Ok(());
246 }
247
248 let database_url = config
249 .llm_database_url
250 .as_ref()
251 .ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;
252
253 let db_options = db::ConnectOptions::new(database_url.clone());
254 let db = LlmDatabase::new(db_options, Executor::Production).await?;
255
256 let migrations_path = config
257 .llm_database_migrations_path
258 .as_deref()
259 .unwrap_or_else(|| {
260 #[cfg(feature = "sqlite")]
261 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations_llm.sqlite");
262 #[cfg(not(feature = "sqlite"))]
263 let default_migrations = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations_llm");
264
265 Path::new(default_migrations)
266 });
267
268 let migrations = run_database_migrations(db.options(), migrations_path).await?;
269 for (migration, duration) in migrations {
270 log::info!(
271 "Migrated {} {} {:?}",
272 migration.version,
273 migration.description,
274 duration
275 );
276 }
277
278 Ok(())
279}
280
281async fn handle_root(Extension(mode): Extension<ServiceMode>) -> String {
282 format!(
283 "collab {mode:?} v{VERSION} ({})",
284 REVISION.unwrap_or("unknown")
285 )
286}
287
288async fn handle_liveness_probe(
289 app_state: Option<Extension<Arc<AppState>>>,
290 llm_state: Option<Extension<Arc<LlmState>>>,
291) -> Result<String> {
292 if let Some(state) = app_state {
293 state.db.get_all_users(0, 1).await?;
294 }
295
296 if let Some(_llm_state) = llm_state {}
297
298 Ok("ok".to_string())
299}
300
301pub fn init_tracing(config: &Config) -> Option<()> {
302 use std::str::FromStr;
303 use tracing_subscriber::layer::SubscriberExt;
304
305 let filter = EnvFilter::from_str(config.rust_log.as_deref()?).log_err()?;
306
307 tracing_subscriber::registry()
308 .with(if config.log_json.unwrap_or(false) {
309 Box::new(
310 tracing_subscriber::fmt::layer()
311 .fmt_fields(JsonFields::default())
312 .event_format(
313 tracing_subscriber::fmt::format()
314 .json()
315 .flatten_event(true)
316 .with_span_list(false),
317 )
318 .with_filter(filter),
319 ) as Box<dyn Layer<_> + Send + Sync>
320 } else {
321 Box::new(
322 tracing_subscriber::fmt::layer()
323 .event_format(tracing_subscriber::fmt::format().pretty())
324 .with_filter(filter),
325 )
326 })
327 .init();
328
329 None
330}