use std::future::Future;
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;
use chrono::{DateTime, Utc};
use axess_cache::ClockTtlCache;
use axess_clock::{Clock, SystemClock};
use crate::authn::ids::{DeviceId, TenantId, UserId};
use crate::device::store::{DeviceStore, SweepCounts};
use crate::device::types::{Device, DeviceTrustLevel, FingerprintHash};
const DEFAULT_CAPACITY: usize = 10_000;
const DEFAULT_TTL_SECS: u64 = 60;
type CacheKey = (TenantId, DeviceId);
pub struct CachedDeviceStore<S>
where
S: DeviceStore,
{
inner: S,
cache: Arc<ClockTtlCache<CacheKey, Device>>,
capacity: NonZeroUsize,
ttl: Duration,
clock: Arc<dyn Clock>,
}
impl<S> Clone for CachedDeviceStore<S>
where
S: DeviceStore,
{
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
cache: self.cache.clone(),
capacity: self.capacity,
ttl: self.ttl,
clock: self.clock.clone(),
}
}
}
impl<S> CachedDeviceStore<S>
where
S: DeviceStore,
{
pub fn new(inner: S) -> Self {
Self::with_options(
inner,
DEFAULT_CAPACITY,
Duration::from_secs(DEFAULT_TTL_SECS),
Arc::new(SystemClock),
)
}
pub fn with_options(inner: S, capacity: usize, ttl: Duration, clock: Arc<dyn Clock>) -> Self {
let capacity = NonZeroUsize::new(capacity.max(1)).expect("capacity >= 1");
let cache = Arc::new(ClockTtlCache::new(capacity, ttl, clock.clone()));
Self {
inner,
cache,
capacity,
ttl,
clock,
}
}
fn rebuilt(mut self) -> Self {
self.cache = Arc::new(ClockTtlCache::new(
self.capacity,
self.ttl,
self.clock.clone(),
));
self
}
pub fn with_capacity(mut self, capacity: usize) -> Self {
self.capacity = NonZeroUsize::new(capacity.max(1)).expect("capacity >= 1");
self.rebuilt()
}
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.ttl = ttl;
self.rebuilt()
}
pub fn with_clock(mut self, clock: Arc<dyn Clock>) -> Self {
self.clock = clock;
self.rebuilt()
}
pub fn stats(&self) -> axess_cache::CacheStats {
self.cache.stats()
}
pub fn invalidate_all(&self) {
self.cache.invalidate_all();
}
pub fn invalidate_tenant(&self, tenant_id: &TenantId) {
let target = *tenant_id;
self.cache.invalidate_by(|k| k.0 == target);
}
}
impl<S> DeviceStore for CachedDeviceStore<S>
where
S: DeviceStore,
{
type Error = S::Error;
fn load(
&self,
tenant_id: &TenantId,
id: &DeviceId,
) -> impl Future<Output = Result<Option<Device>, Self::Error>> + Send {
let key = (*tenant_id, *id);
let cache = self.cache.clone();
let inner = self.inner.clone();
let tenant = *tenant_id;
let device = *id;
async move {
let loaded = cache
.get_or_try_insert_with(key, || async {
match inner.load(&tenant, &device).await {
Ok(Some(d)) => Ok(d),
Ok(None) => Err(None),
Err(e) => Err(Some(e)),
}
})
.await;
match loaded {
Ok(d) => Ok(Some(d)),
Err(None) => Ok(None),
Err(Some(e)) => Err(e),
}
}
}
fn find_by_fingerprint(
&self,
tenant_id: &TenantId,
hash: &FingerprintHash,
) -> impl Future<Output = Result<Option<Device>, Self::Error>> + Send {
self.inner.find_by_fingerprint(tenant_id, hash)
}
fn find_for_user(
&self,
tenant_id: &TenantId,
user_id: &UserId,
limit: usize,
) -> impl Future<Output = Result<Vec<Device>, Self::Error>> + Send {
self.inner.find_for_user(tenant_id, user_id, limit)
}
fn find_by_refresh_family(
&self,
tenant_id: &TenantId,
family_id: &str,
) -> impl Future<Output = Result<Vec<Device>, Self::Error>> + Send {
self.inner.find_by_refresh_family(tenant_id, family_id)
}
fn save(&self, device: &Device) -> impl Future<Output = Result<(), Self::Error>> + Send {
let key = (device.tenant_id, device.id);
let cache = self.cache.clone();
let inner = self.inner.clone();
let device = device.clone();
async move {
let result = inner.save(&device).await;
cache.invalidate(&key);
result
}
}
fn record_sighting(
&self,
tenant_id: &TenantId,
id: &DeviceId,
now: DateTime<Utc>,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
self.inner.record_sighting(tenant_id, id, now)
}
fn set_trust_level(
&self,
tenant_id: &TenantId,
id: &DeviceId,
level: DeviceTrustLevel,
now: DateTime<Utc>,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
let key = (*tenant_id, *id);
let cache = self.cache.clone();
let inner = self.inner.clone();
let tenant = *tenant_id;
let device = *id;
async move {
let result = inner.set_trust_level(&tenant, &device, level, now).await;
cache.invalidate(&key);
result
}
}
fn delete(
&self,
tenant_id: &TenantId,
id: &DeviceId,
) -> impl Future<Output = Result<(), Self::Error>> + Send {
let key = (*tenant_id, *id);
let cache = self.cache.clone();
let inner = self.inner.clone();
let tenant = *tenant_id;
let device = *id;
async move {
let result = inner.delete(&tenant, &device).await;
cache.invalidate(&key);
result
}
}
fn sweep(
&self,
tenant_id: &TenantId,
now: DateTime<Utc>,
) -> impl Future<Output = Result<SweepCounts, Self::Error>> + Send {
let cache = self.cache.clone();
let inner = self.inner.clone();
let tenant = *tenant_id;
async move {
let result = inner.sweep(&tenant, now).await;
cache.invalidate_by(|k| k.0 == tenant);
result
}
}
}
#[cfg(test)]
mod tests;
#[cfg(test)]
mod races;