gitdataai/apps/git-hook/src/main.rs

use clap::Parser;
use config::AppConfig;
use db::cache::AppCache;
use db::database::AppDatabase;
use git::hook::HookService;
use metrics::{describe_counter, Unit};
use metrics_exporter_prometheus::PrometheusHandle;
use observability::{init_tracing_subscriber, install_recorder};
use sea_orm::ConnectionTrait;
use std::sync::Arc;
use tokio::signal;

mod args;

use args::HookArgs;

async fn http_handler(
    db: Arc<AppDatabase>,
    cache: Arc<AppCache>,
    metrics: Arc<PrometheusHandle>,
    req: hyper::Request<hyper::Body>,
) -> Result<hyper::Response<hyper::Body>, std::convert::Infallible> {
    match req.uri().path() {
        "/health" => {
            let writer_ok = db
                .writer()
                .execute_unprepared("SELECT 1")
                .await
                .is_ok();
            let reader_ok = db
                .reader()
                .execute_unprepared("SELECT 1")
                .await
                .is_ok();
            let db_ok = writer_ok && reader_ok;
            let cache_ok = cache.conn().await.is_ok();

            let body = serde_json::json!({
                "status": if db_ok && cache_ok { "ok" } else { "unhealthy" },
                "db": if db_ok { "ok" } else { "error" },
                "cache": if cache_ok { "ok" } else { "error" },
            });

            let status = if db_ok && cache_ok { 200 } else { 503 };
            let body_bytes = match serde_json::to_string(&body) {
                Ok(s) => hyper::Body::from(s),
                Err(e) => return Ok(hyper::Response::builder()
                    .status(500)
                    .body(hyper::Body::from(format!("serialize error: {}", e)))
                    .expect("static response")),
            };
            Ok(hyper::Response::builder()
                .status(status)
                .header("content-type", "application/json")
                .body(body_bytes)
                .expect("static response"))
        }
        "/metrics" => {
            let body = metrics.render();
            Ok(hyper::Response::builder()
                .status(200)
                .header("content-type", "text/plain; version=0.0.4; charset=utf-8")
                .body(hyper::Body::from(body))
                .expect("static response"))
        }
        _ => Ok(hyper::Response::builder()
            .status(404)
            .body(hyper::Body::from("not found"))
            .expect("static response")),
    }
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // 1. Load configuration
    let cfg = AppConfig::load();

    // 2. Init tracing + metrics
    let log_level = cfg.log_level().unwrap_or_else(|_| "info".to_string());
    init_tracing_subscriber(&log_level, false);

    // Pre-register all hook metrics so they appear in /metrics even before first increment.
    describe_counter!("hook_tasks_total", Unit::Count, "Total hook tasks dequeued");
    describe_counter!("hook_tasks_success_total", Unit::Count, "Hook tasks completed successfully");
    describe_counter!("hook_tasks_failed_total", Unit::Count, "Hook tasks that failed");
    describe_counter!("hook_tasks_locked_total", Unit::Count, "Hook tasks re-queued due to repo lock");
    describe_counter!("hook_tasks_retried_total", Unit::Count, "Hook tasks that entered retry");
    describe_counter!("hook_tasks_exhausted_total", Unit::Count, "Hook tasks that exhausted retries");
    describe_counter!("hook_sync_branches_changed_total", Unit::Count, "Branches changed during sync");
    describe_counter!("hook_sync_tags_changed_total", Unit::Count, "Tags changed during sync");

    let metrics_handle = Arc::new(install_recorder());

    // 3. Connect to database
    let db = Arc::new(AppDatabase::init(&cfg).await?);
    tracing::info!("database connected");

    // 4. Connect to Redis cache (also provides the cluster pool for hook queue)
    let cache = Arc::new(AppCache::init(&cfg).await?);
    tracing::info!("cache connected");

    // 5. Parse CLI args
    let _args = HookArgs::parse();

    tracing::info!("git-hook worker starting");

    // 6. Build and start git hook service
    let hooks = HookService::new(
        (*db).clone(),
        (*cache).clone(),
        cache.redis_pool().clone(),
        cfg,
    );

    let cancel = hooks.start_worker();
    let cancel_signal = cancel.clone();

    // 7. Start health/metrics server on a dedicated port
    let health_db = db.clone();
    let health_cache = cache.clone();
    let health_metrics = metrics_handle.clone();
    let health_addr: std::net::SocketAddr =
        ([0, 0, 0, 0], 8083).into();
    let health_service = hyper::service::make_service_fn(move |_| {
        let db = health_db.clone();
        let cache = health_cache.clone();
        let metrics = health_metrics.clone();
        let service = hyper::service::service_fn(move |req| {
            http_handler(db.clone(), cache.clone(), metrics.clone(), req)
        });
        async move { Ok::<_, std::convert::Infallible>(service) }
    });

    let health_server = hyper::Server::bind(&health_addr).serve(health_service);
    tracing::info!(port = 8083, "health/metrics server started");
    tokio::spawn(async move {
        if let Err(e) = health_server.await {
            tracing::error!("health check server error: {}", e);
        }
    });

    // Spawn signal handler that cancels on SIGINT/SIGTERM
    tokio::spawn(async move {
        let ctrl_c = async {
            signal::ctrl_c()
                .await
                .expect("failed to install CTRL+C handler");
        };

        #[cfg(unix)]
        let term = async {
            use tokio::signal::unix::{SignalKind, signal};
            let mut sig =
                signal(SignalKind::terminate()).expect("failed to install SIGTERM handler");
            sig.recv().await;
        };

        #[cfg(not(unix))]
        let term = std::future::pending::<()>();

        tokio::select! {
            _ = ctrl_c => {
                tracing::info!("received SIGINT, initiating shutdown");
            }
            _ = term => {
                tracing::info!("received SIGTERM, initiating shutdown");
            }
        }
        cancel_signal.cancel();
    });

    // Wait until the worker is cancelled (by signal handler or otherwise)
    cancel.cancelled().await;
    tracing::info!("git-hook worker stopped");
    Ok(())
}