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