use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use aion_store::NamespaceStore;
#[derive(Clone, Copy)]
struct CachedQuota {
ceiling: u32,
fetched_at: Instant,
}
#[derive(Clone)]
pub struct QuotaCache {
store: Arc<dyn NamespaceStore>,
platform_default: u32,
ttl: Duration,
entries: Arc<Mutex<HashMap<String, CachedQuota>>>,
}
impl std::fmt::Debug for QuotaCache {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("QuotaCache")
.field("platform_default", &self.platform_default)
.field("ttl", &self.ttl)
.finish_non_exhaustive()
}
}
impl QuotaCache {
#[must_use]
pub fn new(store: Arc<dyn NamespaceStore>, platform_default: u32, ttl: Duration) -> Self {
Self {
store,
platform_default,
ttl,
entries: Arc::new(Mutex::new(HashMap::new())),
}
}
pub async fn ceiling(&self, namespace: &str) -> u32 {
if let Some(hit) = self.fresh_hit(namespace) {
return hit;
}
let ceiling = match self.store.get_namespace(namespace).await {
Ok(Some(record)) => record
.config
.max_in_flight_activities
.unwrap_or(self.platform_default),
Ok(None) | Err(_) => self.platform_default,
};
self.store_entry(namespace, ceiling);
ceiling
}
fn fresh_hit(&self, namespace: &str) -> Option<u32> {
let entries = self.entries.lock().ok()?;
let entry = entries.get(namespace)?;
if entry.fetched_at.elapsed() < self.ttl {
Some(entry.ceiling)
} else {
None
}
}
fn store_entry(&self, namespace: &str, ceiling: u32) {
if let Ok(mut entries) = self.entries.lock() {
entries.insert(
namespace.to_owned(),
CachedQuota {
ceiling,
fetched_at: Instant::now(),
},
);
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::expect_used)]
use std::sync::Arc;
use std::time::Duration;
use aion_store::{InMemoryStore, NamespaceOrigin, NamespaceRecord, NamespaceStore};
use chrono::Utc;
use super::QuotaCache;
async fn register_with_quota(
store: &Arc<dyn NamespaceStore>,
namespace: &str,
quota: Option<u32>,
) -> Result<(), Box<dyn std::error::Error>> {
let mut record =
NamespaceRecord::new_minted(namespace, NamespaceOrigin::Explicit, Utc::now());
record.config.max_in_flight_activities = quota;
store.put_namespace(record).await?;
Ok(())
}
#[tokio::test]
async fn explicit_override_is_read_from_store_and_cached()
-> Result<(), Box<dyn std::error::Error>> {
let store: Arc<dyn NamespaceStore> = Arc::new(InMemoryStore::default());
register_with_quota(&store, "capped", Some(7)).await?;
let cache = QuotaCache::new(Arc::clone(&store), 1024, Duration::from_secs(60));
assert_eq!(
cache.ceiling("capped").await,
7,
"the explicit override is read"
);
register_with_quota(&store, "capped", Some(99)).await?;
assert_eq!(
cache.ceiling("capped").await,
7,
"a fresh cache hit must not re-read the mutated durable record"
);
Ok(())
}
#[tokio::test]
async fn unset_override_resolves_to_platform_default() -> Result<(), Box<dyn std::error::Error>>
{
let store: Arc<dyn NamespaceStore> = Arc::new(InMemoryStore::default());
register_with_quota(&store, "uncapped", None).await?;
let cache = QuotaCache::new(store, 1024, Duration::from_secs(60));
assert_eq!(
cache.ceiling("uncapped").await,
1024,
"a namespace with no override resolves to the generous platform default"
);
Ok(())
}
#[tokio::test]
async fn absent_namespace_resolves_to_platform_default()
-> Result<(), Box<dyn std::error::Error>> {
let store: Arc<dyn NamespaceStore> = Arc::new(InMemoryStore::default());
let cache = QuotaCache::new(store, 512, Duration::from_secs(60));
assert_eq!(
cache.ceiling("never-seen").await,
512,
"an absent registry row resolves to the platform default, never zero"
);
Ok(())
}
#[tokio::test]
async fn refreshes_after_ttl_expiry() -> Result<(), Box<dyn std::error::Error>> {
let store: Arc<dyn NamespaceStore> = Arc::new(InMemoryStore::default());
let cache = QuotaCache::new(Arc::clone(&store), 1024, Duration::ZERO);
assert_eq!(
cache.ceiling("capped").await,
1024,
"an absent record resolves to the platform default"
);
register_with_quota(&store, "capped", Some(7)).await?;
assert_eq!(
cache.ceiling("capped").await,
7,
"an expired entry re-reads the durable record and sees the new override"
);
Ok(())
}
}