zhenyi
89 lines
2.5 KiB
Rust
89 lines
2.5 KiB
Rust
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()
|
|
}
|
|
}
|