//! 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>, 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(), }) } 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> { dotenvy::dotenv().ok(); let mut config = gitks::GitksConfig::from_env()?; // Overlay etcd config (etcd > env > default) 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()); 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(()) }