//! Copyright (c) 2022-2026 GitDataAi All rights reserved. use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use gitks::disk_cache::DiskCache; use gitks::hooks::HookManager; use gitks::metrics; use gitks::server::{GitksService, serve_with_shutdown}; use etcd_client::{Client, PutOptions}; use tokio::sync::Mutex; use tracing_subscriber::EnvFilter; use tracing_subscriber::fmt; use tracing_subscriber::prelude::*; const DEFAULT_HOST: &str = "0.0.0.0"; const DEFAULT_PORT: &str = "50051"; const DEFAULT_STORAGE_NAME: &str = "default"; /// etcd-backed config reader. Priority: etcd > env > default. struct EtcdConfig { client: Arc>, prefix: String, } impl EtcdConfig { async fn connect(endpoints: Vec, prefix: &str) -> Result { let client = Client::connect(endpoints, None) .await .map_err(|e| format!("etcd connect: {e}"))?; Ok(Self { client: Arc::new(Mutex::new(client)), prefix: prefix.to_string(), }) } /// Get config: etcd first, env second, default last. async fn get(&self, key: &str, default: &str) -> String { let etcd_key = format!("{}config/{}", self.prefix, key); if let Ok(mut c) = self.client.try_lock() && let Ok(resp) = c.get(etcd_key.as_str(), None).await && let Some(kv) = resp.kvs().first() && let Ok(v) = kv.value_str() && !v.is_empty() { return v.to_string(); } std::env::var(key).unwrap_or_else(|_| default.to_string()) } /// Register this service under the common prefix for discovery by other services. async fn register(&self, service_name: &str, addr: &str) -> Result<(), String> { let instance_id = uuid::Uuid::now_v7().to_string(); let addr = addr.to_string(); let key = format!("{}services/{}/{}", self.prefix, service_name, instance_id); let instance = serde_json::json!({"addr": &addr, "port": 0, "version": env!("CARGO_PKG_VERSION")}); let value = serde_json::to_string(&instance).map_err(|e| format!("json: {e}"))?; let lease = { let mut c = self.client.lock().await; c.lease_grant(15, None) .await .map_err(|e| format!("lease: {e}"))? }; { let mut c = self.client.lock().await; let opts = PutOptions::new().with_lease(lease.id()); c.put(key.clone(), value, Some(opts)) .await .map_err(|e| format!("put: {e}"))?; } tracing::info!(service = service_name, addr = %addr, "registered in etcd"); let c = self.client.clone(); tokio::spawn(async move { loop { let r = { let mut cl = c.lock().await; cl.lease_keep_alive(lease.id()).await }; drop(r); tokio::time::sleep(std::time::Duration::from_secs(5)).await; if let Ok(lr) = { let mut cl = c.lock().await; cl.lease_grant(15, None).await } { let inst = serde_json::json!({"addr": &addr, "port": 0, "version": env!("CARGO_PKG_VERSION")}); if let Ok(v) = serde_json::to_string(&inst) { let mut cl = c.lock().await; let _ = cl .put(key.clone(), v, Some(PutOptions::new().with_lease(lr.id()))) .await; } } } }); Ok(()) } } fn env_or(key: &str, default: &str) -> String { std::env::var(key).unwrap_or_else(|_| default.into()) } fn env_bool(key: &str, default: bool) -> bool { match std::env::var(key).as_deref() { Ok("true" | "1" | "yes") => true, Ok("false" | "0" | "no") => false, _ => default, } } fn env_u64(key: &str, default: u64) -> u64 { std::env::var(key) .ok() .and_then(|v| v.parse().ok()) .unwrap_or(default) } fn init_tracing() -> Option { let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); let log_format = env_or("GITKS_LOG_FORMAT", "pretty"); let fmt_layer = match log_format.as_str() { "json" => fmt::layer() .json() .with_target(true) .with_file(true) .with_line_number(true) .with_thread_ids(true) .with_span_events(fmt::format::FmtSpan::NEW | fmt::format::FmtSpan::CLOSE) .boxed(), _ => fmt::layer() .pretty() .with_target(true) .with_file(true) .with_line_number(true) .boxed(), }; if let Ok(log_dir) = std::env::var("GITKS_LOG_DIR") { let rotation = match env_or("GITKS_LOG_ROTATION", "daily").as_str() { "hourly" => tracing_appender::rolling::Rotation::HOURLY, "never" => tracing_appender::rolling::Rotation::NEVER, _ => tracing_appender::rolling::Rotation::DAILY, }; let retention = env_u64("GITKS_LOG_RETENTION", 7) as usize; let mut builder = tracing_appender::rolling::Builder::new() .rotation(rotation) .filename_prefix("gitks") .filename_suffix("log"); if retention > 0 { builder = builder.max_log_files(retention); } let file_appender = match builder.build(&log_dir) { Ok(file_appender) => file_appender, Err(err) => { eprintln!("failed to create log directory '{log_dir}': {err}"); tracing_subscriber::registry() .with(env_filter) .with(fmt_layer) .init(); return None; } }; let (non_blocking, guard) = tracing_appender::non_blocking(file_appender); let file_layer = fmt::layer() .json() .with_target(true) .with_file(true) .with_line_number(true) .with_writer(non_blocking) .with_filter(EnvFilter::new("info")) .boxed(); tracing_subscriber::registry() .with(env_filter) .with(fmt_layer) .with(file_layer) .init(); Some(guard) } else { tracing_subscriber::registry() .with(env_filter) .with(fmt_layer) .init(); None } } #[tokio::main] async fn main() -> Result<(), Box> { dotenvy::dotenv().ok(); let _log_guard = init_tracing(); tracing::info!( version = env!("CARGO_PKG_VERSION"), log_format = %env_or("GITKS_LOG_FORMAT", "pretty"), "gitks starting up" ); let host = env_or("GITKS_HOST", DEFAULT_HOST); let port = env_or("GITKS_PORT", DEFAULT_PORT); let storage_name = env_or("STORAGE_NAME", DEFAULT_STORAGE_NAME); let etcd_endpoints: Vec = std::env::var("GITKS_ETCD_ENDPOINTS") .ok() .filter(|s| !s.is_empty()) .map(|s| s.split(',').map(str::trim).map(String::from).collect()) .unwrap_or_else(|| vec!["http://localhost:2379".to_string()]); let etcd_prefix = std::env::var("ETCD_KEY_PREFIX").unwrap_or_else(|_| "/appks/".to_string()); let etcd = EtcdConfig::connect(etcd_endpoints, &etcd_prefix).await.ok(); let host = if let Some(ref e) = etcd { e.get("GITKS_HOST", &host).await } else { host }; let port = if let Some(ref e) = etcd { e.get("GITKS_PORT", &port).await } else { port }; let storage_name = if let Some(ref e) = etcd { e.get("GITKS_STORAGE_NAME", &storage_name).await } else { storage_name }; let grpc_addr = std::env::var("GITKS_ADVERTISE_ADDR").unwrap_or_else(|_| format!("http://{host}:{port}")); if let Some(ref e) = etcd { let addr_str = format!("{host}:{port}"); e.register("gitks", &addr_str).await.ok(); } let repo_prefix = std::env::var("REPO_PREFIX_PATH") .map_err(|_| "REPO_PREFIX_PATH environment variable is required (e.g. /data/repos)")?; let repo_prefix = PathBuf::from(&repo_prefix); if !repo_prefix.is_absolute() { return Err("REPO_PREFIX_PATH must be an absolute path".into()); } if !repo_prefix.exists() { tracing::info!(path = %repo_prefix.display(), "creating repo prefix directory"); std::fs::create_dir_all(&repo_prefix)?; } let disk_cache_enabled = env_bool("GITKS_DISK_CACHE_ENABLED", false); let disk_cache_max_age = env_u64("GITKS_DISK_CACHE_MAX_AGE", 300); let disk_cache = DiskCache::new( repo_prefix.clone(), env!("CARGO_PKG_VERSION").to_string(), disk_cache_max_age, disk_cache_enabled, ); if disk_cache_enabled { tracing::info!(max_age_secs = disk_cache_max_age, "disk cache enabled"); disk_cache.cleanup_on_startup()?; gitks::disk_cache::start_cache_cleanup_task(disk_cache.clone(), Duration::from_secs(300)); } else { tracing::info!("disk cache disabled"); } let pack_cache_enabled = env_bool("GITKS_PACK_CACHE_ENABLED", false); let pack_backpressure = env_bool("GITKS_PACK_CACHE_BACKPRESSURE", true); let pack_cache = if disk_cache_enabled { tracing::info!( pack_objects_cache = pack_cache_enabled, backpressure = pack_backpressure, "pack cache wrapper enabled" ); Some(gitks::pack_cache::PackCache::new( disk_cache.clone(), pack_backpressure, )) } else { None }; let hooks_enabled = env_bool("GITKS_HOOKS_ENABLED", true); let server_hooks_dir = std::env::var("GITKS_SERVER_HOOKS_DIR") .ok() .map(PathBuf::from); let hook_callback_addr = std::env::var("GITKS_HOOK_CALLBACK_ADDR").ok(); let hook_timeout = env_u64("GITKS_HOOK_TIMEOUT", 30); let allow_custom_hooks = env_bool("GITKS_ALLOW_CUSTOM_HOOKS", true); let hook_manager = if hooks_enabled { tracing::info!( timeout_secs = hook_timeout, custom_hooks = allow_custom_hooks, "hooks enabled" ); Some(HookManager::new( repo_prefix.clone(), server_hooks_dir, hook_callback_addr, Duration::from_secs(hook_timeout), allow_custom_hooks, )) } else { tracing::info!("hooks disabled"); None }; let metrics_port = env_u64("GITKS_METRICS_PORT", 9100) as u16; let http_cancel = tokio_util::sync::CancellationToken::new(); metrics::set_http_cancel_token(http_cancel.clone()); let _metrics_handle = metrics::start_metrics_server(metrics_port); tracing::info!(port = metrics_port, "metrics server started"); let _semaphore_cleanup = gitks::rate_limit::start_semaphore_cleanup_task(); let slow_request_threshold = env_u64("GITKS_SLOW_REQUEST_THRESHOLD_MS", 5000); metrics::set_slow_request_threshold(slow_request_threshold); tracing::info!( threshold_ms = slow_request_threshold, "slow request detection configured" ); let addr: std::net::SocketAddr = format!("{host}:{port}").parse()?; let mut svc = GitksService::new(repo_prefix.clone()); if disk_cache_enabled { svc = svc.with_disk_cache(disk_cache); } if let Some(pc) = pack_cache { svc = svc.with_pack_cache(pc); } if let Some(hm) = hook_manager { svc = svc.with_hook_manager(hm); } let svc = svc.with_grpc_addr(grpc_addr.clone()); tracing::info!( addr = %addr, repo_prefix = %repo_prefix.display(), storage = %storage_name, advertise = %grpc_addr, "starting gitks gRPC server" ); metrics::set_ready(true); serve_with_shutdown(addr, svc, shutdown_signal()).await?; metrics::set_ready(false); http_cancel.cancel(); tracing::info!("gitks shut down complete"); Ok(()) } /// Resolves when the process receives SIGTERM or SIGINT (Ctrl+C). async fn shutdown_signal() { let ctrl_c = async { tokio::signal::ctrl_c() .await .expect("failed to install Ctrl+C handler"); }; #[cfg(unix)] let terminate = async { tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) .expect("failed to install SIGTERM handler") .recv() .await; }; #[cfg(not(unix))] let terminate = std::future::pending::<()>(); tokio::select! { _ = ctrl_c => { tracing::info!("received Ctrl+C, starting graceful shutdown"); } _ = terminate => { tracing::info!("received SIGTERM, starting graceful shutdown"); } } }