1use anyhow::anyhow;
2use axum::{routing::get, 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};
10use tokio::signal::unix::SignalKind;
11use tracing_log::LogTracer;
12use tracing_subscriber::{filter::EnvFilter, fmt::format::JsonFields, Layer};
13use util::ResultExt;
14
15const VERSION: &'static str = env!("CARGO_PKG_VERSION");
16
17#[tokio::main]
18async fn main() -> Result<()> {
19 if let Err(error) = env::load_dotenv() {
20 eprintln!(
21 "error loading .env.toml (this is expected in production): {}",
22 error
23 );
24 }
25
26 match args().skip(1).next().as_deref() {
27 Some("version") => {
28 println!("collab v{VERSION}");
29 }
30 Some("migrate") => {
31 let config = envy::from_env::<MigrateConfig>().expect("error loading config");
32 let mut db_options = db::ConnectOptions::new(config.database_url.clone());
33 db_options.max_connections(5);
34 let db = Database::new(db_options).await?;
35
36 let migrations_path = config
37 .migrations_path
38 .as_deref()
39 .unwrap_or_else(|| Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")));
40
41 let migrations = db.migrate(&migrations_path, false).await?;
42 for (migration, duration) in migrations {
43 println!(
44 "Ran {} {} {:?}",
45 migration.version, migration.description, duration
46 );
47 }
48
49 return Ok(());
50 }
51 Some("serve") => {
52 let config = envy::from_env::<Config>().expect("error loading config");
53 init_tracing(&config);
54
55 let state = AppState::new(config).await?;
56
57 let listener = TcpListener::bind(&format!("0.0.0.0:{}", state.config.http_port))
58 .expect("failed to bind TCP listener");
59
60 let epoch = state
61 .db
62 .create_server(&state.config.zed_environment)
63 .await?;
64 let rpc_server = collab::rpc::Server::new(epoch, state.clone(), Executor::Production);
65 rpc_server.start().await?;
66
67 let app = collab::api::routes(rpc_server.clone(), state.clone())
68 .merge(collab::rpc::routes(rpc_server.clone()))
69 .merge(Router::new().route("/", get(handle_root)));
70
71 axum::Server::from_tcp(listener)?
72 .serve(app.into_make_service_with_connect_info::<SocketAddr>())
73 .with_graceful_shutdown(async move {
74 let mut sigterm = tokio::signal::unix::signal(SignalKind::terminate())
75 .expect("failed to listen for interrupt signal");
76 let mut sigint = tokio::signal::unix::signal(SignalKind::interrupt())
77 .expect("failed to listen for interrupt signal");
78 let sigterm = sigterm.recv();
79 let sigint = sigint.recv();
80 futures::pin_mut!(sigterm);
81 futures::pin_mut!(sigint);
82 futures::future::select(sigterm, sigint).await;
83 tracing::info!("Received interrupt signal");
84 rpc_server.teardown();
85 })
86 .await?;
87 }
88 _ => {
89 Err(anyhow!("usage: collab <version | migrate | serve>"))?;
90 }
91 }
92 Ok(())
93}
94
95async fn handle_root() -> String {
96 format!("collab v{VERSION}")
97}
98
99pub fn init_tracing(config: &Config) -> Option<()> {
100 use std::str::FromStr;
101 use tracing_subscriber::layer::SubscriberExt;
102 let rust_log = config.rust_log.clone()?;
103
104 LogTracer::init().log_err()?;
105
106 let subscriber = tracing_subscriber::Registry::default()
107 .with(if config.log_json.unwrap_or(false) {
108 Box::new(
109 tracing_subscriber::fmt::layer()
110 .fmt_fields(JsonFields::default())
111 .event_format(
112 tracing_subscriber::fmt::format()
113 .json()
114 .flatten_event(true)
115 .with_span_list(true),
116 ),
117 ) as Box<dyn Layer<_> + Send + Sync>
118 } else {
119 Box::new(
120 tracing_subscriber::fmt::layer()
121 .event_format(tracing_subscriber::fmt::format().pretty()),
122 )
123 })
124 .with(EnvFilter::from_str(rust_log.as_str()).log_err()?);
125
126 tracing::subscriber::set_global_default(subscriber).unwrap();
127
128 None
129}