use bytes::Bytes;
use moka::sync::Cache;
use moka::Expiry;
use std::time::{Duration, Instant};
#[derive(Clone)]
struct L1Entry {
data: Bytes,
ttl: Duration,
created_at: Instant,
freshness_jitter: f64,
}
struct L1Expiry;
impl Expiry<String, L1Entry> for L1Expiry {
fn expire_after_create(
&self,
_key: &String,
value: &L1Entry,
_created_at: std::time::Instant,
) -> Option<Duration> {
Some(value.ttl)
}
fn expire_after_update(
&self,
_key: &String,
value: &L1Entry,
_updated_at: std::time::Instant,
_duration_until_expiry: Option<Duration>,
) -> Option<Duration> {
Some(value.ttl)
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum L1SwrRead {
Fresh(Vec<u8>),
Stale(Vec<u8>),
Miss,
}
pub(crate) enum SharedSwrRead {
Fresh(Bytes),
Stale(Bytes),
Miss,
}
#[derive(Clone)]
pub struct L1Cache {
store: Cache<String, L1Entry>,
}
impl L1Cache {
pub fn new(capacity: usize) -> Self {
Self {
store: Cache::builder()
.max_capacity(u64::try_from(capacity).unwrap_or(u64::MAX))
.expire_after(L1Expiry)
.build(),
}
}
pub fn get(&self, key: &str) -> Option<Vec<u8>> {
self.get_shared(key).map(|data| data.to_vec())
}
pub(crate) fn get_shared(&self, key: &str) -> Option<Bytes> {
self.store.get(key).map(|entry| entry.data)
}
pub(crate) fn contains(&self, key: &str) -> bool {
self.store.contains_key(key)
}
pub fn get_with_swr(&self, key: &str, threshold_ratio: f64) -> L1SwrRead {
match self.get_with_swr_shared(key, threshold_ratio) {
SharedSwrRead::Fresh(data) => L1SwrRead::Fresh(data.to_vec()),
SharedSwrRead::Stale(data) => L1SwrRead::Stale(data.to_vec()),
SharedSwrRead::Miss => L1SwrRead::Miss,
}
}
pub(crate) fn get_with_swr_shared(&self, key: &str, threshold_ratio: f64) -> SharedSwrRead {
let Some(entry) = self.store.get(key) else {
return SharedSwrRead::Miss;
};
let threshold = entry.ttl.as_secs_f64() * threshold_ratio * entry.freshness_jitter;
if entry.created_at.elapsed().as_secs_f64() > threshold {
SharedSwrRead::Stale(entry.data)
} else {
SharedSwrRead::Fresh(entry.data)
}
}
pub fn set(&self, key: &str, value: &[u8], ttl: Duration) {
self.set_shared(key, Bytes::copy_from_slice(value), ttl);
}
pub(crate) fn set_shared(&self, key: &str, value: Bytes, ttl: Duration) {
self.store.insert(
key.to_string(),
L1Entry {
data: value,
ttl,
created_at: Instant::now(),
freshness_jitter: 0.9 + crate::random_unit() * 0.2,
},
);
}
pub fn delete(&self, key: &str) {
self.store.invalidate(key);
}
pub fn entry_count(&self) -> u64 {
self.store.run_pending_tasks();
self.store.entry_count()
}
pub fn run_pending_tasks(&self) {
self.store.run_pending_tasks();
}
}
#[cfg(test)]
mod tests {
use super::L1Cache;
use std::time::Duration;
#[test]
fn contains_is_false_after_ttl_expiry() {
let cache = L1Cache::new(16);
cache.set("k", b"v", Duration::from_millis(50));
assert!(cache.contains("k"));
std::thread::sleep(Duration::from_millis(100));
assert!(!cache.contains("k"));
}
}