feat: init

etcd/discovery.rs

@@ -0,0 +1,160 @@
+use etcd_client::{GetOptions, WatchOptions};
+use tokio_stream::StreamExt;
+use uuid::Uuid;
+
+use crate::error::{AppError, AppResult};
+use crate::pb::{EmailClient, RepoClient};
+
+use super::types::ServiceInstance;
+use super::{EtcdRegistry, EtcdRegistryInner};
+
+impl EtcdRegistry {
+    pub async fn start_discovery(&self) -> AppResult<()> {
+        self.load_initial("git").await?;
+        self.load_initial("mail").await?;
+        self.spawn_watch("git");
+        self.spawn_watch("mail");
+        Ok(())
+    }
+
+    async fn load_initial(&self, service: &str) -> AppResult<()> {
+        let prefix = self.service_prefix(service);
+        let resp = {
+            let mut client = self.inner.client.lock().await;
+            client
+                .get(prefix.as_str(), Some(GetOptions::new().with_prefix()))
+                .await
+                .map_err(|e| AppError::Config(format!("etcd get {prefix} failed: {e}")))?
+        };
+
+        for kv in resp.kvs() {
+            let key = kv.key_str().unwrap_or_default();
+            let value = kv.value_str().unwrap_or_default();
+            if let Ok(instance) = serde_json::from_str::<ServiceInstance>(value) {
+                Self::upsert_instance(&self.inner, service, key, &instance);
+            }
+        }
+
+        tracing::info!(
+            service = service,
+            prefix = prefix.as_str(),
+            "etcd initial discovery complete"
+        );
+        Ok(())
+    }
+
+    fn spawn_watch(&self, service: &str) {
+        let prefix = self.service_prefix(service);
+        let inner = self.inner.clone();
+        let service = service.to_string();
+
+        tokio::spawn(async move {
+            loop {
+                match Self::watch_loop(&inner, &prefix, &service).await {
+                    Ok(()) => break,
+                    Err(e) => {
+                        tracing::warn!(
+                            service = service.as_str(),
+                            error = %e,
+                            "etcd watch disconnected, retrying in 3s"
+                        );
+                        tokio::time::sleep(std::time::Duration::from_secs(3)).await;
+                    }
+                }
+            }
+        });
+    }
+
+    async fn watch_loop(inner: &EtcdRegistryInner, prefix: &str, service: &str) -> AppResult<()> {
+        let (mut watcher, mut stream) = {
+            let mut client = inner.client.lock().await;
+            client
+                .watch(prefix, Some(WatchOptions::new().with_prefix()))
+                .await
+                .map_err(|e| AppError::Config(format!("etcd watch {prefix} failed: {e}")))?
+        };
+
+        let _keep = &mut watcher;
+
+        while let Some(resp) = stream.next().await {
+            let resp =
+                resp.map_err(|e| AppError::Config(format!("etcd watch stream error: {e}")))?;
+
+            for event in resp.events() {
+                let Some(kv) = event.kv() else { continue };
+                let key = kv.key_str().unwrap_or_default();
+
+                match event.event_type() {
+                    etcd_client::EventType::Put => {
+                        let value = kv.value_str().unwrap_or_default();
+                        if let Ok(instance) = serde_json::from_str::<ServiceInstance>(value) {
+                            Self::upsert_instance(inner, service, key, &instance);
+                            tracing::info!(service = service, key = key, "etcd service upserted");
+                        }
+                    }
+                    etcd_client::EventType::Delete => {
+                        Self::remove_instance(inner, service, key);
+                        tracing::info!(service = service, key = key, "etcd service removed");
+                    }
+                }
+            }
+        }
+        Ok(())
+    }
+
+    pub(crate) fn service_prefix(&self, service: &str) -> String {
+        format!("{}services/{service}/", self.inner.key_prefix)
+    }
+
+    fn extract_id_from_key(key: &str) -> Option<Uuid> {
+        key.rsplit('/').next()?.parse().ok()
+    }
+
+    fn upsert_instance(
+        inner: &EtcdRegistryInner,
+        service: &str,
+        key: &str,
+        instance: &ServiceInstance,
+    ) {
+        let Some(node_id) = Self::extract_id_from_key(key) else {
+            tracing::warn!(key = key, "etcd key has no valid UUID suffix");
+            return;
+        };
+        let addr = instance.addr.clone();
+
+        match service {
+            "git" => match RepoClient::lazy_connect(&addr) {
+                Ok(client) => {
+                    inner.git_nodes.insert(node_id, client);
+                }
+                Err(e) => {
+                    tracing::error!(key = key, addr = addr.as_str(), error = %e, "git client connect failed");
+                }
+            },
+            "mail" => match EmailClient::lazy_connect(&addr) {
+                Ok(client) => {
+                    inner.mail_nodes.insert(node_id, client);
+                }
+                Err(e) => {
+                    tracing::error!(key = key, addr = addr.as_str(), error = %e, "mail client connect failed");
+                }
+            },
+            _ => {}
+        }
+    }
+
+    fn remove_instance(inner: &EtcdRegistryInner, service: &str, key: &str) {
+        let Some(node_id) = Self::extract_id_from_key(key) else {
+            return;
+        };
+        match service {
+            "git" => {
+                inner.git_nodes.remove(&node_id);
+            }
+            "mail" => {
+                inner.mail_nodes.remove(&node_id);
+            }
+            _ => {}
+        }
+    }
+}

etcd/mod.rs

@@ -0,0 +1,88 @@
+mod discovery;
+mod register;
+mod types;
+
+pub use types::ServiceInstance;
+
+use std::sync::Arc;
+use std::sync::atomic::AtomicI64;
+
+use dashmap::DashMap;
+use etcd_client::Client;
+use tokio::sync::Mutex;
+use uuid::Uuid;
+
+use crate::config::AppConfig;
+use crate::error::{AppError, AppResult};
+use crate::pb::{EmailClient, RepoClient};
+
+#[derive(Clone)]
+pub struct EtcdRegistry {
+    pub(crate) inner: Arc<EtcdRegistryInner>,
+}
+
+pub(crate) struct EtcdRegistryInner {
+    pub client: Mutex<Client>,
+    pub config: AppConfig,
+    pub key_prefix: String,
+    pub git_nodes: DashMap<Uuid, RepoClient>,
+    pub mail_nodes: DashMap<Uuid, EmailClient>,
+    pub lease_id: AtomicI64,
+}
+
+impl EtcdRegistry {
+    pub async fn connect(config: &AppConfig) -> AppResult<Self> {
+        let endpoints = config.etcd_endpoints()?;
+        let timeout = config.etcd_connect_timeout()?;
+
+        let opts = etcd_client::ConnectOptions::new()
+            .with_connect_timeout(std::time::Duration::from_secs(timeout));
+
+        let client = Client::connect(&endpoints, Some(opts))
+            .await
+            .map_err(|e| AppError::Config(format!("etcd connect failed: {e}")))?;
+
+        if let (Some(user), Some(pass)) = (config.etcd_username()?, config.etcd_password()?) {
+            let auth_resp = client
+                .auth_client()
+                .authenticate(user, pass)
+                .await
+                .map_err(|e| AppError::Config(format!("etcd auth failed: {e}")))?;
+            let token = auth_resp.token().to_string();
+            tracing::info!(token_len = token.len(), "etcd authenticated");
+        }
+
+        let key_prefix = config.etcd_key_prefix()?;
+
+        Ok(Self {
+            inner: Arc::new(EtcdRegistryInner {
+                client: Mutex::new(client),
+                config: config.clone(),
+                key_prefix,
+                git_nodes: DashMap::new(),
+                mail_nodes: DashMap::new(),
+                lease_id: AtomicI64::new(0),
+            }),
+        })
+    }
+
+    pub fn get_git_client(&self, node_id: &Uuid) -> Option<RepoClient> {
+        self.inner.git_nodes.get(node_id).map(|c| c.clone())
+    }
+
+    pub fn git_node_ids(&self) -> Vec<Uuid> {
+        self.inner.git_nodes.iter().map(|e| *e.key()).collect()
+    }
+
+    pub fn get_email_client(&self) -> Option<EmailClient> {
+        self.inner
+            .mail_nodes
+            .iter()
+            .next()
+            .map(|e| e.value().clone())
+    }
+
+    pub fn has_git_nodes(&self) -> bool {
+        !self.inner.git_nodes.is_empty()
+    }
+}

etcd/register.rs

@@ -0,0 +1,113 @@
+use std::collections::HashMap;
+use std::sync::atomic::Ordering;
+
+use etcd_client::PutOptions;
+use tokio_stream::StreamExt;
+
+use crate::error::{AppError, AppResult};
+
+use super::EtcdRegistry;
+use super::types::ServiceInstance;
+
+impl EtcdRegistry {
+    pub async fn register_self(&self, service_name: &str) -> AppResult<()> {
+        let ttl = self.inner.config.etcd_lease_ttl()?;
+        let listen_addr = self.inner.config.rpc_self_listen_addr()?;
+
+        let instance_id = uuid::Uuid::now_v7().to_string();
+        let key = format!(
+            "{}services/{service_name}/{instance_id}",
+            self.inner.key_prefix
+        );
+
+        let instance = ServiceInstance {
+            addr: listen_addr.clone(),
+            metadata: HashMap::new(),
+        };
+        let value = serde_json::to_string(&instance)?;
+
+        let lease_resp = {
+            let mut client = self.inner.client.lock().await;
+            client
+                .lease_grant(ttl as i64, None)
+                .await
+                .map_err(|e| AppError::Config(format!("etcd lease_grant failed: {e}")))?
+        };
+        let lease_id = lease_resp.id();
+        self.inner.lease_id.store(lease_id, Ordering::SeqCst);
+
+        {
+            let mut client = self.inner.client.lock().await;
+            let opts = PutOptions::new().with_lease(lease_id);
+            client
+                .put(key.clone(), value, Some(opts))
+                .await
+                .map_err(|e| AppError::Config(format!("etcd put failed: {e}")))?;
+        }
+
+        tracing::info!(
+            service = service_name,
+            addr = listen_addr.as_str(),
+            lease_id = lease_id,
+            "registered self in etcd"
+        );
+
+        self.spawn_keep_alive(lease_id, key);
+
+        Ok(())
+    }
+
+    fn spawn_keep_alive(&self, lease_id: i64, key: String) {
+        let inner = self.inner.clone();
+        let interval = self.inner.config.etcd_keep_alive_interval().unwrap_or(10);
+
+        tokio::spawn(async move {
+            loop {
+                let result = {
+                    let mut client = inner.client.lock().await;
+                    client.lease_keep_alive(lease_id).await
+                };
+
+                match result {
+                    Ok((_keeper, mut stream)) => {
+                        while let Some(resp) = stream.next().await {
+                            if let Err(e) = resp {
+                                tracing::warn!(lease_id = lease_id, error = %e, "keep-alive stream error");
+                                break;
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        tracing::warn!(lease_id = lease_id, error = %e, "keep-alive failed");
+                    }
+                }
+
+                tokio::time::sleep(std::time::Duration::from_secs(interval)).await;
+
+                let re_grant = {
+                    let mut client = inner.client.lock().await;
+                    client
+                        .lease_grant(inner.config.etcd_lease_ttl().unwrap_or(15) as i64, None)
+                        .await
+                };
+
+                if let Ok(current) = re_grant {
+                    let new_lease = current.id();
+                    inner.lease_id.store(new_lease, Ordering::SeqCst);
+
+                    let instance = ServiceInstance {
+                        addr: inner.config.rpc_self_listen_addr().unwrap_or_default(),
+                        metadata: HashMap::new(),
+                    };
+
+                    if let Ok(value) = serde_json::to_string(&instance) {
+                        let mut client = inner.client.lock().await;
+                        let opts = PutOptions::new().with_lease(new_lease);
+                        let _ = client.put(key.clone(), value, Some(opts)).await;
+                    }
+                    tracing::info!(old = lease_id, new = new_lease, "etcd lease renewed");
+                }
+            }
+        });
+    }
+}

etcd/types.rs

@@ -0,0 +1,10 @@
+use std::collections::HashMap;
+
+use serde::{Deserialize, Serialize};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ServiceInstance {
+    pub addr: String,
+    #[serde(default)]
+    pub metadata: HashMap<String, String>,
+}