use crate::core::balancer::LoadBalancer;
use crate::core::health::HealthMonitor;
use crate::protocols::ProtocolRegistry;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
#[derive(Debug, Clone)]
pub enum RouteResult {
Matched(String, String), Fallback(String), Discarded,
}
pub struct Router {
registry: Arc<ProtocolRegistry>,
balancer: Arc<RwLock<LoadBalancer>>,
}
impl Router {
#[must_use]
pub const fn new(registry: Arc<ProtocolRegistry>, balancer: Arc<RwLock<LoadBalancer>>) -> Self {
Self { registry, balancer }
}
pub async fn route(&self, data: &[u8]) -> Option<RouteResult> {
let balancer_guard = self.balancer.read().await;
if let Some(m) = self.registry.probe(data) {
let metadata = m.metadata.as_deref();
if let Some(addr) = balancer_guard.next_available(&m.name, metadata).await {
Some(RouteResult::Matched(m.name, addr.clone()))
} else if let Some(addr) = balancer_guard.next_available("fallback", None).await {
Some(RouteResult::Fallback(addr.clone()))
} else {
Some(RouteResult::Discarded)
}
} else {
None
}
}
pub async fn route_fallback(&self) -> RouteResult {
let balancer_guard = self.balancer.read().await;
balancer_guard
.next_available("fallback", None)
.await
.map_or(RouteResult::Discarded, |addr| {
RouteResult::Fallback(addr.clone())
})
}
pub async fn update_balancer(
&self,
routes: HashMap<String, Vec<String>>,
health: Arc<HealthMonitor>,
) {
let mut balancer_guard = self.balancer.write().await;
*balancer_guard = LoadBalancer::new(routes, health);
}
pub async fn get_health_status(&self) -> HashMap<String, Vec<(String, bool)>> {
let balancer_guard = self.balancer.read().await;
balancer_guard.get_status().await
}
}