1use bincode::config::{BigEndian, Configuration};
37use bincode::serde::{decode_from_slice, encode_to_vec};
38use dashmap::DashMap;
39use serde::{Deserialize, Serialize};
40use sha2::{Digest, Sha256};
41use std::collections::HashMap;
42use std::fs::{self, read_dir, File, OpenOptions};
43use std::io::{Read, Write};
44use std::ops::Deref;
45use std::path::{Path, PathBuf};
46use std::sync::{Arc, Mutex};
47use std::time::{Duration, SystemTime, UNIX_EPOCH};
48
49fn now() -> u128 {
50 SystemTime::now()
51 .duration_since(UNIX_EPOCH)
52 .expect("Time went backwards")
53 .as_millis()
54}
55#[derive(Serialize, Deserialize)]
56struct CacheEntry {
57 value: Vec<u8>,
58 expires_at: u128,
59}
60
61#[derive(Serialize, Deserialize)]
62struct PersistentCache {
63 entries: HashMap<String, CacheEntry>,
64}
65
66#[derive(Clone)]
67pub struct CacheConfig {
68 pub persistent: bool,
69 pub hash_prefix_length: usize,
70 pub cleanup_interval: Duration,
71 pub dir_path: String,
72}
73
74impl Default for CacheConfig {
75 fn default() -> Self {
76 Self {
77 persistent: true,
78 hash_prefix_length: 2,
79 cleanup_interval: Duration::from_secs(60),
80 dir_path: "cache_data".to_string(),
81 }
82 }
83}
84
85#[derive(Clone)]
86pub struct Cache {
87 entries: Arc<DashMap<String, CacheEntry>>,
88 key_locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
89 file_locks: Arc<DashMap<String, Arc<Mutex<()>>>>,
90 config: CacheConfig,
91}
92
93impl Cache {
94 fn config() -> Configuration<BigEndian> {
95 bincode::config::standard()
96 .with_big_endian()
97 .with_variable_int_encoding()
98 }
99
100 fn get_file_path(&self, key: &str) -> PathBuf {
101 let mut hasher = Sha256::new();
102 hasher.update(key.as_bytes());
103 let hash = hasher.finalize();
104 let prefix_len = self.config.hash_prefix_length.min(hash.len());
105 let prefix = hash[..prefix_len]
106 .iter()
107 .map(|b| format!("{:02x}", b).get(0..1).unwrap().to_string())
108 .collect::<String>();
109
110 Path::new(&self.config.dir_path).join(format!("cache_{}.bin", prefix))
111 }
112
113 pub fn new(config: CacheConfig) -> Result<Self, Box<dyn std::error::Error>> {
114 if config.persistent {
115 fs::create_dir_all(&config.dir_path)?;
116 }
117
118 let cache = Self {
119 entries: Arc::new(DashMap::new()),
120 key_locks: Arc::new(DashMap::new()),
121 file_locks: Arc::new(DashMap::new()),
122 config,
123 };
124
125 if cache.config.persistent {
126 cache.load_persistent_data()?;
127 }
128
129 let cache_clone = cache.clone();
130 std::thread::spawn(move || loop {
131 std::thread::sleep(cache_clone.config.cleanup_interval);
132 cache_clone.cleanup();
133 });
134
135 Ok(cache)
136 }
137
138 fn load_persistent_data(&self) -> Result<(), Box<dyn std::error::Error>> {
139 for entry in read_dir(&self.config.dir_path)? {
140 let entry = entry?;
141 let path = entry.path();
142 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
143 let file_key = path.to_string_lossy().to_string();
144 let file_lock = self
145 .file_locks
146 .entry(file_key)
147 .or_insert(Arc::new(Mutex::new(())))
148 .clone();
149 let _guard = file_lock.lock();
150
151 let mut file = File::open(&path)?;
152 let mut buffer = Vec::new();
153 file.read_to_end(&mut buffer)?;
154
155 let persistent_cache: HashMap<String, CacheEntry> =
156 decode_from_slice(&buffer, Self::config())?.0;
157
158 for (key, entry) in persistent_cache {
159 let key_lock = self
160 .key_locks
161 .entry(key.clone())
162 .or_insert(Arc::new(Mutex::new(())))
163 .clone();
164 let _guard = key_lock.lock();
165
166 self.entries.insert(
167 key,
168 CacheEntry {
169 value: entry.value,
170 expires_at: entry.expires_at,
171 },
172 );
173 }
174 }
175 }
176 Ok(())
177 }
178
179 fn cleanup(&self) {
180 let now = now();
181 let mut rm =vec![];
182 for i in self.entries.iter() {
183 if i.expires_at <= now {
184 rm.push(i.key().to_string());
185 }
186 }
187 for key in rm{
188 let _=self.remove(&key);
189 }
190 }
191
192 pub fn get_key_lock(&self, key: &str) -> Arc<Mutex<()>> {
193 self.key_locks
194 .entry(key.to_string())
195 .or_insert(Arc::new(Mutex::new(())))
196 .clone()
197 }
198 pub fn set<V: Serialize>(
199 &self,
200 key: &str,
201 value: V,
202 ttl: Duration,
203 ) -> Result<(), Box<dyn std::error::Error>> {
204 let serialized = encode_to_vec(&value, Self::config())?;
205 let expires_at = now() + ttl.as_millis();
206
207 let key_lock = self
209 .key_locks
210 .entry(key.to_string())
211 .or_insert(Arc::new(Mutex::new(())))
212 .clone();
213 let _guard = key_lock.lock();
215
216 self.entries.insert(
218 key.to_string(),
219 CacheEntry {
220 value: serialized,
221 expires_at,
222 },
223 );
224
225 if self.config.persistent {
227 self.persist_key(key)?;
228 }
229
230 Ok(())
231 }
232
233 pub fn set_without_guard<V: Serialize>(
234 &self,
235 key: &str,
236 value: V,
237 ttl: Duration,
238 ) -> Result<(), Box<dyn std::error::Error>> {
239 let serialized = encode_to_vec(&value, Self::config())?;
240 let expires_at = now() + ttl.as_millis();
241
242 self.entries.insert(
244 key.to_string(),
245 CacheEntry {
246 value: serialized,
247 expires_at,
248 },
249 );
250
251 if self.config.persistent {
253 self.persist_key(key)?;
254 }
255
256 Ok(())
257 }
258 fn persist_key(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
259 let file_path = self.get_file_path(key);
260 let file_key = file_path.to_string_lossy().to_string();
261
262 let file_lock = self
264 .file_locks
265 .entry(file_key)
266 .or_insert(Arc::new(Mutex::new(())))
267 .clone();
268
269 let _guard = file_lock.lock();
271
272 let mut persistent_entries = if file_path.exists() {
274 let mut file = File::open(&file_path)?;
275 let mut buffer = Vec::new();
276 file.read_to_end(&mut buffer)?;
277 let r: PersistentCache = decode_from_slice(&buffer, Self::config())?.0;
278 r.entries
279 } else {
280 HashMap::new()
281 };
282
283 if let Some(entry) = self.entries.get(key) {
285 persistent_entries.insert(
286 key.to_string(),
287 CacheEntry {
288 value: entry.value.clone(),
289 expires_at: entry.expires_at,
290 },
291 );
292 } else {
293 persistent_entries.remove(key);
294 }
295
296 if persistent_entries.len() == 0 {
297 fs::remove_file(file_path)?;
298 return Ok(());
299 }
300 let persistent_cache = PersistentCache {
302 entries: persistent_entries,
303 };
304 let serialized = encode_to_vec(&persistent_cache, Self::config())?;
305
306 let mut file = OpenOptions::new()
307 .create(true)
308 .write(true)
309 .truncate(true)
310 .open(&file_path)?;
311 file.write_all(&serialized)?;
312
313 Ok(())
314 }
315
316 pub fn get<V: for<'de> Deserialize<'de>>(&self, key: &str) -> Option<V> {
317 let now = now();
318 self.entries.get(key).and_then(|entry| {
319 if entry.expires_at > now {
320 decode_from_slice(&entry.value, Self::config())
321 .ok()
322 .map(|(v, _)| v)
323 } else {
324 None
325 }
326 })
327 }
328
329 pub fn expire(&self, key: &str) -> Option<Duration> {
330 let entries = self.entries.deref();
331 let now = now();
332 entries.get(key).and_then(|entry| {
333 let a = entry.expires_at - now;
334 if entry.expires_at > 0 {
335 Some(Duration::from_millis(a as u64))
336 } else {
337 None
338 }
339 })
340 }
341 pub fn remove(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
342 {
343 let key_lock = self
345 .key_locks
346 .entry(key.to_string())
347 .or_insert(Arc::new(Mutex::new(())))
348 .clone();
349 let _guard = key_lock.lock();
350
351 self.entries.remove(key);
353
354 if self.config.persistent {
356 self.persist_key(key)?;
357 }
358 }
359 self.key_locks.remove(key);
360
361 Ok(())
362 }
363 pub fn remove_without_guard(&self, key: &str) -> Result<(), Box<dyn std::error::Error>> {
364 self.entries.remove(key);
366
367 if self.config.persistent {
369 self.persist_key(key)?;
370 }
371
372 Ok(())
373 }
374
375 pub fn clear(&self) -> Result<(), Box<dyn std::error::Error>> {
376 self.entries.clear();
378 self.key_locks.clear();
379 self.file_locks.clear();
380 if self.config.persistent {
382 for entry in read_dir(&self.config.dir_path)? {
383 let entry = entry?;
384 let path = entry.path();
385 if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("bin") {
386 fs::remove_file(path)?;
387 }
388 }
389 }
390
391 Ok(())
392 }
393
394 pub fn len(&self) -> usize {
395 self.entries.len()
396 }
397
398 pub fn is_empty(&self) -> bool {
399 self.entries.is_empty()
400 }
401}