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