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