94 lines
2.8 KiB
Rust
94 lines
2.8 KiB
Rust
use std::collections::{HashMap, HashSet};
|
|
use std::sync::Arc;
|
|
|
|
use chrono::{DateTime, Utc};
|
|
use harmony_reconciler_contracts::DeviceGroupSource;
|
|
use tokio::sync::RwLock;
|
|
|
|
pub type DeviceGroups = HashMap<String, HashSet<String>>;
|
|
|
|
#[derive(Default)]
|
|
struct CachedGroups {
|
|
groups: Option<Arc<DeviceGroups>>,
|
|
updated_at: Option<DateTime<Utc>>,
|
|
}
|
|
|
|
/// Failed refreshes retain the last successful snapshot; before the first
|
|
/// success there is no snapshot, so startup remains fail-closed.
|
|
pub struct DeviceGroupCache {
|
|
source: Arc<dyn DeviceGroupSource>,
|
|
cached: RwLock<CachedGroups>,
|
|
}
|
|
|
|
impl DeviceGroupCache {
|
|
pub fn new(source: Arc<dyn DeviceGroupSource>) -> Self {
|
|
Self {
|
|
source,
|
|
cached: RwLock::new(CachedGroups::default()),
|
|
}
|
|
}
|
|
|
|
pub async fn refresh(&self) -> Option<Arc<DeviceGroups>> {
|
|
match self.source.device_groups().await {
|
|
Ok(groups) => {
|
|
let groups = Arc::new(groups);
|
|
let mut cached = self.cached.write().await;
|
|
cached.groups = Some(groups.clone());
|
|
cached.updated_at = Some(Utc::now());
|
|
Some(groups)
|
|
}
|
|
Err(error) => {
|
|
let cached = self.cached.read().await;
|
|
let age_seconds = cached
|
|
.updated_at
|
|
.map(|updated_at| Utc::now().signed_duration_since(updated_at).num_seconds());
|
|
tracing::warn!(
|
|
%error,
|
|
?age_seconds,
|
|
"device group refresh failed; retaining last known state"
|
|
);
|
|
cached.groups.clone()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub async fn snapshot(&self) -> Option<Arc<DeviceGroups>> {
|
|
self.cached.read().await.groups.clone()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
|
|
use async_trait::async_trait;
|
|
use harmony_reconciler_contracts::GroupSourceError;
|
|
|
|
use super::*;
|
|
|
|
struct OneSuccess(AtomicUsize);
|
|
|
|
#[async_trait]
|
|
impl DeviceGroupSource for OneSuccess {
|
|
async fn device_groups(&self) -> Result<DeviceGroups, GroupSourceError> {
|
|
if self.0.fetch_add(1, Ordering::Relaxed) == 0 {
|
|
Ok(HashMap::from([(
|
|
"device-1".into(),
|
|
HashSet::from(["edge-a".into()]),
|
|
)]))
|
|
} else {
|
|
Err(GroupSourceError::Source("unavailable".into()))
|
|
}
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn failed_refresh_retains_last_snapshot() {
|
|
let cache = DeviceGroupCache::new(Arc::new(OneSuccess(AtomicUsize::new(0))));
|
|
let first = cache.refresh().await.unwrap();
|
|
let stale = cache.refresh().await.unwrap();
|
|
|
|
assert!(Arc::ptr_eq(&first, &stale));
|
|
}
|
|
}
|