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