gitks/main.rs

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

use std::sync::Arc;

use etcd_client::{Client, PutOptions};
use tokio::sync::Mutex;

/// etcd-backed config reader. Priority: etcd > env > default.
struct EtcdConfig {
    client: Arc<Mutex<Client>>,
    prefix: String,
}

impl EtcdConfig {
    async fn connect(endpoints: Vec<String>, prefix: &str) -> Result<Self, String> {
        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(),
        })
    }

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

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

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenvy::dotenv().ok();

    let mut config = gitks::GitksConfig::from_env()?;

    // Overlay etcd config (etcd > env > default)
    let etcd_endpoints: Vec<String> = 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());

    if let Ok(etcd) = EtcdConfig::connect(etcd_endpoints, &etcd_prefix).await {
        config.host = etcd.get("GITKS_HOST", &config.host).await;
        config.port = etcd.get("GITKS_PORT", &config.port).await;
        config.storage_name = etcd.get("GITKS_STORAGE_NAME", &config.storage_name).await;
        config.grpc_addr = etcd
            .get("GITKS_ADVERTISE_ADDR", &config.grpc_addr)
            .await;

        let addr_str = format!("{}:{}", config.host, config.port);
        etcd.register("gitks", &addr_str).await.ok();
    }

    let server = gitks::GitksServer::builder().config(config).build()?;
    server.serve().await?;

    Ok(())
}