use rust_webx_core::cache::options::DistributedCacheEntryOptions;
use rust_webx_core::cache::trait_def::{CacheError, IDistributedCache, Result};
use std::collections::{HashMap, VecDeque};
use std::time::Instant;
use tokio::sync::RwLock;
struct CacheEntry {
data: Vec<u8>,
expires_at: Option<Instant>,
sliding_ttl: Option<std::time::Duration>,
}
impl CacheEntry {
fn new(data: Vec<u8>, options: &DistributedCacheEntryOptions) -> Self {
let expires_at = options
.absolute_expiration_relative_to_now
.map(|d| Instant::now() + d);
let sliding_ttl = options.sliding_expiration;
let expires_at = if expires_at.is_some() {
expires_at
} else {
sliding_ttl.map(|d| Instant::now() + d)
};
Self {
data,
expires_at,
sliding_ttl,
}
}
fn is_expired(&self) -> bool {
self.expires_at.is_some_and(|t| Instant::now() >= t)
}
fn refresh(&mut self) {
if let Some(ttl) = self.sliding_ttl {
self.expires_at = Some(Instant::now() + ttl);
}
}
}
struct CacheInner {
entries: HashMap<String, CacheEntry>,
insertion_order: VecDeque<String>,
}
impl CacheInner {
fn evict_expired(&mut self) {
let expired: Vec<String> = self
.entries
.iter()
.filter(|(_, v)| v.is_expired())
.map(|(k, _)| k.clone())
.collect();
for k in expired {
self.entries.remove(&k);
self.insertion_order.retain(|x| x != &k);
}
}
}
pub struct MemoryCache {
inner: RwLock<CacheInner>,
max_entries: usize,
}
impl MemoryCache {
pub fn new() -> Self {
Self {
inner: RwLock::new(CacheInner {
entries: HashMap::new(),
insertion_order: VecDeque::new(),
}),
max_entries: 0,
}
}
pub fn with_max_entries(mut self, n: usize) -> Self {
self.max_entries = n;
self
}
pub async fn compact(&self) {
let mut inner = self.inner.write().await;
inner.evict_expired();
}
pub async fn count(&self) -> usize {
self.inner.read().await.entries.len()
}
}
impl Default for MemoryCache {
fn default() -> Self {
Self::new()
}
}
#[async_trait::async_trait]
impl IDistributedCache for MemoryCache {
async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
{
let inner = self.inner.read().await;
match inner.entries.get(key) {
Some(e) if !e.is_expired() => {
let data = e.data.clone();
let needs_refresh = e.sliding_ttl.is_some();
drop(inner);
if needs_refresh {
let mut inner = self.inner.write().await;
if let Some(e) = inner.entries.get_mut(key) {
e.refresh();
}
}
return Ok(Some(data));
}
Some(_) => {}
None => return Ok(None),
}
}
let mut inner = self.inner.write().await;
match inner.entries.get_mut(key) {
Some(e) if !e.is_expired() => {
e.refresh();
Ok(Some(e.data.clone()))
}
Some(_) => {
inner.entries.remove(key);
inner.insertion_order.retain(|x| x != key);
Ok(None)
}
None => Ok(None),
}
}
async fn set(
&self,
key: &str,
value: Vec<u8>,
options: Option<&DistributedCacheEntryOptions>,
) -> Result<()> {
let opts = options.cloned().unwrap_or_default();
if opts.size_limit > 0 && value.len() > opts.size_limit {
return Err(CacheError::Message(format!(
"size {} exceeds limit {}",
value.len(),
opts.size_limit
)));
}
let mut inner = self.inner.write().await;
let is_new = !inner.entries.contains_key(key);
if is_new && self.max_entries > 0 && inner.entries.len() >= self.max_entries {
inner.evict_expired();
while inner.entries.len() >= self.max_entries {
match inner.insertion_order.pop_front() {
Some(k) => {
inner.entries.remove(&k);
}
None => break,
}
}
}
inner
.entries
.insert(key.to_string(), CacheEntry::new(value, &opts));
if is_new {
inner.insertion_order.push_back(key.to_string());
}
Ok(())
}
async fn remove(&self, key: &str) -> Result<()> {
let mut inner = self.inner.write().await;
inner.entries.remove(key);
inner.insertion_order.retain(|x| x != key);
Ok(())
}
async fn refresh(&self, key: &str) -> Result<()> {
let mut inner = self.inner.write().await;
if let Some(e) = inner.entries.get_mut(key) {
e.refresh();
}
Ok(())
}
async fn exists(&self, key: &str) -> Result<bool> {
{
let inner = self.inner.read().await;
match inner.entries.get(key) {
Some(e) if !e.is_expired() => {
let needs_refresh = e.sliding_ttl.is_some();
drop(inner);
if needs_refresh {
let mut inner = self.inner.write().await;
if let Some(e) = inner.entries.get_mut(key) {
e.refresh();
}
}
return Ok(true);
}
Some(_) => {}
None => return Ok(false),
}
}
let mut inner = self.inner.write().await;
match inner.entries.get_mut(key) {
Some(e) if !e.is_expired() => {
e.refresh();
Ok(true)
}
Some(_) => {
inner.entries.remove(key);
inner.insertion_order.retain(|x| x != key);
Ok(false)
}
None => Ok(false),
}
}
async fn clear(&self) -> Result<()> {
let mut inner = self.inner.write().await;
inner.entries.clear();
inner.insertion_order.clear();
Ok(())
}
}