use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use async_trait::async_trait;
#[async_trait]
pub trait ModuleService: Send + Sync {
fn service_id(&self) -> &'static str;
fn service_type(&self) -> &'static str;
fn module_name(&self) -> &'static str;
async fn health_check(&self) -> Result<ServiceHealth, String>;
fn as_any(&self) -> Option<&dyn Any> {
None
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct ServiceHealth {
pub status: HealthStatus,
pub message: Option<String>,
pub last_check: chrono::DateTime<chrono::Utc>,
pub details: HashMap<String, String>,
}
impl ServiceHealth {
pub fn healthy() -> Self {
Self {
status: HealthStatus::Healthy,
message: None,
last_check: chrono::Utc::now(),
details: HashMap::new(),
}
}
pub fn unhealthy(message: impl Into<String>) -> Self {
Self {
status: HealthStatus::Unhealthy,
message: Some(message.into()),
last_check: chrono::Utc::now(),
details: HashMap::new(),
}
}
pub fn degraded(message: impl Into<String>) -> Self {
Self {
status: HealthStatus::Degraded,
message: Some(message.into()),
last_check: chrono::Utc::now(),
details: HashMap::new(),
}
}
pub fn with_detail(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.details.insert(key.into(), value.into());
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "lowercase")]
pub enum HealthStatus {
Healthy,
Degraded,
Unhealthy,
}
#[derive(Clone)]
pub struct ServiceRegistry {
services: Arc<RwLock<HashMap<String, Arc<dyn ModuleService>>>>,
}
impl ServiceRegistry {
pub fn new() -> Self {
Self {
services: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn register<S: ModuleService + 'static>(&self, service: S) {
let service_id = service.service_id().to_string();
tracing::info!(
service_id = %service_id,
module = %service.module_name(),
service_type = %service.service_type(),
"Registering service"
);
self.services.write().await.insert(service_id, Arc::new(service));
}
pub async fn register_arc(&self, service: Arc<dyn ModuleService>) {
let service_id = service.service_id().to_string();
tracing::info!(
service_id = %service_id,
module = %service.module_name(),
service_type = %service.service_type(),
"Registering service (Arc)"
);
self.services.write().await.insert(service_id, service);
}
pub async fn get(&self, service_id: &str) -> Option<Arc<dyn ModuleService>> {
self.services.read().await.get(service_id).cloned()
}
pub async fn unregister(&self, service_id: &str) -> bool {
self.services.write().await.remove(service_id).is_some()
}
pub async fn list(&self) -> Vec<String> {
self.services.read().await.keys().cloned().collect()
}
pub async fn list_by_module(&self, module_name: &str) -> Vec<Arc<dyn ModuleService>> {
self.services
.read()
.await
.values()
.filter(|s| s.module_name() == module_name)
.cloned()
.collect()
}
pub async fn list_by_type(&self, service_type: &str) -> Vec<Arc<dyn ModuleService>> {
self.services
.read()
.await
.values()
.filter(|s| s.service_type() == service_type)
.cloned()
.collect()
}
pub async fn has(&self, service_id: &str) -> bool {
self.services.read().await.contains_key(service_id)
}
pub async fn count(&self) -> usize {
self.services.read().await.len()
}
pub async fn health_check_all(&self) -> RegistryHealth {
let services = self.services.read().await;
let mut results = HashMap::new();
let mut healthy_count = 0;
let mut unhealthy_count = 0;
for (id, service) in services.iter() {
match service.health_check().await {
Ok(health) => {
if health.status == HealthStatus::Healthy {
healthy_count += 1;
} else {
unhealthy_count += 1;
}
results.insert(id.clone(), Ok(health));
}
Err(e) => {
unhealthy_count += 1;
results.insert(id.clone(), Err(e));
}
}
}
let overall_status = if unhealthy_count == 0 {
HealthStatus::Healthy
} else if healthy_count > 0 {
HealthStatus::Degraded
} else {
HealthStatus::Unhealthy
};
RegistryHealth {
status: overall_status,
total_services: services.len(),
healthy_count,
unhealthy_count,
services: results,
checked_at: chrono::Utc::now(),
}
}
pub async fn statistics(&self) -> RegistryStatistics {
let services = self.services.read().await;
let mut by_module: HashMap<String, usize> = HashMap::new();
let mut by_type: HashMap<String, usize> = HashMap::new();
for service in services.values() {
*by_module.entry(service.module_name().to_string()).or_default() += 1;
*by_type.entry(service.service_type().to_string()).or_default() += 1;
}
RegistryStatistics {
total_services: services.len(),
services_by_module: by_module,
services_by_type: by_type,
}
}
}
impl Default for ServiceRegistry {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct RegistryHealth {
pub status: HealthStatus,
pub total_services: usize,
pub healthy_count: usize,
pub unhealthy_count: usize,
#[serde(skip_serializing_if = "HashMap::is_empty")]
pub services: HashMap<String, Result<ServiceHealth, String>>,
pub checked_at: chrono::DateTime<chrono::Utc>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct RegistryStatistics {
pub total_services: usize,
pub services_by_module: HashMap<String, usize>,
pub services_by_type: HashMap<String, usize>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct ServiceDescriptor {
pub service_id: String,
pub service_type: String,
pub module_name: String,
pub description: Option<String>,
pub version: Option<String>,
}
impl ServiceDescriptor {
pub fn from_service(service: &dyn ModuleService) -> Self {
Self {
service_id: service.service_id().to_string(),
service_type: service.service_type().to_string(),
module_name: service.module_name().to_string(),
description: None,
version: None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
struct TestService {
id: &'static str,
healthy: bool,
}
#[async_trait]
impl ModuleService for TestService {
fn service_id(&self) -> &'static str {
self.id
}
fn service_type(&self) -> &'static str {
"test"
}
fn module_name(&self) -> &'static str {
"test_module"
}
async fn health_check(&self) -> Result<ServiceHealth, String> {
if self.healthy {
Ok(ServiceHealth::healthy())
} else {
Err("Service unhealthy".to_string())
}
}
}
#[tokio::test]
async fn test_register_and_get() {
let registry = ServiceRegistry::new();
let service = TestService {
id: "test.service",
healthy: true,
};
registry.register(service).await;
assert!(registry.has("test.service").await);
assert!(!registry.has("nonexistent").await);
let retrieved = registry.get("test.service").await;
assert!(retrieved.is_some());
assert_eq!(retrieved.unwrap().service_id(), "test.service");
}
#[tokio::test]
async fn test_list_services() {
let registry = ServiceRegistry::new();
registry.register(TestService { id: "test.a", healthy: true }).await;
registry.register(TestService { id: "test.b", healthy: true }).await;
let list = registry.list().await;
assert_eq!(list.len(), 2);
assert!(list.contains(&"test.a".to_string()));
assert!(list.contains(&"test.b".to_string()));
}
#[tokio::test]
async fn test_unregister() {
let registry = ServiceRegistry::new();
registry.register(TestService { id: "test.service", healthy: true }).await;
assert!(registry.has("test.service").await);
let removed = registry.unregister("test.service").await;
assert!(removed);
assert!(!registry.has("test.service").await);
}
#[tokio::test]
async fn test_health_check_all() {
let registry = ServiceRegistry::new();
registry.register(TestService { id: "test.healthy", healthy: true }).await;
registry.register(TestService { id: "test.unhealthy", healthy: false }).await;
let health = registry.health_check_all().await;
assert_eq!(health.total_services, 2);
assert_eq!(health.healthy_count, 1);
assert_eq!(health.unhealthy_count, 1);
assert_eq!(health.status, HealthStatus::Degraded);
}
#[tokio::test]
async fn test_statistics() {
let registry = ServiceRegistry::new();
registry.register(TestService { id: "test.a", healthy: true }).await;
registry.register(TestService { id: "test.b", healthy: true }).await;
let stats = registry.statistics().await;
assert_eq!(stats.total_services, 2);
assert_eq!(stats.services_by_module.get("test_module"), Some(&2));
assert_eq!(stats.services_by_type.get("test"), Some(&2));
}
#[tokio::test]
async fn test_list_by_module() {
let registry = ServiceRegistry::new();
registry.register(TestService { id: "test.a", healthy: true }).await;
let by_module = registry.list_by_module("test_module").await;
assert_eq!(by_module.len(), 1);
let by_other = registry.list_by_module("other_module").await;
assert!(by_other.is_empty());
}
#[test]
fn test_service_health_builders() {
let healthy = ServiceHealth::healthy();
assert_eq!(healthy.status, HealthStatus::Healthy);
assert!(healthy.message.is_none());
let unhealthy = ServiceHealth::unhealthy("Something went wrong");
assert_eq!(unhealthy.status, HealthStatus::Unhealthy);
assert_eq!(unhealthy.message, Some("Something went wrong".to_string()));
let degraded = ServiceHealth::degraded("Slow response")
.with_detail("latency_ms", "500");
assert_eq!(degraded.status, HealthStatus::Degraded);
assert_eq!(degraded.details.get("latency_ms"), Some(&"500".to_string()));
}
}