gitks/lib.rs

//! Copyright (c) 2022-2026 GitDataAi All rights reserved.

pub mod archive;
pub mod bare;
pub mod blame;
pub mod blob;
pub mod branch;
pub mod commit;
pub mod config;
pub mod diff;
pub mod disk_cache;
pub mod error;
pub mod hooks;
pub mod init;
pub mod macros;
pub mod merge;
pub mod metrics;
pub mod oid;
pub mod pack;
pub mod pack_cache;
pub mod paginate;
pub mod pb;
pub mod rate_limit;
pub mod refs;
pub mod remote;
pub mod repository;
pub mod sanitize;
pub mod server;
pub mod snapshot;
pub mod tag;
pub mod tree;

use std::future::Future;
use std::net::SocketAddr;
use std::path::PathBuf;
use std::time::Duration;

use server::GitksService;

/// Configuration for building a [`GitksServer`].
///
/// Use [`GitksConfig::from_env`] to read from environment variables,
/// or construct manually when embedding as a library.
#[derive(Debug, Clone)]
pub struct GitksConfig {
    /// Repository storage prefix path (required).
    pub repo_prefix: PathBuf,
    /// gRPC listen host.
    pub host: String,
    /// gRPC listen port.
    pub port: String,
    /// Storage name used in repository headers.
    pub storage_name: String,
    /// Advertised gRPC address for other services.
    pub grpc_addr: String,
    /// Enable disk-based cache.
    pub disk_cache_enabled: bool,
    /// Max age for disk cache entries (seconds).
    pub disk_cache_max_age: u64,
    /// Enable pack-objects cache wrapper.
    pub pack_cache_enabled: bool,
    /// Enable pack cache backpressure.
    pub pack_cache_backpressure: bool,
    /// Enable hook support.
    pub hooks_enabled: bool,
    /// Server-side hooks directory.
    pub server_hooks_dir: Option<PathBuf>,
    /// Callback address for hooks.
    pub hook_callback_addr: Option<String>,
    /// Hook execution timeout (seconds).
    pub hook_timeout: u64,
    /// Allow custom hooks from repository config.
    pub allow_custom_hooks: bool,
    /// Prometheus metrics port.
    pub metrics_port: u16,
    /// Slow request detection threshold (ms).
    pub slow_request_threshold: u64,
    /// Log format: "pretty" or "json".
    pub log_format: String,
    /// Optional log directory for file output.
    pub log_dir: Option<String>,
    /// Log rotation strategy.
    pub log_rotation: String,
    /// Max log files to retain.
    pub log_retention: usize,
}

/// A ready-to-run gitks gRPC server.
pub struct GitksServer {
    service: GitksService,
    addr: SocketAddr,
    http_cancel: tokio_util::sync::CancellationToken,
    _log_guard: Option<tracing_appender::non_blocking::WorkerGuard>,
    _metrics_handle: tokio::task::JoinHandle<()>,
    _semaphore_cleanup: tokio::task::JoinHandle<()>,
}

/// Builder for [`GitksServer`].
///
/// # Examples
///
/// ```no_run
/// use gitks::{GitksServer, GitksConfig};
///
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// let config = GitksConfig::from_env()?;
/// let server = GitksServer::builder()
///     .config(config)
///     .build()?;
/// server.serve().await?;
/// # Ok(())
/// # }
/// ```
pub struct GitksServerBuilder {
    config: Option<GitksConfig>,
}

impl GitksConfig {
    /// Build config from environment variables.
    pub fn from_env() -> Result<Self, Box<dyn std::error::Error>> {
        dotenvy::dotenv().ok();
        Ok(Self {
            repo_prefix: PathBuf::from(
                std::env::var("REPO_PREFIX_PATH")
                    .map_err(|_| "REPO_PREFIX_PATH is required (e.g. /data/repos)")?,
            ),
            host: env_or("GITKS_HOST", "0.0.0.0"),
            port: env_or("GITKS_PORT", "50051"),
            storage_name: env_or("STORAGE_NAME", "default"),
            grpc_addr: std::env::var("GITKS_ADVERTISE_ADDR")
                .unwrap_or_else(|_| format!("http://{}:{}", env_or("GITKS_HOST", "0.0.0.0"), env_or("GITKS_PORT", "50051"))),
            disk_cache_enabled: env_bool("GITKS_DISK_CACHE_ENABLED", false),
            disk_cache_max_age: env_u64("GITKS_DISK_CACHE_MAX_AGE", 300),
            pack_cache_enabled: env_bool("GITKS_PACK_CACHE_ENABLED", false),
            pack_cache_backpressure: env_bool("GITKS_PACK_CACHE_BACKPRESSURE", true),
            hooks_enabled: env_bool("GITKS_HOOKS_ENABLED", true),
            server_hooks_dir: std::env::var("GITKS_SERVER_HOOKS_DIR").ok().map(PathBuf::from),
            hook_callback_addr: std::env::var("GITKS_HOOK_CALLBACK_ADDR").ok(),
            hook_timeout: env_u64("GITKS_HOOK_TIMEOUT", 30),
            allow_custom_hooks: env_bool("GITKS_ALLOW_CUSTOM_HOOKS", true),
            metrics_port: env_u64("GITKS_METRICS_PORT", 9100) as u16,
            slow_request_threshold: env_u64("GITKS_SLOW_REQUEST_THRESHOLD_MS", 5000),
            log_format: env_or("GITKS_LOG_FORMAT", "pretty"),
            log_dir: std::env::var("GITKS_LOG_DIR").ok(),
            log_rotation: env_or("GITKS_LOG_ROTATION", "daily"),
            log_retention: env_u64("GITKS_LOG_RETENTION", 7) as usize,
        })
    }
}

impl GitksServer {
    /// Create a new builder.
    pub fn builder() -> GitksServerBuilder {
        GitksServerBuilder::default()
    }

    /// Start the gRPC server and block until Ctrl+C (or SIGTERM on Unix).
    pub async fn serve(self) -> Result<(), Box<dyn std::error::Error>> {
        self.serve_with_shutdown(shutdown_signal()).await
    }

    /// Start the gRPC server and block until the provided `shutdown` future resolves.
    pub async fn serve_with_shutdown(
        self,
        shutdown: impl Future<Output = ()>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        metrics::set_ready(true);

        server::serve_with_shutdown(self.addr, self.service, shutdown).await?;

        metrics::set_ready(false);
        self.http_cancel.cancel();

        tracing::info!("gitks shut down complete");
        Ok(())
    }
}

impl GitksServerBuilder {
    /// Set the server configuration.
    pub fn config(mut self, config: GitksConfig) -> Self {
        self.config = Some(config);
        self
    }

    /// Build the server.
    pub fn build(self) -> Result<GitksServer, Box<dyn std::error::Error>> {
        let config = self.config.unwrap_or_else(|| {
            GitksConfig::from_env().expect("failed to load gitks config")
        });

        // Validate repo_prefix
        if !config.repo_prefix.is_absolute() {
            return Err("REPO_PREFIX_PATH must be an absolute path".into());
        }
        if !config.repo_prefix.exists() {
            tracing::info!(path = %config.repo_prefix.display(), "creating repo prefix directory");
            std::fs::create_dir_all(&config.repo_prefix)?;
        }

        // Init logging if log_dir is set
        let log_guard = init_tracing_from_config(&config);

        tracing::info!(
            version = env!("CARGO_PKG_VERSION"),
            log_format = %config.log_format,
            "gitks starting up"
        );

        // Disk cache
        let disk_cache = disk_cache::DiskCache::new(
            config.repo_prefix.clone(),
            env!("CARGO_PKG_VERSION").to_string(),
            config.disk_cache_max_age,
            config.disk_cache_enabled,
        );

        if config.disk_cache_enabled {
            tracing::info!(max_age_secs = config.disk_cache_max_age, "disk cache enabled");
            disk_cache.cleanup_on_startup()?;
            disk_cache::start_cache_cleanup_task(disk_cache.clone(), Duration::from_secs(300));
        } else {
            tracing::info!("disk cache disabled");
        }

        // Pack cache
        let pack_cache = if config.disk_cache_enabled {
            tracing::info!(
                pack_objects_cache = config.pack_cache_enabled,
                backpressure = config.pack_cache_backpressure,
                "pack cache wrapper enabled"
            );
            Some(pack_cache::PackCache::new(
                disk_cache.clone(),
                config.pack_cache_backpressure,
            ))
        } else {
            None
        };

        // Hook manager
        let hook_manager = if config.hooks_enabled {
            tracing::info!(
                timeout_secs = config.hook_timeout,
                custom_hooks = config.allow_custom_hooks,
                "hooks enabled"
            );
            Some(hooks::HookManager::new(
                config.repo_prefix.clone(),
                config.server_hooks_dir,
                config.hook_callback_addr,
                Duration::from_secs(config.hook_timeout),
                config.allow_custom_hooks,
            ))
        } else {
            tracing::info!("hooks disabled");
            None
        };

        // Metrics
        let http_cancel = tokio_util::sync::CancellationToken::new();
        metrics::set_http_cancel_token(http_cancel.clone());
        let metrics_handle = metrics::start_metrics_server(config.metrics_port);
        tracing::info!(port = config.metrics_port, "metrics server started");

        // Rate limiter cleanup
        let semaphore_cleanup = rate_limit::start_semaphore_cleanup_task();

        // Slow request detection
        metrics::set_slow_request_threshold(config.slow_request_threshold);
        tracing::info!(
            threshold_ms = config.slow_request_threshold,
            "slow request detection configured"
        );

        // Build service
        let addr: SocketAddr = format!("{}:{}", config.host, config.port).parse()?;
        let mut svc = GitksService::new(config.repo_prefix.clone());

        if config.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(config.grpc_addr.clone());

        tracing::info!(
            addr = %addr,
            repo_prefix = %config.repo_prefix.display(),
            storage = %config.storage_name,
            advertise = %config.grpc_addr,
            "starting gitks gRPC server"
        );

        Ok(GitksServer {
            service: svc,
            addr,
            http_cancel,
            _log_guard: log_guard,
            _metrics_handle: metrics_handle,
            _semaphore_cleanup: semaphore_cleanup,
        })
    }
}

impl Default for GitksServerBuilder {
    fn default() -> Self {
        Self { config: None }
    }
}

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_from_config(
    config: &GitksConfig,
) -> Option<tracing_appender::non_blocking::WorkerGuard> {
    use tracing_subscriber::EnvFilter;
    use tracing_subscriber::fmt;
    use tracing_subscriber::prelude::*;

    let env_filter =
        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));

    let fmt_layer = match config.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 Some(ref log_dir) = config.log_dir {
        let rotation = match config.log_rotation.as_str() {
            "hourly" => tracing_appender::rolling::Rotation::HOURLY,
            "never" => tracing_appender::rolling::Rotation::NEVER,
            _ => tracing_appender::rolling::Rotation::DAILY,
        };
        let retention = config.log_retention;

        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
    }
}

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");
        }
    }
}