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::ServiceMode;
11use collab::api::CloudflareIpCountryHeader;
12use collab::{
13 AppState, Config, Result, api::fetch_extensions_from_blob_store_periodically, db, env,
14 executor::Executor,
15};
16use db::Database;
17use std::{
18 env::args,
19 net::{SocketAddr, TcpListener},
20 sync::Arc,
21 time::Duration,
22};
23#[cfg(unix)]
24use tokio::signal::unix::SignalKind;
25use tower_http::trace::TraceLayer;
26use tracing_subscriber::{
27 Layer, filter::EnvFilter, fmt::format::JsonFields, util::SubscriberInitExt,
28};
29use util::ResultExt as _;
30
31const VERSION: &str = env!("CARGO_PKG_VERSION");
32const REVISION: Option<&'static str> = option_env!("GITHUB_SHA");
33
34#[expect(clippy::result_large_err)]
35#[tokio::main]
36async fn main() -> Result<()> {
37 if let Err(error) = env::load_dotenv() {
38 eprintln!(
39 "error loading .env.toml (this is expected in production): {}",
40 error
41 );
42 }
43
44 let mut args = args().skip(1);
45 match args.next().as_deref() {
46 Some("version") => {
47 println!("collab v{} ({})", VERSION, REVISION.unwrap_or("unknown"));
48 }
49 Some("seed") => {
50 let config = envy::from_env::<Config>().expect("error loading config");
51 let db_options = db::ConnectOptions::new(config.database_url.clone());
52
53 let mut db = Database::new(db_options).await?;
54 db.initialize_notification_kinds().await?;
55
56 collab::seed::seed(&config, &db, false).await?;
57 }
58 Some("serve") => {
59 let mode = match args.next().as_deref() {
60 Some("collab") => ServiceMode::Collab,
61 Some("api") => ServiceMode::Api,
62 Some("all") => ServiceMode::All,
63 _ => {
64 return Err(anyhow!(
65 "usage: collab <version | seed | serve <api|collab|all>>"
66 ))?;
67 }
68 };
69
70 let config = envy::from_env::<Config>().expect("error loading config");
71 init_tracing(&config);
72 init_panic_hook();
73
74 let mut app = Router::new()
75 .route("/", get(handle_root))
76 .route("/healthz", get(handle_liveness_probe))
77 .layer(Extension(mode));
78
79 let listener = TcpListener::bind(format!("0.0.0.0:{}", config.http_port))
80 .expect("failed to bind TCP listener");
81
82 let mut on_shutdown = None;
83
84 if mode.is_collab() || mode.is_api() {
85 setup_app_database(&config).await?;
86
87 let state = AppState::new(config, Executor::Production).await?;
88
89 if mode.is_collab() {
90 let epoch = state
91 .db
92 .create_server(&state.config.zed_environment)
93 .await?;
94 let rpc_server = collab::rpc::Server::new(epoch, state.clone());
95 rpc_server.start().await?;
96
97 app = app
98 .merge(collab::api::routes(rpc_server.clone()))
99 .merge(collab::rpc::routes(rpc_server.clone()));
100
101 on_shutdown = Some(Box::new(move || rpc_server.teardown()));
102 }
103
104 if mode.is_api() {
105 fetch_extensions_from_blob_store_periodically(state.clone());
106
107 app = app
108 .merge(collab::api::events::router())
109 .merge(collab::api::extensions::router())
110 }
111
112 app = app.layer(Extension(state.clone()));
113 }
114
115 app = app.layer(
116 TraceLayer::new_for_http()
117 .make_span_with(|request: &Request<_>| {
118 let matched_path = request
119 .extensions()
120 .get::<MatchedPath>()
121 .map(MatchedPath::as_str);
122
123 let geoip_country_code = request
124 .headers()
125 .typed_get::<CloudflareIpCountryHeader>()
126 .map(|header| header.to_string());
127
128 tracing::info_span!(
129 "http_request",
130 method = ?request.method(),
131 matched_path,
132 geoip_country_code,
133 user_id = tracing::field::Empty,
134 login = tracing::field::Empty,
135 authn.jti = tracing::field::Empty,
136 is_staff = tracing::field::Empty
137 )
138 })
139 .on_response(
140 |response: &Response<_>, latency: Duration, _: &tracing::Span| {
141 let duration_ms = latency.as_micros() as f64 / 1000.;
142 tracing::info!(
143 duration_ms,
144 status = response.status().as_u16(),
145 "finished processing request"
146 );
147 },
148 ),
149 );
150
151 #[cfg(unix)]
152 let signal = async move {
153 let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
154 .expect("failed to listen for interrupt signal");
155 let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
156 .expect("failed to listen for interrupt signal");
157 let sigterm = sigterm.recv();
158 let sigint = sigint.recv();
159 futures::pin_mut!(sigterm, sigint);
160 futures::future::select(sigterm, sigint).await;
161 };
162
163 #[cfg(windows)]
164 let signal = async move {
165 // todo(windows):
166 // `ctrl_close` does not work well, because tokio's signal handler always returns soon,
167 // but system terminates the application soon after returning CTRL+CLOSE handler.
168 // So we should implement blocking handler to treat CTRL+CLOSE signal.
169 let mut ctrl_break = tokio::signal::windows::ctrl_break()
170 .expect("failed to listen for interrupt signal");
171 let mut ctrl_c = tokio::signal::windows::ctrl_c()
172 .expect("failed to listen for interrupt signal");
173 let ctrl_break = ctrl_break.recv();
174 let ctrl_c = ctrl_c.recv();
175 futures::pin_mut!(ctrl_break, ctrl_c);
176 futures::future::select(ctrl_break, ctrl_c).await;
177 };
178
179 axum::Server::from_tcp(listener)
180 .map_err(|e| anyhow!(e))?
181 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
182 .with_graceful_shutdown(async move {
183 signal.await;
184 tracing::info!("Received interrupt signal");
185
186 if let Some(on_shutdown) = on_shutdown {
187 on_shutdown();
188 }
189 })
190 .await
191 .map_err(|e| anyhow!(e))?;
192 }
193 _ => {
194 Err(anyhow!(
195 "usage: collab <version | migrate | seed | serve <api|collab|llm|all>>"
196 ))?;
197 }
198 }
199 Ok(())
200}
201
202async fn setup_app_database(config: &Config) -> Result<()> {
203 let db_options = db::ConnectOptions::new(config.database_url.clone());
204 let mut db = Database::new(db_options).await?;
205
206 db.initialize_notification_kinds().await?;
207
208 if config.seed_path.is_some() {
209 collab::seed::seed(config, &db, false).await?;
210 }
211
212 Ok(())
213}
214
215async fn handle_root(Extension(mode): Extension<ServiceMode>) -> String {
216 format!("zed:{mode} v{VERSION} ({})", REVISION.unwrap_or("unknown"))
217}
218
219async fn handle_liveness_probe(app_state: Option<Extension<Arc<AppState>>>) -> Result<String> {
220 if let Some(state) = app_state {
221 state.db.get_all_users(0, 1).await?;
222 }
223
224 Ok("ok".to_string())
225}
226
227pub fn init_tracing(config: &Config) -> Option<()> {
228 use std::str::FromStr;
229 use tracing_subscriber::layer::SubscriberExt;
230
231 let filter = EnvFilter::from_str(config.rust_log.as_deref()?).log_err()?;
232
233 tracing_subscriber::registry()
234 .with(if config.log_json.unwrap_or(false) {
235 Box::new(
236 tracing_subscriber::fmt::layer()
237 .fmt_fields(JsonFields::default())
238 .event_format(
239 tracing_subscriber::fmt::format()
240 .json()
241 .flatten_event(true)
242 .with_span_list(false),
243 )
244 .with_filter(filter),
245 ) as Box<dyn Layer<_> + Send + Sync>
246 } else {
247 Box::new(
248 tracing_subscriber::fmt::layer()
249 .event_format(tracing_subscriber::fmt::format().pretty())
250 .with_filter(filter),
251 )
252 })
253 .init();
254
255 None
256}
257
258fn init_panic_hook() {
259 std::panic::set_hook(Box::new(move |panic_info| {
260 let panic_message = match panic_info.payload().downcast_ref::<&'static str>() {
261 Some(message) => *message,
262 None => match panic_info.payload().downcast_ref::<String>() {
263 Some(message) => message.as_str(),
264 None => "Box<Any>",
265 },
266 };
267 let backtrace = std::backtrace::Backtrace::force_capture();
268 let location = panic_info
269 .location()
270 .map(|loc| format!("{}:{}", loc.file(), loc.line()));
271 tracing::error!(panic = true, ?location, %panic_message, %backtrace, "Server Panic");
272 }));
273}