use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, error, info};
use crate::connector::{ServiceConnector, ServiceInfo};
#[derive(Debug, Clone)]
pub struct CachedSnapshot {
pub services: Vec<ServiceInfo>,
#[allow(dead_code)]
pub refreshed_at: Instant,
}
impl CachedSnapshot {
pub fn url_map(&self) -> HashMap<String, String> {
self.services
.iter()
.filter_map(|s| s.url.as_ref().map(|u| (s.id.clone(), u.clone())))
.collect()
}
}
#[derive(Clone, Debug)]
pub struct PollerCache {
inner: Arc<RwLock<Option<CachedSnapshot>>>,
}
impl Default for PollerCache {
fn default() -> Self {
Self::new()
}
}
impl PollerCache {
pub fn new() -> Self {
Self {
inner: Arc::new(RwLock::new(None)),
}
}
pub async fn snapshot(&self) -> Option<CachedSnapshot> {
self.inner.read().await.clone()
}
pub async fn poll_once(
&self,
connectors: Arc<Vec<Box<dyn ServiceConnector>>>,
) -> CachedSnapshot {
let c = Arc::clone(&connectors);
let services: Vec<ServiceInfo> =
tokio::task::spawn_blocking(move || c.iter().map(|conn| conn.detect()).collect())
.await
.unwrap_or_else(|e| {
error!("poller: detection task panicked: {e}");
vec![]
});
let snap = CachedSnapshot {
services,
refreshed_at: Instant::now(),
};
*self.inner.write().await = Some(snap.clone());
snap
}
}
async fn poll_loop(
cache: &PollerCache,
connectors: &Arc<Vec<Box<dyn ServiceConnector>>>,
interval: Duration,
) {
loop {
let snap = cache.poll_once(Arc::clone(connectors)).await;
let running_count = snap.services.iter().filter(|s| s.url.is_some()).count();
debug!(
"poller: refreshed {} services, {} running",
snap.services.len(),
running_count
);
tokio::time::sleep(interval).await;
}
}
pub fn start(
cache: PollerCache,
connectors: Arc<Vec<Box<dyn ServiceConnector>>>,
interval: Duration,
) {
tokio::spawn(async move {
info!(
"poller: starting background health-poll (interval={}s)",
interval.as_secs()
);
poll_loop(&cache, &connectors, interval).await;
error!("poller: background health-poll loop exited unexpectedly — cache will not refresh");
});
}
#[cfg(test)]
mod tests {
use super::*;
use crate::connector::{ServiceInfo, ServiceStatus};
struct StubConnector {
id: &'static str,
url: Option<&'static str>,
}
impl ServiceConnector for StubConnector {
fn id(&self) -> &'static str {
self.id
}
fn display_name(&self) -> &'static str {
"Stub"
}
fn detect(&self) -> ServiceInfo {
ServiceInfo {
id: self.id.to_string(),
display_name: "Stub".to_string(),
status: if self.url.is_some() {
ServiceStatus::Running
} else {
ServiceStatus::Absent
},
version: None,
url: self.url.map(|u| u.to_string()),
hint: None,
}
}
}
fn make_connectors() -> Arc<Vec<Box<dyn ServiceConnector>>> {
Arc::new(vec![
Box::new(StubConnector {
id: "trusty-search",
url: Some("http://127.0.0.1:7878"),
}),
Box::new(StubConnector {
id: "trusty-memory",
url: None,
}),
])
}
#[tokio::test]
async fn test_cache_initialises_with_connectors() {
let cache = PollerCache::new();
let connectors = make_connectors();
assert!(cache.snapshot().await.is_none(), "should start empty");
let snap = cache.poll_once(Arc::clone(&connectors)).await;
assert_eq!(snap.services.len(), 2);
assert_eq!(snap.services[0].id, "trusty-search");
assert_eq!(snap.services[1].id, "trusty-memory");
let cached = cache.snapshot().await.expect("snapshot after poll");
assert_eq!(cached.services.len(), 2);
}
#[test]
fn test_snapshot_url_map() {
let snap = CachedSnapshot {
services: vec![
ServiceInfo {
id: "trusty-search".to_string(),
display_name: "Search".to_string(),
status: ServiceStatus::Running,
version: Some("1.0.0".to_string()),
url: Some("http://127.0.0.1:7878".to_string()),
hint: None,
},
ServiceInfo {
id: "trusty-memory".to_string(),
display_name: "Memory".to_string(),
status: ServiceStatus::Absent,
version: None,
url: None,
hint: None,
},
],
refreshed_at: Instant::now(),
};
let map = snap.url_map();
assert_eq!(map.len(), 1);
assert_eq!(map["trusty-search"], "http://127.0.0.1:7878");
assert!(!map.contains_key("trusty-memory"));
}
}