use async_trait::async_trait;
use dashmap::DashMap;
use reqwest::Client;
use revoke_core::{HealthStatus, Result, RevokeError, ServiceInfo, ServiceRegistry, Status};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::interval;
use tracing::{debug, error, info};
use uuid::Uuid;
const DEFAULT_CHECK_INTERVAL: Duration = Duration::from_secs(10);
const DEFAULT_CACHE_REFRESH_INTERVAL: Duration = Duration::from_secs(5);
const DEFAULT_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);
const DEFAULT_DEREGISTER_CRITICAL: &str = "30s";
const DEFAULT_TTL: Duration = Duration::from_secs(30);
#[derive(Debug, Clone)]
pub struct ConsulConfig {
pub address: String,
pub check_interval: Duration,
pub cache_refresh_interval: Duration,
pub client_timeout: Duration,
pub deregister_critical: String,
pub ttl: Duration,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ConsulService {
#[serde(rename = "ID")]
pub id: String,
#[serde(rename = "Name")]
pub name: String,
#[serde(rename = "Tags")]
pub tags: Option<Vec<String>>,
#[serde(rename = "Address")]
pub address: String,
#[serde(rename = "Port")]
pub port: u16,
#[serde(rename = "Meta")]
pub meta: Option<HashMap<String, String>>,
#[serde(rename = "EnableTagOverride")]
pub enable_tag_override: Option<bool>,
#[serde(rename = "Check")]
pub check: Option<ConsulCheck>,
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ConsulCheck {
#[serde(rename = "DeregisterCriticalServiceAfter")]
pub deregister_critical_service_after: Option<String>,
#[serde(rename = "TTL")]
pub ttl: Option<String>,
#[serde(rename = "HTTP")]
pub http: Option<String>,
#[serde(rename = "Interval")]
pub interval: Option<String>,
#[serde(rename = "Timeout")]
pub timeout: Option<String>,
}
#[derive(Debug, Deserialize)]
pub struct ConsulHealthService {
#[serde(rename = "Node")]
pub node: ConsulNode,
#[serde(rename = "Service")]
pub service: ConsulServiceInfo,
#[serde(rename = "Checks")]
pub checks: Vec<ConsulHealthCheck>,
}
#[derive(Debug, Deserialize)]
pub struct ConsulNode {
#[serde(rename = "ID")]
pub id: String,
#[serde(rename = "Node")]
pub node: String,
#[serde(rename = "Address")]
pub address: String,
}
#[derive(Debug, Deserialize)]
pub struct ConsulServiceInfo {
#[serde(rename = "ID")]
pub id: String,
#[serde(rename = "Service")]
pub service: String,
#[serde(rename = "Tags")]
pub tags: Option<Vec<String>>,
#[serde(rename = "Address")]
pub address: String,
#[serde(rename = "Port")]
pub port: u16,
#[serde(rename = "Meta")]
pub meta: Option<HashMap<String, String>>,
}
#[derive(Debug, Deserialize)]
pub struct ConsulHealthCheck {
#[serde(rename = "CheckID")]
pub check_id: String,
#[serde(rename = "Status")]
pub status: String,
#[serde(rename = "Output")]
pub output: Option<String>,
}
impl Default for ConsulConfig {
fn default() -> Self {
Self {
address: "localhost:8500".to_string(),
check_interval: DEFAULT_CHECK_INTERVAL,
cache_refresh_interval: DEFAULT_CACHE_REFRESH_INTERVAL,
client_timeout: DEFAULT_CLIENT_TIMEOUT,
deregister_critical: DEFAULT_DEREGISTER_CRITICAL.to_string(),
ttl: DEFAULT_TTL,
}
}
}
impl ConsulConfig {
pub fn new(address: impl Into<String>) -> Self {
Self {
address: address.into(),
..Default::default()
}
}
pub fn with_check_interval(mut self, interval: Duration) -> Self {
self.check_interval = interval;
self
}
pub fn with_cache_refresh_interval(mut self, interval: Duration) -> Self {
self.cache_refresh_interval = interval;
self
}
pub fn with_client_timeout(mut self, timeout: Duration) -> Self {
self.client_timeout = timeout;
self
}
pub fn with_deregister_critical(mut self, critical: impl Into<String>) -> Self {
self.deregister_critical = critical.into();
self
}
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.ttl = ttl;
self
}
}
pub struct ConsulRegistry {
client: Client,
base_url: String,
local_cache: Arc<DashMap<String, Vec<ServiceInfo>>>,
config: ConsulConfig,
}
impl ConsulRegistry {
pub async fn new(config: ConsulConfig) -> Result<Self> {
let base_url = format!("http://{}/v1", config.address);
let client = Client::builder()
.timeout(config.client_timeout)
.build()
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?;
let registry = Self {
client,
base_url,
local_cache: Arc::new(DashMap::new()),
config,
};
registry.start_cache_refresh().await;
Ok(registry)
}
pub async fn from_address(consul_addr: &str) -> Result<Self> {
Self::new(ConsulConfig::new(consul_addr)).await
}
fn convert_to_consul_service(&self, service: &ServiceInfo) -> ConsulService {
ConsulService {
id: service.id.to_string(),
name: service.name.clone(),
tags: Some(vec![
format!("version={}", service.version),
format!("protocol={:?}", service.protocol),
]),
address: service.address.clone(),
port: service.port,
meta: if service.metadata.is_empty() {
None
} else {
Some(service.metadata.clone())
},
enable_tag_override: Some(false),
check: Some(ConsulCheck {
deregister_critical_service_after: Some(self.config.deregister_critical.clone()),
ttl: Some(format!("{}s", self.config.ttl.as_secs())),
http: None,
interval: None,
timeout: None,
}),
}
}
fn convert_from_consul_service(service: ConsulServiceInfo) -> Option<ServiceInfo> {
let id = Uuid::parse_str(&service.id).ok()?;
let metadata = service.meta.unwrap_or_default();
let version = service
.tags
.as_ref()
.and_then(|tags| {
tags.iter()
.find(|t| t.starts_with("version="))
.map(|t| t.trim_start_matches("version=").to_string())
})
.unwrap_or_else(|| "1.0.0".to_string());
let protocol = service
.tags
.as_ref()
.and_then(|tags| {
tags.iter()
.find(|t| t.starts_with("protocol="))
.and_then(|t| match t.trim_start_matches("protocol=") {
"Http" => Some(revoke_core::Protocol::Http),
"Https" => Some(revoke_core::Protocol::Https),
"Grpc" => Some(revoke_core::Protocol::Grpc),
"Tcp" => Some(revoke_core::Protocol::Tcp),
_ => None,
})
})
.unwrap_or(revoke_core::Protocol::Http);
Some(ServiceInfo {
id,
name: service.service,
version,
address: service.address,
port: service.port,
protocol,
metadata,
})
}
async fn refresh_cache(&self) -> Result<()> {
let url = format!("{}/catalog/services", self.base_url);
let services: HashMap<String, Vec<String>> = self
.client
.get(&url)
.send()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?
.json()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?;
self.local_cache.clear();
for (service_name, _) in services {
let health_url = format!(
"{}/health/service/{}?passing=true",
self.base_url, service_name
);
match self.client.get(&health_url).send().await {
Ok(response) => {
if let Ok(health_services) = response.json::<Vec<ConsulHealthService>>().await {
let service_infos: Vec<ServiceInfo> = health_services
.into_iter()
.filter_map(|hs| Self::convert_from_consul_service(hs.service))
.collect();
if !service_infos.is_empty() {
self.local_cache.insert(service_name, service_infos);
}
}
}
Err(e) => {
error!("Failed to get health for service {}: {:?}", service_name, e);
}
}
}
Ok(())
}
pub async fn start_cache_refresh(&self) {
let cache = self.local_cache.clone();
let client = self.client.clone();
let base_url = self.base_url.clone();
let refresh_interval = self.config.cache_refresh_interval;
tokio::spawn(async move {
let mut ticker = interval(refresh_interval);
loop {
ticker.tick().await;
let url = format!("{}/catalog/services", base_url);
let services: HashMap<String, Vec<String>> = match client.get(&url).send().await {
Ok(resp) => match resp.json().await {
Ok(s) => s,
Err(e) => {
error!("Failed to parse services: {:?}", e);
continue;
}
},
Err(e) => {
error!("Failed to get services from Consul: {:?}", e);
continue;
}
};
let mut new_cache: HashMap<String, Vec<ServiceInfo>> = HashMap::new();
for (service_name, _) in services {
let health_url =
format!("{}/health/service/{}?passing=true", base_url, service_name);
match client.get(&health_url).send().await {
Ok(response) => {
if let Ok(health_services) =
response.json::<Vec<ConsulHealthService>>().await
{
let service_infos: Vec<ServiceInfo> = health_services
.into_iter()
.filter_map(|hs| Self::convert_from_consul_service(hs.service))
.collect();
if !service_infos.is_empty() {
new_cache.insert(service_name, service_infos);
}
}
}
Err(e) => {
error!("Failed to get service {} health: {:?}", service_name, e);
}
}
}
cache.clear();
for (name, services) in new_cache {
cache.insert(name, services);
}
}
});
}
async fn keep_alive(&self, service_id: String) {
let client = self.client.clone();
let base_url = self.base_url.clone();
let check_id = format!("service:{}", service_id);
let check_interval = self.config.check_interval;
tokio::spawn(async move {
let mut ticker = interval(check_interval);
loop {
ticker.tick().await;
let url = format!("{}/agent/check/pass/{}", base_url, check_id);
if let Err(e) = client.put(&url).send().await {
error!("Failed to update check TTL for {}: {:?}", service_id, e);
}
}
});
}
}
#[async_trait]
impl ServiceRegistry for ConsulRegistry {
async fn register(&self, service: ServiceInfo) -> Result<()> {
let consul_service = self.convert_to_consul_service(&service);
let service_id = service.id.to_string();
let url = format!("{}/agent/service/register", self.base_url);
self.client
.put(&url)
.json(&consul_service)
.send()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?
.error_for_status()
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?;
info!("Service {} registered with ID {}", service.name, service.id);
self.keep_alive(service_id).await;
self.refresh_cache().await?;
Ok(())
}
async fn deregister(&self, service_id: Uuid) -> Result<()> {
let id = service_id.to_string();
let url = format!("{}/agent/service/deregister/{}", self.base_url, id);
self.client
.put(&url)
.send()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?
.error_for_status()
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?;
info!("Service {} deregistered", service_id);
self.refresh_cache().await?;
Ok(())
}
async fn get_service(&self, name: &str) -> Result<Vec<ServiceInfo>> {
if let Some(services) = self.local_cache.get(name) {
Ok(services.clone())
} else {
let url = format!("{}/health/service/{}?passing=true", self.base_url, name);
let health_services: Vec<ConsulHealthService> = self
.client
.get(&url)
.send()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?
.json()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?;
let service_infos: Vec<ServiceInfo> = health_services
.into_iter()
.filter_map(|hs| Self::convert_from_consul_service(hs.service))
.collect();
if !service_infos.is_empty() {
self.local_cache
.insert(name.to_string(), service_infos.clone());
}
Ok(service_infos)
}
}
async fn update_health(&self, status: HealthStatus) -> Result<()> {
let check_id = format!("service:{}", status.service_id);
let endpoint = match status.status {
Status::Healthy => "pass",
Status::Unhealthy => "fail",
Status::Unknown => "warn",
};
let url = format!("{}/agent/check/{}/{}", self.base_url, endpoint, check_id);
let mut req = self.client.put(&url);
if let Some(message) = &status.message {
req = req.body(message.clone());
}
req.send()
.await
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?
.error_for_status()
.map_err(|e| RevokeError::ConnectionError(e.to_string()))?;
debug!("Health status updated for service {}", status.service_id);
Ok(())
}
}