1use anyhow::anyhow;
2use axum::{routing::get, Extension, Router};
3use collab::{db, env, executor::Executor, AppState, Config, MigrateConfig, Result};
4use db::Database;
5use std::{
6 env::args,
7 net::{SocketAddr, TcpListener},
8 path::Path,
9 sync::Arc,
10};
11use tokio::signal::unix::SignalKind;
12use tracing_log::LogTracer;
13use tracing_subscriber::{filter::EnvFilter, fmt::format::JsonFields, Layer};
14use util::ResultExt;
15
16const VERSION: &'static str = env!("CARGO_PKG_VERSION");
17
18#[tokio::main]
19async fn main() -> Result<()> {
20 if let Err(error) = env::load_dotenv() {
21 eprintln!(
22 "error loading .env.toml (this is expected in production): {}",
23 error
24 );
25 }
26
27 match args().skip(1).next().as_deref() {
28 Some("version") => {
29 println!("collab v{VERSION}");
30 }
31 Some("migrate") => {
32 let config = envy::from_env::<MigrateConfig>().expect("error loading config");
33 let mut db_options = db::ConnectOptions::new(config.database_url.clone());
34 db_options.max_connections(5);
35 let db = Database::new(db_options, Executor::Production).await?;
36
37 let migrations_path = config
38 .migrations_path
39 .as_deref()
40 .unwrap_or_else(|| Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")));
41
42 let migrations = db.migrate(&migrations_path, false).await?;
43 for (migration, duration) in migrations {
44 println!(
45 "Ran {} {} {:?}",
46 migration.version, migration.description, duration
47 );
48 }
49
50 return Ok(());
51 }
52 Some("serve") => {
53 let config = envy::from_env::<Config>().expect("error loading config");
54 init_tracing(&config);
55
56 if config.is_development() {
57 // sanity check database url so even if we deploy a busted ZED_ENVIRONMENT to production
58 // we do not run
59 if config.database_url != "postgres://postgres@localhost/zed" {
60 panic!("about to run development migrations on a non-development database?")
61 }
62 let migrations_path = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations"));
63 let db_options = db::ConnectOptions::new(config.database_url.clone());
64 let db = Database::new(db_options, Executor::Production).await?;
65
66 let migrations = db.migrate(&migrations_path, false).await?;
67 for (migration, duration) in migrations {
68 println!(
69 "Ran {} {} {:?}",
70 migration.version, migration.description, duration
71 );
72 }
73 }
74
75 let state = AppState::new(config).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 epoch = state
81 .db
82 .create_server(&state.config.zed_environment)
83 .await?;
84 let rpc_server = collab::rpc::Server::new(epoch, state.clone(), Executor::Production);
85 rpc_server.start().await?;
86
87 let app = collab::api::routes(rpc_server.clone(), state.clone())
88 .merge(collab::rpc::routes(rpc_server.clone()))
89 .merge(
90 Router::new()
91 .route("/", get(handle_root))
92 .route("/healthz", get(handle_liveness_probe))
93 .layer(Extension(state.clone())),
94 );
95
96 axum::Server::from_tcp(listener)?
97 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
98 .with_graceful_shutdown(async move {
99 let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
100 .expect("failed to listen for interrupt signal");
101 let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
102 .expect("failed to listen for interrupt signal");
103 let sigterm = sigterm.recv();
104 let sigint = sigint.recv();
105 futures::pin_mut!(sigterm, sigint);
106 futures::future::select(sigterm, sigint).await;
107 tracing::info!("Received interrupt signal");
108 rpc_server.teardown();
109 })
110 .await?;
111 }
112 _ => {
113 Err(anyhow!("usage: collab <version | migrate | serve>"))?;
114 }
115 }
116 Ok(())
117}
118
119async fn handle_root() -> String {
120 format!("collab v{VERSION}")
121}
122
123async fn handle_liveness_probe(Extension(state): Extension<Arc<AppState>>) -> Result<String> {
124 state.db.get_all_users(0, 1).await?;
125 Ok("ok".to_string())
126}
127
128pub fn init_tracing(config: &Config) -> Option<()> {
129 use std::str::FromStr;
130 use tracing_subscriber::layer::SubscriberExt;
131 let rust_log = config.rust_log.clone()?;
132
133 LogTracer::init().log_err()?;
134
135 let subscriber = tracing_subscriber::Registry::default()
136 .with(if config.log_json.unwrap_or(false) {
137 Box::new(
138 tracing_subscriber::fmt::layer()
139 .fmt_fields(JsonFields::default())
140 .event_format(
141 tracing_subscriber::fmt::format()
142 .json()
143 .flatten_event(true)
144 .with_span_list(true),
145 ),
146 ) as Box<dyn Layer<_> + Send + Sync>
147 } else {
148 Box::new(
149 tracing_subscriber::fmt::layer()
150 .event_format(tracing_subscriber::fmt::format().pretty()),
151 )
152 })
153 .with(EnvFilter::from_str(rust_log.as_str()).log_err()?);
154
155 tracing::subscriber::set_global_default(subscriber).unwrap();
156
157 None
158}