rust_webx_host/
memory_cache.rs1use rust_webx_core::cache::options::DistributedCacheEntryOptions;
7use rust_webx_core::cache::trait_def::{CacheError, IDistributedCache, Result};
8use std::collections::HashMap;
9use std::time::Instant;
10use tokio::sync::RwLock;
11
12struct CacheEntry {
13 data: Vec<u8>,
14 expires_at: Option<Instant>,
15 sliding_ttl: Option<std::time::Duration>,
16}
17
18impl CacheEntry {
19 fn new(data: Vec<u8>, options: &DistributedCacheEntryOptions) -> Self {
20 let expires_at = options
21 .absolute_expiration_relative_to_now
22 .map(|d| Instant::now() + d);
23 let sliding_ttl = options.sliding_expiration;
24 let expires_at = if expires_at.is_some() {
25 expires_at
26 } else {
27 sliding_ttl.map(|d| Instant::now() + d)
28 };
29 Self {
30 data,
31 expires_at,
32 sliding_ttl,
33 }
34 }
35 fn is_expired(&self) -> bool {
36 self.expires_at.is_some_and(|t| Instant::now() >= t)
37 }
38 fn refresh(&mut self) {
39 if let Some(ttl) = self.sliding_ttl {
40 self.expires_at = Some(Instant::now() + ttl);
41 }
42 }
43}
44
45pub struct MemoryCache {
46 inner: RwLock<HashMap<String, CacheEntry>>,
47 max_entries: usize,
48}
49
50impl MemoryCache {
51 pub fn new() -> Self {
52 Self {
53 inner: RwLock::new(HashMap::new()),
54 max_entries: 0,
55 }
56 }
57 pub fn with_max_entries(mut self, n: usize) -> Self {
58 self.max_entries = n;
59 self
60 }
61 pub async fn compact(&self) {
62 self.inner.write().await.retain(|_, v| !v.is_expired());
63 }
64 pub async fn count(&self) -> usize {
65 self.inner.read().await.len()
66 }
67}
68
69impl Default for MemoryCache {
70 fn default() -> Self {
71 Self::new()
72 }
73}
74
75#[async_trait::async_trait]
76impl IDistributedCache for MemoryCache {
77 async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
78 let mut inner = self.inner.write().await;
79 match inner.get_mut(key) {
80 Some(e) if !e.is_expired() => {
81 e.refresh();
82 Ok(Some(e.data.clone()))
83 }
84 Some(_) => {
85 inner.remove(key);
86 Ok(None)
87 }
88 None => Ok(None),
89 }
90 }
91
92 async fn set(
93 &self,
94 key: &str,
95 value: Vec<u8>,
96 options: Option<&DistributedCacheEntryOptions>,
97 ) -> Result<()> {
98 let opts = options.cloned().unwrap_or_default();
99 if opts.size_limit > 0 && value.len() > opts.size_limit {
100 return Err(CacheError::Message(format!(
101 "size {} exceeds limit {}",
102 value.len(),
103 opts.size_limit
104 )));
105 }
106 let mut inner = self.inner.write().await;
107 if self.max_entries > 0 && inner.len() >= self.max_entries && !inner.contains_key(key) {
108 inner.retain(|_, v| !v.is_expired());
109 if inner.len() >= self.max_entries {
110 if let Some(k) = inner.keys().next().cloned() {
111 inner.remove(&k);
112 }
113 }
114 }
115 inner.insert(key.to_string(), CacheEntry::new(value, &opts));
116 Ok(())
117 }
118
119 async fn remove(&self, key: &str) -> Result<()> {
120 self.inner.write().await.remove(key);
121 Ok(())
122 }
123
124 async fn refresh(&self, key: &str) -> Result<()> {
125 let mut inner = self.inner.write().await;
126 if let Some(e) = inner.get_mut(key) {
127 e.refresh();
128 }
129 Ok(())
130 }
131
132 async fn exists(&self, key: &str) -> Result<bool> {
133 let mut inner = self.inner.write().await;
134 match inner.get_mut(key) {
135 Some(e) if !e.is_expired() => {
136 e.refresh();
137 Ok(true)
138 }
139 Some(_) => {
140 inner.remove(key);
141 Ok(false)
142 }
143 None => Ok(false),
144 }
145 }
146
147 async fn clear(&self) -> Result<()> {
148 self.inner.write().await.clear();
149 Ok(())
150 }
151}